Skip to content

UN-3815 [FIX] Apply organization scoping to prompt-studio child models - #2213

Open
athul-rs wants to merge 8 commits into
mainfrom
UN-3794-org-scoping
Open

UN-3815 [FIX] Apply organization scoping to prompt-studio child models#2213
athul-rs wants to merge 8 commits into
mainfrom
UN-3794-org-scoping

Conversation

@athul-rs

@athul-rs athul-rs commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

What

  • Applies organization scoping to the five prompt-studio child models — DocumentManager, IndexManager, PromptStudioOutputManager, ToolStudioPrompt, ProfileManager — via the existing OrgAwareManager.
  • Scopes the three lookups that take an id straight from the request: delete_for_ide, get_output_for_tool_default, make_profile_default.
  • Pins the FK path each model uses to reach Organization instead of re-deriving it by BFS on every fresh process.
  • Removes the unused file/delete route and action.
  • Adds select_for_update(of=("self",)) where the new org filter introduces joins.

Why

  • Custom actions bypass the global org filter. OrganizationFilterBackend runs in filter_queryset(), which custom DRF @action methods never call. These five models have no organization FK and used a plain BaseModelManager, so a raw .objects lookup inside a custom action carried no organization predicate at all — roughly 44 call sites relying on the caller to pass a correct id. OrgAwareManager already existed for exactly this shape but was only wired to one model (ExecutionLog).
  • BFS path resolution is order-dependent. get_org_path returns the shortest FK chain to Organization and breaks ties by field declaration order. Reordering two fields on a model can swap in a different path of the same length. If that path runs through a nullable FK, Django turns the filter into an INNER JOIN and silently drops every row with a NULL — data loss that presents as missing records, not as an error. This applies to OrganizationFilterBackend in production today, independent of anything else in this PR.
  • file/delete is dead and shaped wrong. No caller anywhere in the frontend or backend; prompt-studio/file/<tool_id> (DELETE → delete_for_ide) is the live path. It also performed a delete over GET, which makes it prefetchable.

How

  • objects = OrgAwareManager() on the four models with no custom manager; ProfileManagerModelManager now extends OrgAwareManager instead of BaseModelManager. No migration — no manager sets use_in_migrations, so swapping objects serializes nothing (makemigrations --check --dry-run is clean).
  • ORG_PATH_OVERRIDES in org_path_discovery.py, keyed by model label and checked before BFS, so both consumers (OrgAwareManager and OrganizationFilterBackend) are frozen on the same value. Each pin is set to the path BFS resolves today, so this changes no behaviour on its own.
    • ProfileManager pins to vector_store__organization, not prompt_studio_tool__organization: BFS reaches AdapterInstance (which carries the organization FK) before CustomTool, and prompt_studio_tool is nullable, so pinning there would drop tool-less profiles. AdapterInstance is org-owned and the serializer's FK queryset uses the org-scoped default manager, so this scopes to the same organization.
  • The three request-id lookups gain an explicit predicate and use get_object_or_404, so a non-matching id is a 404 rather than an unhandled Model.DoesNotExist — which middleware.exception.drf_logging_exc_handler does not map, and would surface as a 500.
  • select_for_update(of=("self",)) in prompt_studio_index_helper and migration_utils: the org filter adds INNER JOINs, and Postgres FOR UPDATE without of= locks rows in every joined table (DocumentManager, CustomTool, AdapterInstance).

Can this PR break any existing features. If yes, please list possible items. If no, please explain why. (PS: Admins do not merge the PR without this section filled)

Yes — three areas, each covered by a test:

  1. Worker and Celery paths. Organization context is set in worker, internal-API and scheduler paths (internal_api_auth.py, scheduler/tasks.py, workflow_helper.py), so OrgAwareManager filters there too — it is not a no-op outside requests. Indexing and execution pass because the worker's org matches the data's org. test_worker_context_sees_its_own_org covers this. Any path that legitimately spans organizations would now return empty; none was found.
  2. get_or_create under a filtering manager. If the get half is filtered out while the row exists, the create half hits the unique constraint. Only reachable across organizations, but the failure mode changes from silently-wrong to IntegrityError.
  3. select_for_update lock scope. Addressed with of=("self",); without it the joins would widen the lock. ProfileManager and IndexManager are the two affected call sites.

Management commands and shell keep full access: UserContext.get_organization() returns None outside a request and the manager fails open, unchanged. test_no_org_context_is_unfiltered pins that.

file/delete removal is the one behaviour change with no in-repo caller to break. Any external API consumer of that endpoint is unknowable from this repo — worth a release note.

Database Migrations

None. makemigrations --check --dry-run is clean; no manager sets use_in_migrations, so replacing objects does not produce a migration.

Env Config

None.

Relevant Docs

None.

Related Issues or PRs

UN-3815

Dependencies Versions

Unchanged.

Notes on Testing

  • backend/prompt_studio/tests/test_cross_org_isolation.py — two fully populated organizations, then per-model checks that org A cannot reach org B's rows, that org A's own rows stay visible, that worker context still sees its own org, and that the no-org path stays unfiltered. Every isolation assertion was confirmed to fail against main before the fix, so the tests actually bite.
  • backend/utils/tests/test_org_path_discovery.py — asserts each pin still matches what BFS resolves, and that no pin traverses a nullable FK (with one documented exception, ToolStudioPrompt.tool_id, which is the path already in force).
  • Full backend suite run against this branch and against main: identical failure sets (36, all pre-existing in workflow_manager/execution/tests/test_pg_finalization_fixes.py), zero new.

