UN-3815 [FIX] Apply organization scoping to prompt-studio child models - #2213
Open
athul-rs wants to merge 6 commits into
Open
UN-3815 [FIX] Apply organization scoping to prompt-studio child models#2213athul-rs wants to merge 6 commits into
athul-rs wants to merge 6 commits into
Conversation
Contributor
Summary by CodeRabbit
WalkthroughThe 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. ChangesOrganization isolation and access control
Endpoint and transaction changes
Estimated code review effort: 4 (Complex) | ~45 minutes Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
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. Comment |
Contributor
|
| Filename | Overview |
|---|---|
| backend/utils/models/org_path_discovery.py | Pins stable organization relationship paths for the five newly scoped Prompt Studio child models. |
| backend/prompt_studio/prompt_profile_manager_v2/models.py | Makes profile queries organization-aware while preserving the existing user-sharing manager behavior. |
| backend/prompt_studio/prompt_studio_core_v2/views.py | Tool-scopes request-supplied document and profile identifiers and makes default-profile replacement atomic. |
| backend/utils/organization_utils.py | Changes the organization-scoping helper to fail closed when organization context is absent or unresolvable. |
| backend/prompt_studio/tests/test_cross_org_isolation.py | Adds regression coverage for cross-organization manager access, explicit parent scoping, worker context, and the removed route. |
Flowchart
%%{init: {'theme': 'neutral'}}%%
flowchart LR
Request[Request or worker context] --> Org[Current organization]
Org --> Manager[OrgAwareManager]
Manager --> Path[get_org_path]
Path --> Override[Stable pinned FK path]
Override --> Children[Prompt Studio child queryset]
Children --> Tool[Organization-owned CustomTool or adapter]
Action[Custom action with request-supplied ID] --> ParentScope[Explicit tool or organization predicate]
ParentScope --> Children
Reviews (6): Last reviewed commit: "Merge branch 'main' into UN-3794-org-sco..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
backend/file_management/views.py (1)
28-33: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winStale 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 winNo
tearDownto resetUserContextafter mutating tests.
test_no_org_context_is_unfiltered(Line 143) sets the org identifier toNone, andtest_worker_context_sees_its_own_org(Line 151) sets it to org B; neither is restored. SinceUserContextlooks like process-level/thread-local state (not something Django's transactionalTestCaserolls 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
📒 Files selected for processing (15)
backend/file_management/urls.pybackend/file_management/views.pybackend/prompt_studio/prompt_profile_manager_v2/models.pybackend/prompt_studio/prompt_studio_core_v2/migration_utils.pybackend/prompt_studio/prompt_studio_core_v2/views.pybackend/prompt_studio/prompt_studio_document_manager_v2/models.pybackend/prompt_studio/prompt_studio_index_manager_v2/models.pybackend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.pybackend/prompt_studio/prompt_studio_output_manager_v2/models.pybackend/prompt_studio/prompt_studio_output_manager_v2/views.pybackend/prompt_studio/prompt_studio_v2/models.pybackend/prompt_studio/tests/__init__.pybackend/prompt_studio/tests/test_cross_org_isolation.pybackend/utils/models/org_path_discovery.pybackend/utils/tests/test_org_path_discovery.py
💤 Files with no reviewable changes (1)
- backend/file_management/urls.py
athul-rs
changed the title
athul-rs
marked this pull request as draft
July 27, 2026 19:16
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
force-pushed
the
UN-3794-org-scoping
branch
from
July 29, 2026 09:42
65613a0 to
14f94cd
Compare
athul-rs
marked this pull request as ready for review
July 29, 2026 09:42
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (1)
backend/prompt_studio/tests/test_cross_org_isolation.py (1)
163-207: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExercise the actions, not only their ORM predicates.
These tests recreate the intended lookups directly, so they cannot catch a regression in
delete_for_ideormake_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
📒 Files selected for processing (15)
backend/file_management/urls.pybackend/file_management/views.pybackend/prompt_studio/prompt_profile_manager_v2/models.pybackend/prompt_studio/prompt_studio_core_v2/migration_utils.pybackend/prompt_studio/prompt_studio_core_v2/views.pybackend/prompt_studio/prompt_studio_document_manager_v2/models.pybackend/prompt_studio/prompt_studio_index_manager_v2/models.pybackend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.pybackend/prompt_studio/prompt_studio_output_manager_v2/models.pybackend/prompt_studio/prompt_studio_output_manager_v2/views.pybackend/prompt_studio/prompt_studio_v2/models.pybackend/prompt_studio/tests/__init__.pybackend/prompt_studio/tests/test_cross_org_isolation.pybackend/utils/models/org_path_discovery.pybackend/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
…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.
Contributor
|
@greptileai please review this |
Contributor
Author
|
@greptileai re-review this PR |
|
Contributor
There was a problem hiding this comment.
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 winReturn 400 when
default_profileis missing.
request.data["default_profile"]raisesKeyErrorwhen 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 winMark the class configuration as
ClassVar.Ruff reports RUF012 for both mutable class attributes. Add
ClassVarannotations 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
📒 Files selected for processing (1)
backend/prompt_studio/prompt_studio_core_v2/views.py
Contributor
Unstract test resultsPer-group results
Critical paths
|
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
DocumentManager,IndexManager,PromptStudioOutputManager,ToolStudioPrompt,ProfileManager— via the existingOrgAwareManager.delete_for_ide,get_output_for_tool_default,make_profile_default.Organizationinstead of re-deriving it by BFS on every fresh process.file/deleteroute and action.select_for_update(of=("self",))where the new org filter introduces joins.Why
OrganizationFilterBackendruns infilter_queryset(), which custom DRF@actionmethods never call. These five models have noorganizationFK and used a plainBaseModelManager, so a raw.objectslookup inside a custom action carried no organization predicate at all — roughly 44 call sites relying on the caller to pass a correct id.OrgAwareManageralready existed for exactly this shape but was only wired to one model (ExecutionLog).get_org_pathreturns the shortest FK chain toOrganizationand 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 toOrganizationFilterBackendin production today, independent of anything else in this PR.file/deleteis 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;ProfileManagerModelManagernow extendsOrgAwareManagerinstead ofBaseModelManager. No migration — no manager setsuse_in_migrations, so swappingobjectsserializes nothing (makemigrations --check --dry-runis clean).ORG_PATH_OVERRIDESinorg_path_discovery.py, keyed by model label and checked before BFS, so both consumers (OrgAwareManagerandOrganizationFilterBackend) are frozen on the same value. Each pin is set to the path BFS resolves today, so this changes no behaviour on its own.ProfileManagerpins tovector_store__organization, notprompt_studio_tool__organization: BFS reachesAdapterInstance(which carries the organization FK) beforeCustomTool, andprompt_studio_toolis nullable, so pinning there would drop tool-less profiles.AdapterInstanceis org-owned and the serializer's FK queryset uses the org-scoped default manager, so this scopes to the same organization.get_object_or_404, so a non-matching id is a 404 rather than an unhandledModel.DoesNotExist— whichmiddleware.exception.drf_logging_exc_handlerdoes not map, and would surface as a 500.select_for_update(of=("self",))inprompt_studio_index_helperandmigration_utils: the org filter adds INNER JOINs, and PostgresFOR UPDATEwithoutof=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:
internal_api_auth.py,scheduler/tasks.py,workflow_helper.py), soOrgAwareManagerfilters 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_orgcovers this. Any path that legitimately spans organizations would now return empty; none was found.get_or_createunder a filtering manager. If thegethalf is filtered out while the row exists, thecreatehalf hits the unique constraint. Only reachable across organizations, but the failure mode changes from silently-wrong toIntegrityError.select_for_updatelock scope. Addressed withof=("self",); without it the joins would widen the lock.ProfileManagerandIndexManagerare the two affected call sites.Management commands and shell keep full access:
UserContext.get_organization()returnsNoneoutside a request and the manager fails open, unchanged.test_no_org_context_is_unfilteredpins that.file/deleteremoval 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-runis clean; no manager setsuse_in_migrations, so replacingobjectsdoes 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 againstmainbefore 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).main: identical failure sets (36, all pre-existing inworkflow_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.