Not covered by automation: a real two-org Prompt Studio cycle (upload → index → run → delete) in a compose stack. Worth doing manually before merge.

Screenshots

Checklist

I have read and understood the Contribution Guidelines.

@athul-rs
athul-rs requested review from jaseemjaskp and ritwik-g July 27, 2026 04:36
@coderabbitai

coderabbitai Bot commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

Summary by CodeRabbit

  • Bug Fixes

    • Strengthened organization-level data isolation across prompt studio resources.
    • Prevented unauthorized cross-organization access through direct resource lookups.
    • Improved not-found handling with appropriate 404 responses.
    • Organization filtering now fails safely when context is missing or invalid.
    • Improved concurrency handling during profile and index updates.
    • Removed the file deletion endpoint; file listing, downloading, and uploading remain available.
  • Tests

    • Added coverage for organization isolation, scoped lookups, safe filtering, and removed file deletion routes.

Walkthrough

The changes enforce organization-aware querying, constrain Prompt Studio lookups, remove file deletion routes, narrow row-locking behavior, and add organization-path and cross-organization isolation tests.

Changes

Organization isolation and access control

Layer / File(s) Summary
Organization path and fail-closed filtering
backend/utils/models/org_path_discovery.py, backend/utils/organization_utils.py, backend/utils/tests/*
Organization paths are pinned and validated. Missing or unresolved organization context returns empty querysets.
Prompt Studio manager and lookup scoping
backend/prompt_studio/prompt_*/models.py, backend/prompt_studio/prompt_studio_core_v2/views.py, backend/prompt_studio/prompt_studio_output_manager_v2/views.py, backend/prompt_studio/tests/test_cross_org_isolation.py
Models use organization-aware managers. Profile, document, prompt, and output lookups enforce organization or tool scope. Isolation tests cover these paths.

Endpoint and transaction changes

Layer / File(s) Summary
File deletion endpoint removal
backend/file_management/urls.py, backend/file_management/views.py
The file deletion action and route are removed. Listing, download, and upload routes remain.
Self-only row locking
backend/prompt_studio/prompt_studio_core_v2/migration_utils.py, backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py
Migration and extraction-status operations restrict locks to target model rows.

Estimated code review effort: 4 (Complex) | ~45 minutes

Suggested reviewers: jaseemjaskp, ritwik-g, chandrasekharan-zipstack, muhammad-ali-e

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Write docstrings for the functions missing them to satisfy the coverage threshold.
✅ Passed checks (4 passed)
Check name Status Explanation
Title check ✅ Passed The title clearly identifies the organization-scoping fix for Prompt Studio child models.
Description check ✅ Passed The description covers the template sections and clearly documents scope, risks, migrations, testing, and related issue details.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches 💡 1
📝 Generate docstrings 💡
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch UN-3794-org-scoping

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@greptile-apps

greptile-apps Bot commented Jul 27, 2026

Copy link
Copy Markdown
Contributor

Greptile Summary

This PR applies organization-aware scoping throughout Prompt Studio’s child-model graph.

  • Adds OrgAwareManager to documents, indexes, outputs, prompts, and profiles.
  • Pins deterministic organization paths and adds isolation and fail-closed behavior tests.
  • Constrains request-supplied document, profile, and tool identifiers to the authorized organization or parent tool.
  • Preserves narrow PostgreSQL row-lock scope where organization filtering introduces joins.
  • Removes the unused state-changing file/delete GET endpoint.

Confidence Score: 5/5

The PR appears safe to merge.

No blocking failure remains; the previously questioned profile path is consistent with the enforced adapter ownership model and the current changes preserve organization isolation across the affected Prompt Studio models.

Important Files Changed

Filename Overview
backend/utils/models/org_path_discovery.py Pins stable organization traversal paths for Prompt Studio child models, including profile scoping through its required vector-store adapter.
backend/utils/models/org_aware_manager.py Applies organization filtering when context exists while retaining unfiltered access for context-free operational paths.
backend/prompt_studio/prompt_profile_manager_v2/models.py Moves profile queries onto the organization-aware manager without changing the persisted schema.
backend/prompt_studio/prompt_studio_core_v2/views.py Tool-scopes request-supplied document and profile lookups and makes default-profile switching atomic.
backend/prompt_studio/tests/test_cross_org_isolation.py Exercises cross-organization denial, same-organization visibility, worker context, default switching, and removal of the obsolete delete route.

Flowchart

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    Org[Organization]
    Tool[CustomTool]
    Adapter[AdapterInstance]
    Profile[ProfileManager]
    Document[DocumentManager]
    Prompt[ToolStudioPrompt]
    Index[IndexManager]
    Output[PromptStudioOutputManager]

    Org --> Tool
    Org --> Adapter
    Adapter -->|pinned organization path| Profile
    Tool --> Document
    Tool --> Prompt
    Tool --> Profile
    Document --> Index
    Profile --> Index
    Document --> Output
    Profile --> Output
    Prompt --> Output
Loading

Reviews (7): Last reviewed commit: "Merge remote-tracking branch 'origin/mai..." | Re-trigger Greptile

Comment thread backend/utils/models/org_path_discovery.py

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🧹 Nitpick comments (2)
backend/file_management/views.py (1)

28-33: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Stale docstring still advertises DELETE.

The class docstring still says the viewset "Handles GET,POST,PUT,PATCH and DELETE" but the delete action (and its URL route) is now gone.

✏️ Proposed docstring fix
     """FileManagement view.

-    Handles GET,POST,PUT,PATCH and DELETE
+    Handles GET, POST, PUT, and PATCH
     """
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@backend/file_management/views.py` around lines 28 - 33, Update the
FileManagementViewSet class docstring to remove DELETE from the listed supported
operations, leaving only the methods and actions still exposed by the viewset.
backend/prompt_studio/tests/test_cross_org_isolation.py (1)

100-104: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

No tearDown to reset UserContext after mutating tests.

test_no_org_context_is_unfiltered (Line 143) sets the org identifier to None, and test_worker_context_sees_its_own_org (Line 151) sets it to org B; neither is restored. Since UserContext looks like process-level/thread-local state (not something Django's transactional TestCase rolls back), whichever of these runs last leaves stale org context for the next test class in the same run.

♻️ Proposed fix
     def setUp(self) -> None:
         self.a = OrgFixture(f"org-a-{secrets.token_hex(3)}")
         self.b = OrgFixture(f"org-b-{secrets.token_hex(3)}")
         # End state: acting as org A, as a request would.
         UserContext.set_organization_identifier(self.a.org.organization_id)
+
+    def tearDown(self) -> None:
+        UserContext.set_organization_identifier(None)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@backend/prompt_studio/tests/test_cross_org_isolation.py` around lines 100 -
104, Update the test fixture class containing setUp, OrgFixture, and the
affected isolation tests with a tearDown method that clears or restores
UserContext’s organization identifier after every test. Ensure tests that mutate
the context, including test_no_org_context_is_unfiltered and
test_worker_context_sees_its_own_org, cannot leak state into subsequent tests.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@backend/prompt_studio/prompt_studio_core_v2/views.py`:
- Around line 446-453: The default-profile update flow should resolve the target
ProfileManager before clearing the current default, so invalid or cross-tool IDs
leave existing state unchanged. In the relevant view method, move the
get_object_or_404 lookup for prompt_tool and request.data["default_profile"]
ahead of the reset, then wrap target validation and both default updates in
transaction.atomic().

In `@backend/prompt_studio/prompt_studio_output_manager_v2/views.py`:
- Around line 127-132: Update fetch_default_output_response() after the
organization-scoped ToolStudioPrompt.objects.filter() lookup to explicitly
detect an empty queryset and raise the existing tool-not-found error. Preserve
the scoped tool_id and organization filters, and continue using the queryset for
valid tools.

---

Nitpick comments:
In `@backend/file_management/views.py`:
- Around line 28-33: Update the FileManagementViewSet class docstring to remove
DELETE from the listed supported operations, leaving only the methods and
actions still exposed by the viewset.

In `@backend/prompt_studio/tests/test_cross_org_isolation.py`:
- Around line 100-104: Update the test fixture class containing setUp,
OrgFixture, and the affected isolation tests with a tearDown method that clears
or restores UserContext’s organization identifier after every test. Ensure tests
that mutate the context, including test_no_org_context_is_unfiltered and
test_worker_context_sees_its_own_org, cannot leak state into subsequent tests.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 99c5a213-933f-481f-a966-d43ed583fd72

📥 Commits

Reviewing files that changed from the base of the PR and between 023b140 and 65613a0.

📒 Files selected for processing (15)
  • backend/file_management/urls.py
  • backend/file_management/views.py
  • backend/prompt_studio/prompt_profile_manager_v2/models.py
  • backend/prompt_studio/prompt_studio_core_v2/migration_utils.py
  • backend/prompt_studio/prompt_studio_core_v2/views.py
  • backend/prompt_studio/prompt_studio_document_manager_v2/models.py
  • backend/prompt_studio/prompt_studio_index_manager_v2/models.py
  • backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py
  • backend/prompt_studio/prompt_studio_output_manager_v2/models.py
  • backend/prompt_studio/prompt_studio_output_manager_v2/views.py
  • backend/prompt_studio/prompt_studio_v2/models.py
  • backend/prompt_studio/tests/__init__.py
  • backend/prompt_studio/tests/test_cross_org_isolation.py
  • backend/utils/models/org_path_discovery.py
  • backend/utils/tests/test_org_path_discovery.py
💤 Files with no reviewable changes (1)
  • backend/file_management/urls.py

Comment thread backend/prompt_studio/prompt_studio_core_v2/views.py Outdated
Comment thread backend/prompt_studio/prompt_studio_output_manager_v2/views.py Outdated
@athul-rs athul-rs changed the title UN-3794 [FIX] Apply organization scoping to prompt-studio child models UN-3815 [FIX] Apply organization scoping to prompt-studio child models Jul 27, 2026
@athul-rs
athul-rs marked this pull request as draft July 27, 2026 19:16
athul-rs added 3 commits July 29, 2026 15:10
get_org_path resolves the shortest FK chain from a model to Organization
and breaks ties by field declaration order. Reordering two fields can
therefore swap in a different path of the same length, and if that path
runs through a nullable FK the org filter becomes an INNER JOIN that
silently drops every row with a NULL — which reads as missing records
rather than as an error.

Pin the five prompt-studio models to their currently resolved paths so
both consumers (OrgAwareManager and OrganizationFilterBackend) are frozen
on the same value, and add tests that fail if a pin drifts from discovery
or starts traversing a nullable FK.

ProfileManager resolves to vector_store__organization rather than
prompt_studio_tool__organization: BFS reaches AdapterInstance (which
carries the organization FK) before CustomTool, and prompt_studio_tool is
nullable, so pinning there would drop tool-less profiles.
Custom DRF @action methods never call filter_queryset(), so
OrganizationFilterBackend does not run on them and a raw .objects lookup
inside one carries no organization predicate. Five prompt-studio models
have no organization FK and used a plain manager, leaving roughly 44 such
call sites relying on the caller to pass a correct id.

- Scope at the model layer: OrgAwareManager on DocumentManager,
  IndexManager, PromptStudioOutputManager, ToolStudioPrompt and
  ProfileManager. No migration — no manager sets use_in_migrations, so
  swapping objects serializes nothing.
- Scope the lookups that take an id straight from the request:
  delete_for_ide now requires the document to belong to the tool the
  caller already passed authz on, get_output_for_tool_default filters
  prompts by organization, and make_profile_default constrains its
  secondary lookup to the same tool. All three use get_object_or_404 so a
  non-matching id is a 404 rather than an unhandled DoesNotExist, which
  the DRF handler would turn into a 500.
- Drop the file/delete route and action: it has no caller, and it deleted
  a document over GET.
- select_for_update(of=("self",)) where the org filter now adds joins, so
  Postgres does not also lock rows in DocumentManager, CustomTool or
  AdapterInstance.

Tests cover the org isolation matrix, same-org access, worker context
(org is set there, so the manager filters) and the no-org fail-open path.
…aults

make_profile_default cleared is_default across every profile on the tool
and only then resolved the id from the request body. A non-matching id
left the tool with no default at all, and the two writes were not in a
transaction.

Resolve first, then clear and set inside a single transaction, so a
rejected id changes nothing. Adds a regression test for that, plus a
tearDown resetting the thread-local UserContext (TestCase rollback does
not clear it, so the org-switching tests leaked into later classes) and
drops DELETE from the FileManagement docstring now the route is gone.
@athul-rs
athul-rs force-pushed the UN-3794-org-scoping branch from 65613a0 to 14f94cd Compare July 29, 2026 09:42
@athul-rs
athul-rs marked this pull request as ready for review July 29, 2026 09:42

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
backend/prompt_studio/tests/test_cross_org_isolation.py (1)

163-207: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Exercise the actions, not only their ORM predicates.

These tests recreate the intended lookups directly, so they cannot catch a regression in delete_for_ide or make_profile_default’s HTTP 404 mapping or mutation order. Add authenticated action requests that assert 404 and preserve the original default profile.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@backend/prompt_studio/tests/test_cross_org_isolation.py` around lines 163 -
207, Add authenticated HTTP action tests covering delete_for_ide and
make_profile_default, rather than only direct DocumentManager/ProfileManager
lookups. Use cross-organization or cross-tool IDs, assert each endpoint returns
404, and verify the target tool’s existing default profile remains unchanged
after each rejected request.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@backend/prompt_studio/tests/test_cross_org_isolation.py`:
- Around line 163-207: Add authenticated HTTP action tests covering
delete_for_ide and make_profile_default, rather than only direct
DocumentManager/ProfileManager lookups. Use cross-organization or cross-tool
IDs, assert each endpoint returns 404, and verify the target tool’s existing
default profile remains unchanged after each rejected request.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 7b3f6192-48dd-4517-b09f-875acf509d01

📥 Commits

Reviewing files that changed from the base of the PR and between 65613a0 and 14f94cd.

📒 Files selected for processing (15)
  • backend/file_management/urls.py
  • backend/file_management/views.py
  • backend/prompt_studio/prompt_profile_manager_v2/models.py
  • backend/prompt_studio/prompt_studio_core_v2/migration_utils.py
  • backend/prompt_studio/prompt_studio_core_v2/views.py
  • backend/prompt_studio/prompt_studio_document_manager_v2/models.py
  • backend/prompt_studio/prompt_studio_index_manager_v2/models.py
  • backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py
  • backend/prompt_studio/prompt_studio_output_manager_v2/models.py
  • backend/prompt_studio/prompt_studio_output_manager_v2/views.py
  • backend/prompt_studio/prompt_studio_v2/models.py
  • backend/prompt_studio/tests/__init__.py
  • backend/prompt_studio/tests/test_cross_org_isolation.py
  • backend/utils/models/org_path_discovery.py
  • backend/utils/tests/test_org_path_discovery.py
💤 Files with no reviewable changes (1)
  • backend/file_management/urls.py
🚧 Files skipped from review as they are similar to previous changes (9)
  • backend/prompt_studio/prompt_studio_output_manager_v2/views.py
  • backend/prompt_studio/prompt_studio_core_v2/migration_utils.py
  • backend/prompt_studio/prompt_studio_document_manager_v2/models.py
  • backend/prompt_studio/prompt_profile_manager_v2/models.py
  • backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py
  • backend/prompt_studio/prompt_studio_index_manager_v2/models.py
  • backend/utils/tests/test_org_path_discovery.py
  • backend/file_management/views.py
  • backend/prompt_studio/prompt_studio_v2/models.py

athul-rs added 2 commits July 31, 2026 00:41
…ol_default

filter() does not raise ObjectDoesNotExist, so the except branch could
never fire and the tool-not-found message was dead. Empty is the right
result here anyway: it covers a missing tool, an out-of-org tool, and a
newly created project that has no prompts yet, which is a normal state
that must not 400.
filter_queryset_by_organization returned the queryset unfiltered when the
request carried no organization context, which is the opposite of what a
scoping helper should do — and its own docstring already claimed it
returned an empty queryset.

Six internal viewsets set skip_org_filter = True, which disables
OrganizationFilterBackend and leaves this helper as their only tenant
boundary across roughly 39 call sites. The internal auth middleware logs
a warning and continues when X-Organization-ID is missing, so any caller
holding the internal service key reached those endpoints without context
by omitting the header, reading across every organization — and through
the file-execution viewset, writing and deleting too.

Return none() instead, and log loudly, so a caller that legitimately has
no context is visible rather than silently served everything.

Deliberately not rejecting header-less /internal/ requests in the
middleware: the leader-elected reaper calls without the header on
purpose, to scan across organizations. It queries the model directly
rather than through this helper, so failing closed leaves it working.
@ritwik-g

ritwik-g commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

@greptileai please review this

@athul-rs

athul-rs commented Aug 4, 2026

Copy link
Copy Markdown
Contributor Author

@greptileai re-review this PR

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
backend/prompt_studio/prompt_studio_core_v2/views.py (1)

453-456: 🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win

Return 400 when default_profile is missing.

request.data["default_profile"] raises KeyError when the field is omitted. Django REST Framework then returns a 500 response. Validate the field with a serializer or use .get() and return a 400 response before the scoped lookup.

Proposed fix
+        default_profile_id = request.data.get("default_profile")
+        if default_profile_id is None:
+            return Response(
+                {"detail": "default_profile is required."},
+                status=status.HTTP_400_BAD_REQUEST,
+            )
+
         profile_manager = get_object_or_404(
             ProfileManager,
-            pk=request.data["default_profile"],
+            pk=default_profile_id,
             prompt_studio_tool=prompt_tool,
         )
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@backend/prompt_studio/prompt_studio_core_v2/views.py` around lines 453 - 456,
Update the view logic around the ProfileManager lookup to validate that
default_profile is present before accessing request.data["default_profile"].
Return a 400 response when it is omitted, while preserving the existing scoped
lookup through prompt_studio_tool for valid values; use the view’s established
validation or error-response pattern.
🧹 Nitpick comments (1)
backend/prompt_studio/prompt_studio_core_v2/views.py (1)

137-139: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Mark the class configuration as ClassVar.

Ruff reports RUF012 for both mutable class attributes. Add ClassVar annotations to make the shared viewset configuration explicit without changing behavior.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@backend/prompt_studio/prompt_studio_core_v2/views.py` around lines 137 - 139,
Annotate the mutable ordering and ordering_fields class attributes in the
surrounding viewset with ClassVar, importing ClassVar from typing if needed.
Preserve their existing list values and behavior while satisfying Ruff RUF012.

Source: Linters/SAST tools

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@backend/prompt_studio/prompt_studio_core_v2/views.py`:
- Around line 453-456: Update the view logic around the ProfileManager lookup to
validate that default_profile is present before accessing
request.data["default_profile"]. Return a 400 response when it is omitted, while
preserving the existing scoped lookup through prompt_studio_tool for valid
values; use the view’s established validation or error-response pattern.

---

Nitpick comments:
In `@backend/prompt_studio/prompt_studio_core_v2/views.py`:
- Around line 137-139: Annotate the mutable ordering and ordering_fields class
attributes in the surrounding viewset with ClassVar, importing ClassVar from
typing if needed. Preserve their existing list values and behavior while
satisfying Ruff RUF012.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro Plus

Run ID: 5ae00466-ec2f-47b0-a872-58ea3d565c73

📥 Commits

Reviewing files that changed from the base of the PR and between 14b7e68 and 0dce94e.

📒 Files selected for processing (1)
  • backend/prompt_studio/prompt_studio_core_v2/views.py

@chandrasekharan-zipstack chandrasekharan-zipstack left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Standardized review — PR #2213

Verdict: REQUEST CHANGES

Summary — Critical: 0 · High: 4 · Medium: 9 · Low: 9 · Lenses run: 16/16

Reviewed under the unstract:standard-review 16-lens rubric. Findings below are deduplicated against the existing CodeRabbit and Greptile threads — anything already raised there is not repeated. Specifically not re-raised:

  • CodeRabbit's "exercise the actions, not only their ORM predicates" on test_cross_org_isolation.py — I agree with it and rate it higher than Trivial; only the part it did not cover (the make_profile_default ordering fix having vacuous coverage) is filed below.
  • Greptile's "shared adapters hide tool profiles" on org_path_discovery.py:47 — the author's rebuttal is correct; AdapterInstanceModelManager does scope every sharing path to the org. Closed on the merits.
  • CodeRabbit's get_output_for_tool_default empty-200 thread, which the author answered and CodeRabbit accepted. Only the third cause of empty that the thread never discussed is filed below.
  • CodeRabbit's stale-DELETE docstring on file_management/views.py. Residual nit: the replacement line now reads Handles GET, POST, PUT and PATCH, but urls.py routes only GET and POST — no update/partial_update exists on the viewset.

Lens checklist

# Lens Result
1 Spec & intent Clean
2 Architectural fit See H2, M7
3 Correctness & edge cases See H4, M1, M3, M5, M6, M9
4 Security See H2
5 Data integrity & migrations See M1, M2
6 Concurrency Clean — of=("self",) rationale verified correct at both sites; positive filters give INNER JOINs, so no nullable-outer-join hazard
7 API & contract compatibility See H1
8 Reliability & resilience See H1
9 Performance & cost Clean
10 Observability See H4
11 Operational safety See H1 — no flag, no deploy gate
12 LLM/agent N/A — no model calls touched
13 Testing See H3, M2
14 Dependencies & build N/A — none changed. Confirmed no migration needed: no manager sets use_in_migrations
15 Code quality Low only
16 Doc & comment accuracy See M8, M9, and Lows

Unanchored findings (outside the diff hunks)

[High] [Lens 8, 11] — validate_tool_instances_internal returns success: true having validated nothing. backend/tool_instance_v2/internal_views.py:337-397. Function-based @api_view, no filter backend, so filter_queryset_by_organization is its only scoping. Header-less, tool_instances is now empty, the loop never runs, validation_errors stays empty, and it returns HTTP 200 {"success": true, "errors": []}. The adapter-ID migration inside that loop (:355-360) is skipped too. Worker side, workers/shared/workflow/execution/tool_validation.py:120-133 then logs ✅ Validated N tool instances successfully using the requested count, not len(validated_instances). This is the sharpest instance of H1 and the reason I would argue H1 up to Critical if any deployed worker can omit the header.

[Medium] [Lens 3] — import_prompts attaches profile_manager=None to every imported prompt. backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py:3014-3016:3062. ProfileManager.objects.filter(...).first() is now org-scoped and returns None rather than raising when the filter empties; there is no None check before :3062 passes it into every ToolStudioPrompt.objects.create(...). The sibling sync_prompts at :3146-3153 does raise on exactly this — the omission looks accidental rather than deliberate.

[Medium] [Lens 5] — Org-scoped .delete() in sync_prompts can leave survivors. prompt_studio_helper.py:3159-3161. ToolStudioPrompt.objects.filter(tool_id=tool).delete() is now filtered by tool_id__organization; prompts the scope misses survive alongside their recreated replacements inside the same transaction.atomic(). deleted_count is only used to decide whether to bump modified_at, never to verify the delete was complete.

[Medium] [Lens 3] — check_files_history reads org from the header but sets it from the body. backend/workflow_manager/internal_views.py:2584-2595 installs request.data["organization_id"] into StateStore, but filter_queryset_by_organization reads request.organization_id, populated only from the header. A body-only caller previously worked and now gets .none()Workflow.DoesNotExist → 404 "not found or access denied", which blames authorization for a context-plumbing mismatch inside one function.

Low (9)

  • backend/prompt_studio/tests/test_cross_org_isolation.py:112, :156, :163, :208 — review-artifact tags (A-1, A-3, A-4, A-5, B1, "the reported call sites") resolve to nothing in the repo. Same for "pinned as-is rather than changed under a security fix" at org_path_discovery.py:41-42. Repo CLAUDE.md asks that comments read correctly without the authoring session's context.
  • backend/file_management/views.py:31 — "Handles GET, POST, PUT and PATCH"; only GET and POST are routed.
  • Three symbols orphaned by the route removal, each had exactly one caller and this PR deleted it: file_management/serializer.py:53 (FileInfoIdeSerializer), file_management/file_management_helper.py:229 (delete_file), prompt_studio_output_manager_v2/constants.py:12 (TOOL_NOT_FOUND).
  • backend/utils/tests/test_org_path_discovery.py:27-30test_pin_is_returned reduces to d.get(k) == d[k]; it can only fail if the short-circuit is deleted outright.
  • backend/utils/tests/test_organization_scoping.py:23-26, :59-62_Request.__init__ guards on is not None, so the None leg of for falsy in ("", None) produces an object with no attribute at all — byte-identical to test_missing_org_context_returns_nothing.
  • test_cross_org_isolation.py:96, test_organization_scoping.py:29@pytest.mark.django_db is a no-op on TestCase subclasses and does not drive tier selection (backend/conftest.py:38-44 marks on either signal).
  • test_cross_org_isolation.py:100-110OrgFixture.__init__ sets thread-local org context as a construction side effect, and unittest skips tearDown when setUp raises. self.addCleanup(UserContext.set_organization_identifier, None) as the first statement of setUp runs even on failure.
  • test_organization_scoping.py:47-53 — depends on Django's private _base_manager MRO resolution; a base_manager_name added to BaseModel later would silently re-scope the queryset and point the failure at the helper.
  • Django admin changelists for all five models now use OrgAwareManager via _default_manager, and /admin/ is not matched by OrganizationMiddleware, so the list depends on whatever StateStore holds on that thread.

Verified clean, for the record

All five pins match what _discover_org_path returns today (BFS field ordering traced per model). _base_manager stays a plain unfiltered Manager (no base_manager_name on BaseModel), so cascade deletes, forward-FK descriptors and the pre_delete receiver at prompt_studio_index_manager_v2/models.py:122-141 are unaffected. Zero references to file/delete across unstract, unstract-cloud and unstract-docs — the UI deletes via DELETE /prompt-studio/file/<tool_id> (ManageDocsModal.jsx:674-687), so the route removal is correct, and the sibling-route guard in test_file_delete_route_removed is a nice touch. tests/groups.yaml collects both new test directories and CI runs them. The CONCURRENCY_MODERuntimeError → fail-open path in StateStore is real but latent — the env var is set in no compose, helm or env file in either repo.

Open questions

  1. Can any currently deployed worker call an internal endpoint without X-Organization-ID? Three in-repo comments say yes during rolling deploys. That answer decides whether H1 is High or Critical.
  2. Is OrgAwareManager's fail-open deliberate policy, or an artifact of it predating the fail-closed backend? This PR pins it in a test and argues the opposite in a docstring, in the same diff.
  3. SELECT count(*) FROM custom_tool WHERE organization_id IS NULL (and adapter_instance) — several Mediums collapse to nothing if that is zero.

Reviewed with unstract:standard-review v0.18.1 (16-lens rubric, 4 specialist agents). Comments are advisory; event: COMMENT, no merge gate.

Comment on lines +103 to +113
if not org_id:
logger.warning(
"Organization scoping requested without organization context on %s; "
"returning no rows. A caller that must span organizations should "
"query the model directly instead of using this helper.",
getattr(request, "path", "<unknown path>"),
)
return queryset.none()

organization = resolve_organization(org_id, raise_on_not_found=False)
if not organization:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[High] [Lens 7, 8, 11] — Fail-closing this helper breaks a documented worker contract, with no rollout gate

Failing closed is the right end state — I am not arguing against the change itself. The problem is shipping it without a deploy gate while the codebase still documents the header as optional.

Failure mode: six viewsets set skip_org_filter = True, so this helper is their only tenant boundary. A worker that omits X-Organization-ID previously got the unfiltered queryset (the leak being fixed); it now gets zero rows. retrieve() becomes 404, list() becomes empty, and metrics endpoints return a well-formed 200 with all counters at zero — indistinguishable from a genuinely idle system.

Evidence that the header is genuinely optional today:

  • backend/middleware/internal_api_auth.py:157-164 returns {"warning": ..., "context_set": False} and the request proceeds.
  • workers/shared/clients/base_client.py:158 initialises self.organization_id = None, clears it again at :589 and :602, and :315-317 only attaches the header if current_org_id:.
  • Two in-repo comments state the contract this diff breaks and are now stale: workflow_manager/workflow_v2/views.py:405-408 ("workers may call without X-Organization-ID during rolling deployments") and notification_v2/internal_views.py:44 ("Backward compat: remove once all workers pass X-Organization-ID").

The sharpest consequence is in tool_instance_v2/internal_views.py — see "Unanchored findings" in the summary for that one, since it is outside this diff.

Suggested fix: either (a) land the fail-closed switch behind a settings flag defaulted to the old behaviour for one release while still emitting the warning, or (b) make InternalAPIAuthMiddleware reject internal requests without a resolvable X-Organization-ID (400/401) so the failure is explicit rather than an empty 200/404 — and update or delete the three backward-compat comments in the same PR.

Confidence: High that the contract changes and the comments are now false. Medium that live workers hit it — confirming every internal client call site passes organization_id would raise this to High.

Lens 7 · 8 · 11

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The three stale comments are fixed — and there were seven, not three: tool_instance_v2/internal_views.py:26, pipeline_v2/internal_api_views.py:16 and workflow_manager/file_execution/internal_views.py:32-35 carried the same claim. All now state that the helper fails closed and that the header is required in practice.

On the rollout gate: shipping fail-closed without one, deliberately. A flag defaulted to the old behaviour keeps the cross-tenant leak open in the default configuration for another release, which is the thing this PR exists to close. Option (b) — rejecting header-less internal requests at the middleware — trades a silent empty result for a hard failure on every internal caller at once, which is a larger blast radius than the finding it fixes.

Evidence for the risk being low: set_organization_context is called at 25 sites across the worker packages, including internal_client.py:1376-1383, which fans it out to all eight sub-clients at once. That is not a proof that no path omits it, and the confidence rating here is fair. If a gate is wanted anyway, it is a one-line settings check and can go in before merge — flagging it as a deployment call rather than deciding it unilaterally.

Comment thread backend/utils/organization_utils.py Outdated
Comment thread backend/utils/organization_utils.py Outdated
Comment thread backend/prompt_studio/tests/test_cross_org_isolation.py Outdated
Comment thread backend/prompt_studio/prompt_studio_document_manager_v2/models.py Outdated
Comment thread backend/prompt_studio/prompt_studio_output_manager_v2/views.py Outdated
Comment thread backend/prompt_studio/prompt_studio_core_v2/views.py
Comment on lines +1201 to +1207
# Scope to the tool the caller already passed authz on — tighter than
# org scope, and this action never runs filter_queryset().
# get_object_or_404 keeps a non-matching id a 404 rather than an
# unhandled DoesNotExist, which the DRF handler turns into a 500.
document: DocumentManager = get_object_or_404(
DocumentManager, pk=document_id, tool=custom_tool
)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[Medium] [Lens 3, 10] — delete_for_ide reports success while silently leaving Redis indexing flags behind

The scoping tightening on this lookup is correct. The issue is the code immediately after it.

IndexManager.objects.filter(document_manager=document_id) at :1187 is now org-scoped. If it comes back empty because the filter hid the rows rather than because none exist, the for loop body never runs, DocumentIndexingService.remove_document_indexing is never called, and execution proceeds straight to document.delete() and a 200 "File deleted succesfully."

The document row and the file are gone, but the Redis indexing flags persist — so a re-upload of the same file is treated as already-indexed. The user is told the delete succeeded, and nothing distinguishes "this document had no index managers" from "the filter hid them".

The except Exception at :1207-1212 compounds it: connector errors, storage errors, Redis errors and ORM errors all collapse into one 400 {"data": "File deletion failed."} with the detail only in logger.error. Worth noting this PR did correctly delete an identical swallow-everything handler over in file_management/views.py — this one, in the surviving path the same diff edits, was left in place.

Suggested fix: log at WARNING when index_managers is empty, including the resolved org, before proceeding. Split the except Exception into the specific failures (ConnectorError, storage exceptions, IntegrityError) with distinct messages and let unexpected types propagate to the DRF handler.

Lens 3 · 10

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The empty-queryset case is fixed: a WARNING now fires when no index managers are visible, carrying the document, tool and resolved org, before the delete proceeds. That makes "the filter hid them" distinguishable from "there were none", which was the part that let Redis flags outlive the document behind a 200.

Not narrowing the except Exception, though — and this is a disagreement rather than an omission. Three subsystems are reachable inside that block: Redis via DocumentIndexingService.remove_document_indexing, the object store via PromptStudioFileHelper.delete_for_ide (fsspec, whose backends surface botocore.ClientError among others), and the database. They share no common base class, so any explicit list turns an outage in whichever one was missed into a 500 for a user who previously got a handled 400. ConnectorError/IntegrityError/storage exceptions does not cover the Redis path.

The diagnosability complaint stands on its own and is fixed directly: the handler now logs the exception type, the document id, the tool id and a stack via exc_info=True, where before it logged one interpolated message and no traceback. Happy to revisit the narrowing separately if the 400-to-500 change for infra outages is considered acceptable.

Comment thread backend/prompt_studio/prompt_studio_core_v2/migration_utils.py
athul-rs and others added 2 commits August 11, 2026 14:33
Fixes the failure modes the newly-scoped managers introduced, and corrects the
comments that described the scoping inaccurately.

- get_or_create now goes through _base_manager at both call sites. Django
  applies a manager's filter to the get half but not the create half, so a row
  the org scope hid made get miss and create collide with the unique
  constraint. Both callers already hold org-verified parents.
- mark_extraction_status: the internal endpoint returns 500 instead of
  200 {"success": false}. The worker never read the body, so a failed write
  was silently dropped and every later Answer Prompt re-ran the full X2Text
  extraction. The bare `except Exception` is narrowed, and the worker logs at
  ERROR with the cost spelled out.
- make_profile_default validates default_profile up front: a missing key was a
  KeyError and a non-UUID value a Django ValidationError, both 500s next to
  the 404 this action already returned. The write is now
  save(update_fields=["is_default"]) so it cannot clobber a concurrent edit
  from its pre-transaction snapshot.
- get_output_for_tool_default and latest_outputs_by_keys validate tool_id as a
  UUID (a non-UUID raised while the query was built, giving a 500) and refuse
  to run with no organization in context, which compiled to
  `organization_id IS NULL` and served a blank project that has real outputs.
- delete_for_ide warns when no index managers are visible: the delete
  otherwise returned 200 while leaving Redis indexing flags behind. Its
  handler keeps the broad catch — Redis, the object store and the database
  are all in play and share no base class — but now logs type, document and
  stack.
- The lazy summarize migration distinguishes "profile is filtered out" from
  "profile does not exist"; the first never self-heals and no longer hides
  behind the same INFO line.
- OrgAwareManager logs when it fails open on an exception. That arm catches
  more than its stated cause: StateStore.get raises RuntimeError for any
  unrecognised CONCURRENCY_MODE. The org-is-None arm stays silent — it is the
  normal state for every Celery query.
- Comment corrections: "six internal viewsets" undercounted a ~35-call-site
  surface; "custom @action methods never call filter_queryset()" is wrong,
  since get_object() does filter and it is the raw .objects lookups beside it
  that do not; the pin comment overstated what the test proves and omitted
  that org_filter_paths outranks the pin at the view layer; and seven
  backward-compat comments still described the header as optional after the
  helper began failing closed.

Tests: the nullable-hop assertion now covers the terminal organization FK,
which is the nullable one on every pin, with the exemptions written down.
make_profile_default is exercised through the view — allow path and rejection
path. Mutation-tested: the rejection case fails only on clear-then-resolve
*without* the transaction, which is what the code did before; reverting the
ordering alone is safe because the 404 rolls the clear back, so the test
docstring says that rather than the reviewer's stronger claim.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
@sonarqubecloud

Copy link
Copy Markdown

@github-actions

Copy link
Copy Markdown
Contributor

Unstract test results

Per-group results

Status Group Tier Passed Failed Errors Skipped Duration (s)
e2e-api-deployment e2e 3 0 0 0 16.1
e2e-coowners e2e 1 0 0 0 1.1
e2e-etl e2e 1 0 0 0 7.9
e2e-login e2e 2 0 0 0 0.9
e2e-prompt-studio e2e 1 0 0 0 4.3
e2e-smoke e2e 2 0 0 0 0.8
e2e-workflow e2e 1 0 0 0 16.3
integration-backend integration 289 0 0 26 43.2
integration-connectors integration 1 0 0 7 7.7
integration-workers integration 140 0 0 1 48.5
unit-backend unit 1013 0 0 1 39.1
unit-connectors unit 63 0 0 0 9.8
unit-core unit 33 0 0 0 1.3
unit-platform-service unit 15 0 0 0 2.6
unit-rig unit 117 0 0 0 5.4
unit-sdk1 unit 480 0 0 0 23.2
unit-workers unit 1335 0 0 1 94.9
TOTAL 3497 0 0 36 323.0

Critical paths

⚠️ Critical paths not yet covered

  • workflow-execution-fan-out — Multi-file workflow execution fans out to file-processing workers and rejoins. (declared coverage: no groups declared)
✅ Covered critical paths
  • auth-login — covered by e2e-login
  • adapter-register-llm — covered by integration-backend
  • workflow-author — covered by integration-backend
  • co-owner-manage — covered by integration-backend, e2e-coowners
  • workflow-create-execute — covered by e2e-workflow
  • api-deployment-provision — covered by integration-backend
  • api-deployment-auth — covered by integration-backend
  • api-deployment-run — covered by e2e-api-deployment
  • mcp-server-auth — covered by integration-backend
  • mcp-platform-auth — covered by integration-backend
  • prompt-studio-author — covered by integration-backend
  • prompt-studio-fetch-response — covered by e2e-prompt-studio
  • connector-register-test — covered by integration-backend
  • pipeline-etl-execute — covered by e2e-etl
  • usage-aggregate-read — covered by integration-backend
  • usage-token-tracking — covered by e2e-api-deployment
  • callback-result-delivery — covered by e2e-api-deployment

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants