From 447d0c940b36d59b669fb137ebd007e263e41b44 Mon Sep 17 00:00:00 2001 From: Athul Date: Mon, 27 Jul 2026 10:02:04 +0530 Subject: [PATCH 1/6] UN-3794 [FIX] Pin organization FK paths instead of relying on BFS order MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- backend/utils/models/org_path_discovery.py | 32 +++++++++ .../utils/tests/test_org_path_discovery.py | 65 +++++++++++++++++++ 2 files changed, 97 insertions(+) create mode 100644 backend/utils/tests/test_org_path_discovery.py diff --git a/backend/utils/models/org_path_discovery.py b/backend/utils/models/org_path_discovery.py index 9a8011c716..e057fedc29 100644 --- a/backend/utils/models/org_path_discovery.py +++ b/backend/utils/models/org_path_discovery.py @@ -19,6 +19,34 @@ _FK_TYPES = (models.ForeignKey, models.OneToOneField) +# Org paths pinned explicitly, checked before BFS. Keyed by model label +# ("app_label.ModelName") so this module stays import-free of the models. +# +# BFS returns the *shortest* path 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 resulting INNER JOIN +# silently drops every row with a NULL — data loss that reads as "missing +# records", not as an error. Pinning freezes the path for both consumers +# (OrgAwareManager and OrganizationFilterBackend); test_org_path_discovery +# asserts each pin still matches BFS and traverses only non-nullable FKs. +ORG_PATH_OVERRIDES: dict[str, str] = { + "prompt_studio_document_manager_v2.DocumentManager": "tool__organization", + "prompt_studio_index_manager_v2.IndexManager": ( + "document_manager__tool__organization" + ), + "prompt_studio_output_manager_v2.PromptStudioOutputManager": ( + "tool_id__organization" + ), + # ToolStudioPrompt.tool_id is nullable — prompts orphaned from their tool + # are excluded. This is the path already in force, pinned as-is rather + # than changed under a security fix. + "prompt_studio_v2.ToolStudioPrompt": "tool_id__organization", + # Deliberately not prompt_studio_tool__organization: that FK is nullable, + # so it would drop tool-less profiles. vector_store is non-null and + # AdapterInstance is org-owned, so it scopes to the same organization. + "prompt_profile_manager_v2.ProfileManager": "vector_store__organization", +} + def get_org_path(model: type) -> str | None: """Get the cached FK path from a model to Organization. @@ -26,6 +54,10 @@ def get_org_path(model: type) -> str | None: Returns the ORM lookup path (e.g., "wf_execution__workflow__organization") or None if no path exists. """ + pinned = ORG_PATH_OVERRIDES.get(model._meta.label) + if pinned: + return pinned + if model in _org_path_cache: return _org_path_cache[model] diff --git a/backend/utils/tests/test_org_path_discovery.py b/backend/utils/tests/test_org_path_discovery.py new file mode 100644 index 0000000000..958c393139 --- /dev/null +++ b/backend/utils/tests/test_org_path_discovery.py @@ -0,0 +1,65 @@ +"""Guards for the pinned organization FK paths. + +These paths decide how every org-scoped queryset is filtered, at both the +manager layer (OrgAwareManager) and the view layer (OrganizationFilterBackend). +A path that changes silently is a cross-tenant leak or silent row loss, so both +properties are asserted here rather than left to review. + +No DB access — path discovery walks the model metadata only. +""" + +import pytest +from django.apps import apps +from utils.models.org_path_discovery import ( + ORG_PATH_OVERRIDES, + _discover_org_path, + get_org_path, +) + +PINS = sorted(ORG_PATH_OVERRIDES.items()) + +# Nullable hops accepted as pre-existing behaviour, not introduced here. +# Rows with a NULL value on these FKs are excluded from every org-scoped +# query. Anything not listed must be non-nullable. +KNOWN_NULLABLE_HOPS = {("prompt_studio_v2.ToolStudioPrompt", "tool_id")} + + +@pytest.mark.parametrize("label,expected", PINS) +def test_pin_is_returned(label, expected): + """get_org_path serves the pin, bypassing BFS.""" + assert get_org_path(apps.get_model(label)) == expected + + +@pytest.mark.parametrize("label,expected", PINS) +def test_pin_matches_discovery(label, expected): + """The pin still agrees with what BFS would pick. + + Fails when a field reorder or a new FK changes the shortest path. That is + the signal to re-derive the pin deliberately, not to update this constant + to make CI green. + """ + assert _discover_org_path(apps.get_model(label)) == expected + + +@pytest.mark.parametrize("label,expected", PINS) +def test_pin_traverses_only_non_nullable_fks(label, expected): + """Every hop before `organization` must be non-nullable. + + Django turns a positive filter over a nullable FK into an INNER JOIN, which + drops rows whose FK is NULL. On an org filter that is invisible data loss. + """ + model = apps.get_model(label) + hops = expected.split("__") + + for hop in hops[:-1]: + field = model._meta.get_field(hop) + assert not field.null or (label, hop) in KNOWN_NULLABLE_HOPS, ( + f"{model._meta.label}.{hop} is nullable: this pin drops every row " + f"with a NULL {hop}. Pick a non-nullable path or add it to " + f"KNOWN_NULLABLE_HOPS with a reason." + ) + model = field.related_model + + # Final hop must actually be the Organization FK. + org_field = model._meta.get_field(hops[-1]) + assert org_field.related_model._meta.label == "account_v2.Organization" From 09d320b99b4ee33535f5e3d1cea4563e872df722 Mon Sep 17 00:00:00 2001 From: Athul Date: Mon, 27 Jul 2026 10:02:25 +0530 Subject: [PATCH 2/6] UN-3794 [FIX] Apply organization scoping to prompt-studio child models MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- backend/file_management/urls.py | 6 - backend/file_management/views.py | 39 +--- .../prompt_profile_manager_v2/models.py | 5 +- .../prompt_studio_core_v2/migration_utils.py | 8 +- .../prompt_studio_core_v2/views.py | 18 +- .../models.py | 5 + .../prompt_studio_index_manager_v2/models.py | 5 + .../prompt_studio_index_helper.py | 15 +- .../prompt_studio_output_manager_v2/models.py | 5 + .../prompt_studio_output_manager_v2/views.py | 7 +- .../prompt_studio/prompt_studio_v2/models.py | 6 + backend/prompt_studio/tests/__init__.py | 0 .../tests/test_cross_org_isolation.py | 190 ++++++++++++++++++ 13 files changed, 250 insertions(+), 59 deletions(-) create mode 100644 backend/prompt_studio/tests/__init__.py create mode 100644 backend/prompt_studio/tests/test_cross_org_isolation.py diff --git a/backend/file_management/urls.py b/backend/file_management/urls.py index 8b0ae2dcf9..995ff3f4d7 100644 --- a/backend/file_management/urls.py +++ b/backend/file_management/urls.py @@ -34,16 +34,10 @@ "get": "list_ide", } ) -file_delete = FileManagementViewSet.as_view( - { - "get": "delete", - } -) urlpatterns = format_suffix_patterns( [ path("file", file_list, name="file-list"), path("file/download", file_downlaod, name="download"), path("file/upload", file_upload, name="upload"), - path("file/delete", file_delete, name="delete"), ] ) diff --git a/backend/file_management/views.py b/backend/file_management/views.py index b2875d4251..01d5a87662 100644 --- a/backend/file_management/views.py +++ b/backend/file_management/views.py @@ -4,12 +4,10 @@ from connector_v2.models import ConnectorInstance from django.http import HttpRequest from oauth2client.client import HttpAccessTokenRefreshError -from prompt_studio.prompt_studio_document_manager_v2.models import DocumentManager -from rest_framework import serializers, status, viewsets +from rest_framework import serializers, viewsets from rest_framework.decorators import action from rest_framework.response import Response from rest_framework.versioning import URLPathVersioning -from utils.user_session import UserSessionUtils from file_management.exceptions import ( ConnectorInstanceNotFound, @@ -18,13 +16,11 @@ ) from file_management.file_management_helper import FileManagerHelper from file_management.serializer import ( - FileInfoIdeSerializer, FileInfoSerializer, FileListRequestSerializer, FileUploadSerializer, ) from unstract.connectors.exceptions import ConnectorError -from unstract.connectors.filesystems.local_storage.local_storage import LocalStorageFS logger = logging.getLogger(__name__) @@ -99,36 +95,3 @@ def upload(self, request: HttpRequest) -> Response: logger.info(f"Uploading file: {file_name}" if file_name else "Uploading file") FileManagerHelper.upload_file(file_system, path, uploaded_file, file_name) return Response({"message": "Files are uploaded successfully!"}) - - @action(detail=True, methods=["get"]) - def delete(self, request: HttpRequest) -> Response: - serializer = FileInfoIdeSerializer(data=request.GET) - serializer.is_valid(raise_exception=True) - document_id: str = serializer.validated_data.get("document_id") - document: DocumentManager = DocumentManager.objects.get(pk=document_id) - file_name: str = document.document_name - tool_id: str = serializer.validated_data.get("tool_id") - file_path = FileManagerHelper.handle_sub_directory_for_tenants( - UserSessionUtils.get_organization_id(request), - is_create=False, - user_id=request.user.user_id, - tool_id=tool_id, - ) - path = file_path - file_system = LocalStorageFS(settings={"path": path}) - try: - # Delete the document record - document.delete() - - # Delete the file - FileManagerHelper.delete_file(file_system, path, file_name) - return Response( - {"data": "File deleted succesfully."}, - status=status.HTTP_200_OK, - ) - except Exception as exc: - logger.error(f"Exception thrown from file deletion, error {exc}") - return Response( - {"data": "File deletion failed."}, - status=status.HTTP_400_BAD_REQUEST, - ) diff --git a/backend/prompt_studio/prompt_profile_manager_v2/models.py b/backend/prompt_studio/prompt_profile_manager_v2/models.py index 10a234f462..fbd7437656 100644 --- a/backend/prompt_studio/prompt_profile_manager_v2/models.py +++ b/backend/prompt_studio/prompt_profile_manager_v2/models.py @@ -5,14 +5,15 @@ from django.db import models from django.db.models import Q from tenant_account_v2.organization_member_service import OrganizationMemberService -from utils.models.base_model import BaseModel, BaseModelManager +from utils.models.base_model import BaseModel +from utils.models.org_aware_manager import OrgAwareManager from utils.user_context import UserContext from prompt_studio.prompt_studio_core_v2.exceptions import DefaultProfileError from prompt_studio.prompt_studio_core_v2.models import CustomTool -class ProfileManagerModelManager(BaseModelManager): +class ProfileManagerModelManager(OrgAwareManager): def for_user(self, user): """Read visibility: profile's own share fields OR parent CustomTool sharing. diff --git a/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py b/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py index b9236d04f0..b21a96dd08 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py +++ b/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py @@ -60,9 +60,11 @@ def migrate_tool_to_adapter_based( # Re-fetch the summarize profile with lock within transaction try: - summarize_profile = ProfileManager.objects.select_for_update().get( - prompt_studio_tool=tool_instance, is_summarize_llm=True - ) + # of=("self",): the org-scoped manager joins through + # AdapterInstance, which would otherwise be locked too. + summarize_profile = ProfileManager.objects.select_for_update( + of=("self",) + ).get(prompt_studio_tool=tool_instance, is_summarize_llm=True) except ObjectDoesNotExist: logger.info( f"No summarize profile found for tool {tool_instance.tool_id}, skipping migration" diff --git a/backend/prompt_studio/prompt_studio_core_v2/views.py b/backend/prompt_studio/prompt_studio_core_v2/views.py index ae11da451a..c4a63b0b3c 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/views.py @@ -14,6 +14,7 @@ from django.db import IntegrityError from django.db.models import Count, OuterRef, QuerySet, Subquery from django.http import HttpRequest, HttpResponse +from django.shortcuts import get_object_or_404 from django.utils import timezone from file_management.constants import FileInformationKey as FileKey from file_management.exceptions import FileNotFound @@ -442,7 +443,14 @@ def make_profile_default(self, request: HttpRequest, pk: Any = None) -> Response is_default=False ) - profile_manager = ProfileManager.objects.get(pk=request.data["default_profile"]) + # The id comes straight from the request body, so scope it to the same + # tool the de-dup update above ran against. get_object_or_404 keeps a + # non-matching id a 404 rather than an unhandled DoesNotExist. + profile_manager = get_object_or_404( + ProfileManager, + pk=request.data["default_profile"], + prompt_studio_tool=prompt_tool, + ) profile_manager.is_default = True profile_manager.save() @@ -1180,7 +1188,13 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response: document_id: str = serializer.validated_data.get(ToolStudioPromptKeys.DOCUMENT_ID) org_id = UserSessionUtils.get_organization_id(request) user_id = custom_tool.created_by.user_id - document: DocumentManager = DocumentManager.objects.get(pk=document_id) + # 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 + ) try: # Delete indexed flags in redis diff --git a/backend/prompt_studio/prompt_studio_document_manager_v2/models.py b/backend/prompt_studio/prompt_studio_document_manager_v2/models.py index 15c76c5087..1f4ea4e6a7 100644 --- a/backend/prompt_studio/prompt_studio_document_manager_v2/models.py +++ b/backend/prompt_studio/prompt_studio_document_manager_v2/models.py @@ -3,6 +3,7 @@ from account_v2.models import User from django.db import models from utils.models.base_model import BaseModel +from utils.models.org_aware_manager import OrgAwareManager from prompt_studio.prompt_studio_core_v2.models import CustomTool @@ -10,6 +11,10 @@ class DocumentManager(BaseModel): """Model to store the document details.""" + # Org scoping lives here because custom @action methods never call + # filter_queryset(), so OrganizationFilterBackend does not run on them. + objects = OrgAwareManager() + document_id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) document_name = models.CharField( diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/models.py b/backend/prompt_studio/prompt_studio_index_manager_v2/models.py index 60ee406304..9c2372ce95 100644 --- a/backend/prompt_studio/prompt_studio_index_manager_v2/models.py +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/models.py @@ -7,6 +7,7 @@ from django.db.models.signals import pre_delete from django.dispatch import receiver from utils.models.base_model import BaseModel +from utils.models.org_aware_manager import OrgAwareManager from utils.user_context import UserContext from prompt_studio.prompt_profile_manager_v2.models import ProfileManager @@ -21,6 +22,10 @@ class IndexManager(BaseModel): """Model to store the index details.""" + # See DocumentManager.objects — custom @action methods bypass the + # OrganizationFilterBackend, so scoping has to be at the manager. + objects = OrgAwareManager() + index_manager_id = models.UUIDField( primary_key=True, default=uuid.uuid4, editable=False ) diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py b/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py index d17c157865..4d90bffed4 100644 --- a/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py @@ -109,12 +109,15 @@ def mark_extraction_status( # Lock the row (or create an empty one) so concurrent callers # merge into the same dict rather than clobbering each other. - index_manager, created = ( - IndexManager.objects.select_for_update().get_or_create( - document_manager=document, - profile_manager=profile_manager, - defaults={"extraction_status": {}}, - ) + # of=("self",) because the org-scoped manager joins through + # DocumentManager and CustomTool; without it Postgres locks + # rows in those tables too. + index_manager, created = IndexManager.objects.select_for_update( + of=("self",) + ).get_or_create( + document_manager=document, + profile_manager=profile_manager, + defaults={"extraction_status": {}}, ) # Merge in place — update_or_create(defaults=...) would replace diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/models.py b/backend/prompt_studio/prompt_studio_output_manager_v2/models.py index 7b8616968f..1f51c94733 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/models.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/models.py @@ -3,6 +3,7 @@ from account_v2.models import User from django.db import models from utils.models.base_model import BaseModel +from utils.models.org_aware_manager import OrgAwareManager from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.models import CustomTool @@ -16,6 +17,10 @@ class PromptStudioOutputManager(BaseModel): By default the tools will be added to private tool hub. """ + # See DocumentManager.objects — custom @action methods bypass the + # OrganizationFilterBackend, so scoping has to be at the manager. + objects = OrgAwareManager() + prompt_output_id = models.UUIDField( primary_key=True, default=uuid.uuid4, editable=False ) diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py index 44111dc744..9c47814e66 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py @@ -124,9 +124,12 @@ def get_output_for_tool_default(self, request: HttpRequest) -> Response: raise ValidationError(detail=tool_validation_message) try: - # Fetch ToolStudioPrompt records based on tool_id + # Fetch ToolStudioPrompt records based on tool_id. + # Custom actions skip filter_queryset(), so OrganizationFilterBackend + # never runs — scope explicitly to prevent cross-tenant reads. tool_studio_prompts = ToolStudioPrompt.objects.filter( - tool_id=tool_id + tool_id=tool_id, + tool_id__organization=UserContext.get_organization(), ).order_by("sequence_number") except ObjectDoesNotExist: raise ValidationError(detail=tool_not_found) diff --git a/backend/prompt_studio/prompt_studio_v2/models.py b/backend/prompt_studio/prompt_studio_v2/models.py index faaf7b0313..47aa284293 100644 --- a/backend/prompt_studio/prompt_studio_v2/models.py +++ b/backend/prompt_studio/prompt_studio_v2/models.py @@ -4,6 +4,7 @@ from django.db import models from django.utils import timezone from utils.models.base_model import BaseModel +from utils.models.org_aware_manager import OrgAwareManager from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.models import CustomTool @@ -15,6 +16,11 @@ class ToolStudioPrompt(BaseModel): It has Many to one relation with CustomTool for ToolStudio. """ + # See DocumentManager.objects — custom @action methods bypass the + # OrganizationFilterBackend, so scoping has to be at the manager. + # tool_id is nullable, so prompts orphaned from their tool are excluded. + objects = OrgAwareManager() + class EnforceType(models.TextChoices): TEXT = "text", "Response sent as Text" NUMBER = "number", "Response sent as number" diff --git a/backend/prompt_studio/tests/__init__.py b/backend/prompt_studio/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/prompt_studio/tests/test_cross_org_isolation.py b/backend/prompt_studio/tests/test_cross_org_isolation.py new file mode 100644 index 0000000000..02137f1dc6 --- /dev/null +++ b/backend/prompt_studio/tests/test_cross_org_isolation.py @@ -0,0 +1,190 @@ +"""Organization isolation for the prompt-studio child models. + +Custom DRF ``@action`` methods never call ``filter_queryset()``, so +``OrganizationFilterBackend`` does not run on them and a raw +``.objects.get()/filter()`` inside one is not org-scoped. These tests pin the +controls that cover that gap: org scoping on the managers, plus explicit +scoping where an id arrives directly from the request. + +Shape of each case: act as org A, pass an org B id, assert the call is +refused and org B's row is untouched. +""" + +import secrets + +import pytest +from account_v2.models import Organization, User +from adapter_processor_v2.models import AdapterInstance +from django.test import TestCase +from django.urls import NoReverseMatch, reverse +from utils.user_context import UserContext + +from prompt_studio.prompt_profile_manager_v2.models import ProfileManager +from prompt_studio.prompt_studio_core_v2.models import CustomTool +from prompt_studio.prompt_studio_document_manager_v2.models import DocumentManager +from prompt_studio.prompt_studio_index_manager_v2.models import IndexManager +from prompt_studio.prompt_studio_output_manager_v2.models import ( + PromptStudioOutputManager, +) +from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt + + +class OrgFixture: + """One organization with a fully populated prompt-studio object graph.""" + + def __init__(self, slug: str): + self.org = Organization.objects.create( + name=slug, display_name=slug, organization_id=slug + ) + UserContext.set_organization_identifier(slug) + + self.user = User.objects.create_user( + username=f"{slug}@example.com", + email=f"{slug}@example.com", + password=secrets.token_urlsafe(), + ) + self.tool = CustomTool.objects.create( + tool_name=f"tool-{slug}", + description="isolation test tool", + organization=self.org, + created_by=self.user, + ) + adapter = self._adapter(slug) + self.profile = ProfileManager.objects.create( + profile_name=f"profile-{slug}", + vector_store=adapter, + embedding_model=adapter, + llm=adapter, + x2text=adapter, + chunk_size=0, + chunk_overlap=0, + section="Default", + retrieval_strategy="simple", + similarity_top_k=3, + prompt_studio_tool=self.tool, + is_default=True, + created_by=self.user, + ) + self.document = DocumentManager.objects.create( + document_name=f"doc-{slug}.pdf", tool=self.tool, created_by=self.user + ) + self.index = IndexManager.objects.create( + document_manager=self.document, profile_manager=self.profile + ) + self.prompt = ToolStudioPrompt.objects.create( + prompt_key=f"key_{slug}", prompt="extract", tool_id=self.tool + ) + self.output = PromptStudioOutputManager.objects.create( + output="secret", + prompt_id=self.prompt, + document_manager=self.document, + profile_manager=self.profile, + tool_id=self.tool, + ) + + def _adapter(self, slug: str) -> AdapterInstance: + return AdapterInstance.objects.create( + adapter_name=f"adapter-{slug}", + adapter_id="openai|test", + adapter_type="LLM", + adapter_metadata={}, + organization=self.org, + created_by=self.user, + ) + + +@pytest.mark.django_db +class CrossOrgIsolationTest(TestCase): + """Org A must not reach org B's prompt-studio rows through any manager.""" + + 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) + + # --- manager scoping: the default-deny layer (A-1) -------------------- + + def test_document_of_other_org_is_not_gettable(self): + """A document id from another org must not resolve.""" + with self.assertRaises(DocumentManager.DoesNotExist): + DocumentManager.objects.get(pk=self.b.document.document_id) + + def test_prompt_of_other_org_is_not_listable(self): + """Prompts must not be listable by another org's tool id.""" + assert not ToolStudioPrompt.objects.filter(tool_id=self.b.tool).exists() + + def test_output_of_other_org_is_not_listable(self): + assert not PromptStudioOutputManager.objects.filter( + tool_id=self.b.tool + ).exists() + + def test_index_of_other_org_is_not_listable(self): + assert not IndexManager.objects.filter( + document_manager=self.b.document + ).exists() + + def test_profile_of_other_org_is_not_gettable(self): + """``make_profile_default`` takes this id straight from the body.""" + with self.assertRaises(ProfileManager.DoesNotExist): + ProfileManager.objects.get(pk=self.b.profile.profile_id) + + # --- same-org access must still work ---------------------------------- + + def test_own_org_rows_remain_visible(self): + assert DocumentManager.objects.get(pk=self.a.document.document_id) + assert ProfileManager.objects.get(pk=self.a.profile.profile_id) + assert ToolStudioPrompt.objects.filter(tool_id=self.a.tool).exists() + assert PromptStudioOutputManager.objects.filter(tool_id=self.a.tool).exists() + assert IndexManager.objects.filter(document_manager=self.a.document).exists() + + def test_no_org_context_is_unfiltered(self): + """Management commands and shell keep full access (fail-open).""" + UserContext.set_organization_identifier(None) + assert DocumentManager.objects.filter( + pk=self.b.document.document_id + ).exists() + + def test_worker_context_sees_its_own_org(self): + """B1 — workers do run with org context set, so the manager filters + there too. Indexing must still find its own org's rows.""" + UserContext.set_organization_identifier(self.b.org.organization_id) + assert IndexManager.objects.filter( + document_manager=self.b.document + ).exists() + assert DocumentManager.objects.get(pk=self.b.document.document_id) + + # --- explicit scoping at the reported call sites (A-3, A-5) ----------- + + def test_delete_for_ide_lookup_is_tool_scoped(self): + """A doc id from another tool in the *same* org is refused too.""" + sibling = CustomTool.objects.create( + tool_name="sibling", + description="second tool, same org", + organization=self.a.org, + created_by=self.a.user, + ) + with self.assertRaises(DocumentManager.DoesNotExist): + DocumentManager.objects.get( + pk=self.a.document.document_id, tool=sibling + ) + + def test_make_profile_default_lookup_is_tool_scoped(self): + """This lookup runs after ``get_object()`` has already passed authz on + the caller's own tool, so org scope alone does not constrain it.""" + with self.assertRaises(ProfileManager.DoesNotExist): + ProfileManager.objects.get( + pk=self.b.profile.profile_id, prompt_studio_tool=self.a.tool + ) + # Victim's default flag untouched. + UserContext.set_organization_identifier(self.b.org.organization_id) + assert ProfileManager.objects.get(pk=self.b.profile.profile_id).is_default + + # --- A-4: the dead, state-changing-over-GET route is gone ------------- + + def test_file_delete_route_removed(self): + """Removed rather than fixed: no caller, and it deleted over GET.""" + # Sibling route still resolves, so a naming change can't fake a pass. + assert reverse("tenant:upload").endswith("/file/upload") + with pytest.raises(NoReverseMatch): + reverse("tenant:delete") From 14f94cdb3212b5ab60e8685a41b56865a07477a8 Mon Sep 17 00:00:00 2001 From: Athul Date: Wed, 29 Jul 2026 15:07:21 +0530 Subject: [PATCH 3/6] UN-3815 [FIX] Resolve the target profile before clearing existing defaults 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. --- backend/file_management/views.py | 2 +- .../prompt_studio_core_v2/views.py | 24 ++++++++++-------- .../tests/test_cross_org_isolation.py | 25 +++++++++++++++++++ 3 files changed, 40 insertions(+), 11 deletions(-) diff --git a/backend/file_management/views.py b/backend/file_management/views.py index 01d5a87662..2a4d684763 100644 --- a/backend/file_management/views.py +++ b/backend/file_management/views.py @@ -28,7 +28,7 @@ class FileManagementViewSet(viewsets.ModelViewSet): """FileManagement view. - Handles GET,POST,PUT,PATCH and DELETE + Handles GET, POST, PUT and PATCH """ versioning_class = URLPathVersioning diff --git a/backend/prompt_studio/prompt_studio_core_v2/views.py b/backend/prompt_studio/prompt_studio_core_v2/views.py index c4a63b0b3c..7af887ba13 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/views.py @@ -11,7 +11,7 @@ from api_v2.models import APIDeployment from celery import signature from celery.result import AsyncResult -from django.db import IntegrityError +from django.db import IntegrityError, transaction from django.db.models import Count, OuterRef, QuerySet, Subquery from django.http import HttpRequest, HttpResponse from django.shortcuts import get_object_or_404 @@ -439,20 +439,24 @@ def make_profile_default(self, request: HttpRequest, pk: Any = None) -> Response self.get_object() ) # Assuming you have a get_object method in your viewset - ProfileManager.objects.filter(prompt_studio_tool=prompt_tool).update( - is_default=False - ) - - # The id comes straight from the request body, so scope it to the same - # tool the de-dup update above ran against. get_object_or_404 keeps a - # non-matching id a 404 rather than an unhandled DoesNotExist. + # Resolve the target before clearing anything: the id comes straight + # from the request body, and clearing first would leave the tool with no + # default at all when it does not match. Scoped to the same tool the + # caller already passed authz on, so another tool's id is a 404. profile_manager = get_object_or_404( ProfileManager, pk=request.data["default_profile"], prompt_studio_tool=prompt_tool, ) - profile_manager.is_default = True - profile_manager.save() + + # Both writes in one transaction so a failure between them cannot leave + # the tool with zero defaults or two. + with transaction.atomic(): + ProfileManager.objects.filter(prompt_studio_tool=prompt_tool).update( + is_default=False + ) + profile_manager.is_default = True + profile_manager.save() return Response( status=status.HTTP_200_OK, diff --git a/backend/prompt_studio/tests/test_cross_org_isolation.py b/backend/prompt_studio/tests/test_cross_org_isolation.py index 02137f1dc6..eac6a72d02 100644 --- a/backend/prompt_studio/tests/test_cross_org_isolation.py +++ b/backend/prompt_studio/tests/test_cross_org_isolation.py @@ -103,6 +103,12 @@ def setUp(self) -> None: # End state: acting as org A, as a request would. UserContext.set_organization_identifier(self.a.org.organization_id) + def tearDown(self) -> None: + # UserContext is thread-local, not DB state, so TestCase's transaction + # rollback does not clear it. Tests below deliberately switch org and + # would otherwise leak that into whatever runs next. + UserContext.set_organization_identifier(None) + # --- manager scoping: the default-deny layer (A-1) -------------------- def test_document_of_other_org_is_not_gettable(self): @@ -169,6 +175,25 @@ def test_delete_for_ide_lookup_is_tool_scoped(self): pk=self.a.document.document_id, tool=sibling ) + def test_rejected_default_leaves_the_existing_default_intact(self): + """A non-matching id must not clear the tool's current default. + + The de-dup update runs against every profile on the tool, so resolving + the target after it would leave the tool with no default at all when the + id turns out to be someone else's. + """ + assert ProfileManager.objects.get(pk=self.a.profile.profile_id).is_default + + with self.assertRaises(ProfileManager.DoesNotExist): + ProfileManager.objects.get( + pk=self.b.profile.profile_id, prompt_studio_tool=self.a.tool + ) + + self.a.profile.refresh_from_db() + assert self.a.profile.is_default, ( + "the tool lost its default profile while rejecting another org's id" + ) + def test_make_profile_default_lookup_is_tool_scoped(self): """This lookup runs after ``get_object()`` has already passed authz on the caller's own tool, so org scope alone does not constrain it.""" From 18a53f9f917cccf6ff2ef2661c088e37b854296a Mon Sep 17 00:00:00 2001 From: Athul Date: Fri, 31 Jul 2026 00:41:01 +0530 Subject: [PATCH 4/6] UN-3815 [FIX] Drop unreachable exception handler in get_output_for_tool_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. --- .../prompt_studio_output_manager_v2/views.py | 25 ++++++++++--------- 1 file changed, 13 insertions(+), 12 deletions(-) diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py index 9c47814e66..4e0a746f3f 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py @@ -1,7 +1,6 @@ import logging from typing import Any -from django.core.exceptions import ObjectDoesNotExist from django.db.models import QuerySet from django.http import HttpRequest from rest_framework import status, viewsets @@ -119,20 +118,22 @@ def get_output_for_tool_default(self, request: HttpRequest) -> Response: tool_id = request.GET.get("tool_id") document_manager_id = request.GET.get("document_manager") tool_validation_message = PromptOutputManagerErrorMessage.TOOL_VALIDATION - tool_not_found = PromptOutputManagerErrorMessage.TOOL_NOT_FOUND if not tool_id: raise ValidationError(detail=tool_validation_message) - try: - # Fetch ToolStudioPrompt records based on tool_id. - # Custom actions skip filter_queryset(), so OrganizationFilterBackend - # never runs — scope explicitly to prevent cross-tenant reads. - tool_studio_prompts = ToolStudioPrompt.objects.filter( - tool_id=tool_id, - tool_id__organization=UserContext.get_organization(), - ).order_by("sequence_number") - except ObjectDoesNotExist: - raise ValidationError(detail=tool_not_found) + # Fetch ToolStudioPrompt records based on tool_id. + # Custom actions skip filter_queryset(), so OrganizationFilterBackend + # never runs — scope explicitly to prevent cross-tenant reads. + # + # No exception handling here: filter() does not raise for a missing or + # out-of-org tool, it returns empty. Empty is also the correct result + # for a tool that simply has no prompts yet, which is the normal state + # of a newly created project — so this stays a 200 with an empty body + # rather than a validation error. + tool_studio_prompts = ToolStudioPrompt.objects.filter( + tool_id=tool_id, + tool_id__organization=UserContext.get_organization(), + ).order_by("sequence_number") # Invoke helper method to frame and fetch default response. result: dict[str, Any] = OutputManagerHelper.fetch_default_output_response( From 14b7e68db0d0d5746b266832bc67eff611877c73 Mon Sep 17 00:00:00 2001 From: Athul Date: Fri, 31 Jul 2026 00:45:32 +0530 Subject: [PATCH 5/6] UN-3815 [FIX] Fail closed when organization context is absent MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- backend/utils/organization_utils.py | 46 ++++++++--- .../utils/tests/test_organization_scoping.py | 80 +++++++++++++++++++ 2 files changed, 114 insertions(+), 12 deletions(-) create mode 100644 backend/utils/tests/test_organization_scoping.py diff --git a/backend/utils/organization_utils.py b/backend/utils/organization_utils.py index 15053684bf..9af586954d 100644 --- a/backend/utils/organization_utils.py +++ b/backend/utils/organization_utils.py @@ -72,7 +72,23 @@ def get_organization_context(organization: Organization) -> dict[str, Any]: def filter_queryset_by_organization(queryset, request, organization_field="organization"): - """Filter a Django queryset by organization context from request. + """Filter a Django queryset by the request's organization context. + + Fails closed. Six internal viewsets set ``skip_org_filter = True``, which + disables OrganizationFilterBackend, leaving this function as their only + tenant boundary — so returning the queryset unfiltered when there is no + organization context hands back every organization's rows. A scoping helper + returns nothing when it cannot scope, never everything. + + The absent-header case is not exotic: ``InternalAPIAuthMiddleware`` logs a + warning and continues when ``X-Organization-ID`` is missing, so any caller + holding the internal service key reaches here without context simply by + omitting it. + + Note for callers that genuinely span organizations — the leader-elected + reaper is one — query the model directly rather than routing through here. + ``recover_stuck_pg_executions`` already does, which is why failing closed + does not affect it. Args: queryset: Django QuerySet to filter @@ -80,16 +96,22 @@ def filter_queryset_by_organization(queryset, request, organization_field="organ organization_field: Field name for organization relationship (default: 'organization') Returns: - Filtered queryset or empty queryset if organization not found + The queryset filtered to the request's organization, or an empty + queryset when the organization is absent or unresolvable. """ org_id = getattr(request, "organization_id", None) - if org_id: - organization = resolve_organization(org_id, raise_on_not_found=False) - if organization: - # Use dynamic field lookup - filter_kwargs = {organization_field: organization} - return queryset.filter(**filter_kwargs) - else: - # Return empty queryset if organization not found - return queryset.none() - return queryset + 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", ""), + ) + return queryset.none() + + organization = resolve_organization(org_id, raise_on_not_found=False) + if not organization: + logger.warning("Organization %s not found; returning no rows.", org_id) + return queryset.none() + + return queryset.filter(**{organization_field: organization}) diff --git a/backend/utils/tests/test_organization_scoping.py b/backend/utils/tests/test_organization_scoping.py new file mode 100644 index 0000000000..ce6e4b89f0 --- /dev/null +++ b/backend/utils/tests/test_organization_scoping.py @@ -0,0 +1,80 @@ +"""``filter_queryset_by_organization`` must fail closed. + +Six internal viewsets set ``skip_org_filter = True``, which disables +OrganizationFilterBackend and leaves this helper as their only tenant +boundary. Returning the queryset unfiltered when there is no organization +context therefore returns every organization's rows, and the absent-header +case is reachable — the internal auth middleware warns and continues rather +than rejecting. +""" + +import secrets + +import pytest +from account_v2.models import Organization +from django.test import TestCase +from utils.organization_utils import filter_queryset_by_organization +from workflow_manager.workflow_v2.models.workflow import Workflow + + +class _Request: + """Stands in for the request object the helper reads context off.""" + + def __init__(self, organization_id=None, path="/internal/test/"): + if organization_id is not None: + self.organization_id = organization_id + self.path = path + + +@pytest.mark.django_db +class FilterQuerysetByOrganizationTest(TestCase): + def setUp(self) -> None: + self.org_a = self._org("a") + self.org_b = self._org("b") + self.wf_a = Workflow.objects.create( + workflow_name=f"wf-a-{secrets.token_hex(3)}", organization=self.org_a + ) + self.wf_b = Workflow.objects.create( + workflow_name=f"wf-b-{secrets.token_hex(3)}", organization=self.org_b + ) + + def _org(self, tag: str) -> Organization: + slug = f"org-{tag}-{secrets.token_hex(3)}" + return Organization.objects.create( + name=slug, display_name=slug, organization_id=slug + ) + + def _filter(self, request): + # _base_manager, not objects: Workflow's default manager is itself + # org-scoped off UserContext, which would empty the queryset before the + # helper ever ran and make these tests pass for the wrong reason. The + # helper's contract is "given a queryset, scope it", so hand it an + # unscoped one. + return filter_queryset_by_organization(Workflow._base_manager.all(), request) + + def test_missing_org_context_returns_nothing(self): + """The header is optional at the middleware, so this is reachable.""" + assert not self._filter(_Request()).exists() + + def test_falsy_org_context_returns_nothing(self): + for falsy in ("", None): + with self.subTest(organization_id=falsy): + assert not self._filter(_Request(organization_id=falsy)).exists() + + def test_unresolvable_org_returns_nothing(self): + request = _Request(organization_id="does-not-exist") + assert not self._filter(request).exists() + + def test_valid_org_returns_only_its_own_rows(self): + rows = self._filter(_Request(organization_id=self.org_a.organization_id)) + assert list(rows) == [self.wf_a] + + def test_other_org_rows_are_never_included(self): + for org, mine, theirs in ( + (self.org_a, self.wf_a, self.wf_b), + (self.org_b, self.wf_b, self.wf_a), + ): + with self.subTest(org=org.organization_id): + rows = list(self._filter(_Request(organization_id=org.organization_id))) + assert mine in rows + assert theirs not in rows From cc419564b333b563031c49920455fadf2827d630 Mon Sep 17 00:00:00 2001 From: Athul Date: Tue, 11 Aug 2026 14:33:48 +0530 Subject: [PATCH 6/6] UN-3815 [FIX] Address review findings on prompt-studio org scoping MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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) --- backend/notification_v2/internal_views.py | 4 +- backend/pipeline_v2/internal_api_views.py | 4 +- .../prompt_studio_core_v2/internal_views.py | 21 +++- .../prompt_studio_core_v2/migration_utils.py | 26 +++- .../prompt_studio_core_v2/views.py | 52 +++++++- .../models.py | 7 +- .../prompt_studio_index_manager_v2/models.py | 3 +- .../prompt_studio_index_helper.py | 35 ++++-- .../prompt_studio_output_manager_v2/models.py | 3 +- .../output_manager_helper.py | 37 +++--- .../prompt_studio_output_manager_v2/views.py | 66 ++++++++-- .../prompt_studio/prompt_studio_v2/models.py | 3 +- .../tests/test_cross_org_isolation.py | 119 ++++++++++++++---- backend/tool_instance_v2/internal_views.py | 4 +- backend/utils/filters/organization_filter.py | 7 ++ backend/utils/models/org_aware_manager.py | 31 +++-- backend/utils/models/org_path_discovery.py | 19 ++- backend/utils/organization_utils.py | 17 ++- .../utils/tests/test_org_path_discovery.py | 31 ++++- .../utils/tests/test_organization_scoping.py | 13 +- .../file_execution/internal_views.py | 8 +- backend/workflow_manager/internal_views.py | 8 +- backend/workflow_manager/workflow_v2/views.py | 8 +- workers/ide_callback/tasks.py | 9 +- 24 files changed, 414 insertions(+), 121 deletions(-) diff --git a/backend/notification_v2/internal_views.py b/backend/notification_v2/internal_views.py index e352f56db7..dfd97ad149 100644 --- a/backend/notification_v2/internal_views.py +++ b/backend/notification_v2/internal_views.py @@ -42,7 +42,9 @@ class WebhookInternalViewSet(viewsets.ReadOnlyModelViewSet): serializer_class = NotificationSerializer lookup_field = "id" - # Backward compat: remove once all workers pass X-Organization-ID. + # OrganizationFilterBackend is off here; get_queryset() scopes instead, via + # filter_queryset_by_organization. That helper fails closed, so a caller + # without X-Organization-ID gets zero rows. skip_org_filter = True def get_queryset(self): diff --git a/backend/pipeline_v2/internal_api_views.py b/backend/pipeline_v2/internal_api_views.py index 5c1471d717..5681667af4 100644 --- a/backend/pipeline_v2/internal_api_views.py +++ b/backend/pipeline_v2/internal_api_views.py @@ -13,7 +13,9 @@ class PipelineInternalViewSet(ViewSet): - # Backward compat: remove once all workers pass X-Organization-ID. + # OrganizationFilterBackend is off here; scoping runs through + # filter_queryset_by_organization, which fails closed, so a caller without + # X-Organization-ID gets zero rows. skip_org_filter = True def retrieve(self, request, pk=None): diff --git a/backend/prompt_studio/prompt_studio_core_v2/internal_views.py b/backend/prompt_studio/prompt_studio_core_v2/internal_views.py index 3ad3a5db16..b4246e0cf7 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/internal_views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/internal_views.py @@ -211,7 +211,26 @@ def extraction_status(request): extracted=extracted, error_message=error_message, ) - return JsonResponse({"success": success}) + if not success: + # A 200 here is indistinguishable from a write that landed: the + # worker only wraps this call in try/except and never reads the + # body, so the status would be silently dropped and every later + # Answer Prompt would re-run the full extraction. Non-2xx makes + # the worker's existing handler log it. + logger.error( + "extraction_status not recorded for document %s profile %s", + document_id, + profile_manager_id, + ) + return JsonResponse( + { + "success": False, + "error": "Extraction status could not be recorded", + }, + status=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + + return JsonResponse({"success": True}) except Exception as e: logger.exception("extraction_status internal API failed") diff --git a/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py b/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py index b21a96dd08..ea838a06bf 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py +++ b/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py @@ -66,9 +66,29 @@ def migrate_tool_to_adapter_based( of=("self",) ).get(prompt_studio_tool=tool_instance, is_summarize_llm=True) except ObjectDoesNotExist: - logger.info( - f"No summarize profile found for tool {tool_instance.tool_id}, skipping migration" - ) + # Two different situations reach here now that + # ProfileManager.objects is scoped through + # vector_store__organization: the profile genuinely does + # not exist, or it exists and the org filter hid it. The + # second is a misconfiguration that never self-heals — this + # lazy migration re-runs and re-skips on every invocation — + # so it must not share an INFO line with the first. + exists_unscoped = ProfileManager._base_manager.filter( + prompt_studio_tool=tool_instance, is_summarize_llm=True + ).exists() + if exists_unscoped: + logger.error( + "Summarize profile for tool %s exists but is not " + "visible in the current organization context; " + "migration skipped and will keep being skipped.", + tool_instance.tool_id, + ) + else: + logger.info( + "No summarize profile found for tool %s, skipping " + "migration", + tool_instance.tool_id, + ) return False # Check if profile has an LLM adapter diff --git a/backend/prompt_studio/prompt_studio_core_v2/views.py b/backend/prompt_studio/prompt_studio_core_v2/views.py index a5f642ace9..e3d53881ee 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/views.py @@ -27,6 +27,7 @@ from plugins import get_plugin from rest_framework import status, viewsets from rest_framework.decorators import action +from rest_framework.exceptions import ValidationError from rest_framework.request import Request from rest_framework.response import Response from rest_framework.versioning import URLPathVersioning @@ -446,24 +447,39 @@ def make_profile_default(self, request: HttpRequest, pk: Any = None) -> Response self.get_object() ) # Assuming you have a get_object method in your viewset + # Validate before looking anything up. A missing key raised KeyError + # and a non-UUID value raised Django's ValidationError; neither is + # mapped by drf_standardized_errors, so both surfaced as 500s next to + # the 404 this action already returns for a valid-but-unmatched id. + default_profile = request.data.get("default_profile") + if not default_profile: + raise ValidationError(detail="'default_profile' is required.") + try: + default_profile = uuid.UUID(str(default_profile)) + except (ValueError, AttributeError, TypeError): + raise ValidationError(detail="'default_profile' must be a valid UUID.") + # Resolve the target before clearing anything: the id comes straight # from the request body, and clearing first would leave the tool with no # default at all when it does not match. Scoped to the same tool the # caller already passed authz on, so another tool's id is a 404. profile_manager = get_object_or_404( ProfileManager, - pk=request.data["default_profile"], + pk=default_profile, prompt_studio_tool=prompt_tool, ) # Both writes in one transaction so a failure between them cannot leave - # the tool with zero defaults or two. + # the tool with zero defaults or two. update_fields so the second write + # touches one column: profile_manager was read before the transaction + # opened, and a bare save() would write every column from that snapshot + # back over any concurrent edit. with transaction.atomic(): ProfileManager.objects.filter(prompt_studio_tool=prompt_tool).update( is_default=False ) profile_manager.is_default = True - profile_manager.save() + profile_manager.save(update_fields=["is_default"]) return Response( status=status.HTTP_200_OK, @@ -1199,7 +1215,8 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response: org_id = UserSessionUtils.get_organization_id(request) user_id = custom_tool.created_by.user_id # Scope to the tool the caller already passed authz on — tighter than - # org scope, and this action never runs filter_queryset(). + # org scope. self.get_object() above is filtered by the backend, but + # this lookup is a raw .objects query and would not be. # 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( @@ -1209,6 +1226,18 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response: try: # Delete indexed flags in redis index_managers = IndexManager.objects.filter(document_manager=document_id) + if not index_managers.exists(): + # Empty means either "never indexed" or "the org filter hid the + # rows". In the second case the Redis indexing flags outlive the + # document, and a re-upload of the same file is treated as + # already indexed — with a 200 telling the user it all worked. + logger.warning( + "No index managers visible for document %s (tool %s, org %s); " + "deleting without clearing Redis indexing flags.", + document_id, + custom_tool.tool_id, + org_id, + ) for index_manager in index_managers: raw_index_id = index_manager.raw_index_id DocumentIndexingService.remove_document_indexing( @@ -1229,7 +1258,20 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response: status=status.HTTP_200_OK, ) except Exception as exc: - logger.error("Exception thrown from file deletion, error: %s", exc) + # Deliberately broad. Three subsystems are in play — Redis via + # DocumentIndexingService, the object store via + # PromptStudioFileHelper, and the database — and their failures do + # not share a base class, so narrowing to any list turns a + # reachable outage in whichever one was missed into a 500. The + # diagnosability problem was the log line, not the catch: it now + # carries the exception type, the document and a stack. + logger.error( + "File deletion failed for document %s (tool %s): %s", + document_id, + custom_tool.tool_id, + exc, + exc_info=True, + ) return Response( {"data": "File deletion failed."}, status=status.HTTP_400_BAD_REQUEST, diff --git a/backend/prompt_studio/prompt_studio_document_manager_v2/models.py b/backend/prompt_studio/prompt_studio_document_manager_v2/models.py index 1f4ea4e6a7..07ee681c4c 100644 --- a/backend/prompt_studio/prompt_studio_document_manager_v2/models.py +++ b/backend/prompt_studio/prompt_studio_document_manager_v2/models.py @@ -11,8 +11,11 @@ class DocumentManager(BaseModel): """Model to store the document details.""" - # Org scoping lives here because custom @action methods never call - # filter_queryset(), so OrganizationFilterBackend does not run on them. + # Org scoping lives at the manager because OrganizationFilterBackend only + # scopes querysets routed through filter_queryset(). A raw Model.objects + # lookup inside a view bypasses it — including inside a custom @action, + # whose own self.get_object() *is* filtered but whose hand-written queries + # are not. objects = OrgAwareManager() document_id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False) diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/models.py b/backend/prompt_studio/prompt_studio_index_manager_v2/models.py index 9c2372ce95..6d84ed832b 100644 --- a/backend/prompt_studio/prompt_studio_index_manager_v2/models.py +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/models.py @@ -22,8 +22,7 @@ class IndexManager(BaseModel): """Model to store the index details.""" - # See DocumentManager.objects — custom @action methods bypass the - # OrganizationFilterBackend, so scoping has to be at the manager. + # See DocumentManager.objects for why scoping lives at the manager. objects = OrgAwareManager() index_manager_id = models.UUIDField( diff --git a/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py b/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py index 4d90bffed4..1c10bf5b1a 100644 --- a/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py +++ b/backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py @@ -1,7 +1,8 @@ import json import logging -from django.db import transaction +from django.core.exceptions import ImproperlyConfigured +from django.db import DatabaseError, transaction from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.exceptions import IndexingAPIError @@ -109,10 +110,16 @@ def mark_extraction_status( # Lock the row (or create an empty one) so concurrent callers # merge into the same dict rather than clobbering each other. - # of=("self",) because the org-scoped manager joins through - # DocumentManager and CustomTool; without it Postgres locks - # rows in those tables too. - index_manager, created = IndexManager.objects.select_for_update( + # of=("self",) keeps the lock on index_manager rows only. + # + # _base_manager, not objects: Django applies a manager's filter + # to the get half of get_or_create but not to the create half. + # Through the org-scoped manager, a row the filter hides makes + # get miss and create insert, which violates + # unique_document_manager_profile_manager_index. The document + # above was already fetched org-scoped, so the scope is checked + # either way and this only removes the failure mode. + index_manager, created = IndexManager._base_manager.select_for_update( of=("self",) ).get_or_create( document_manager=document, @@ -153,12 +160,24 @@ def mark_extraction_status( return True except DocumentManager.DoesNotExist: - logger.error(f"Document with ID {document_id} does not exist.") + # Now reachable two ways: the row is genuinely gone, or the + # org-scoped manager hid it from this caller. Both mean the status + # was not written, which is what the caller has to act on. + logger.error( + "Document %s not found or not visible in the current " + "organization; extraction status not recorded.", + document_id, + ) return False - except Exception as e: + except (DatabaseError, TypeError, ImproperlyConfigured): + # DatabaseError covers IntegrityError/OperationalError, TypeError a + # malformed extraction_status payload, ImproperlyConfigured a bad + # org path pin. Narrowed from a bare `except Exception` so an + # unexpected type propagates instead of being reported as "no such + # document". logger.exception( - f"Unexpected error marking extraction status for document {document_id}: {e}" + "Failed to mark extraction status for document %s", document_id ) return False diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/models.py b/backend/prompt_studio/prompt_studio_output_manager_v2/models.py index 1f51c94733..9420b87e9d 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/models.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/models.py @@ -17,8 +17,7 @@ class PromptStudioOutputManager(BaseModel): By default the tools will be added to private tool hub. """ - # See DocumentManager.objects — custom @action methods bypass the - # OrganizationFilterBackend, so scoping has to be at the manager. + # See DocumentManager.objects for why scoping lives at the manager. objects = OrgAwareManager() prompt_output_id = models.UUIDField( diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py index 699cacb749..54450201c7 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py @@ -76,21 +76,28 @@ def update_or_create_prompt_output( the instance. """ try: - prompt_output, success = PromptStudioOutputManager.objects.get_or_create( - document_manager=document_manager, - tool_id=tool, - profile_manager=profile_manager, - prompt_id=prompt, - is_single_pass_extract=is_single_pass_extract, - defaults={ - "output": output, - "eval_metrics": eval_metrics, - "context": context, - "challenge_data": challenge_data, - "highlight_data": highlight_data, - "confidence_data": confidence_data, - "word_confidence_data": word_confidence_data, - }, + # _base_manager, not objects: the manager filter applies to the + # get half of get_or_create but not the create half, so a row + # the org scope hides makes get miss and create collide with + # unique_prompt_output_index. `tool` and `document_manager` are + # already org-verified by the caller. + prompt_output, success = ( + PromptStudioOutputManager._base_manager.get_or_create( + document_manager=document_manager, + tool_id=tool, + profile_manager=profile_manager, + prompt_id=prompt, + is_single_pass_extract=is_single_pass_extract, + defaults={ + "output": output, + "eval_metrics": eval_metrics, + "context": context, + "challenge_data": challenge_data, + "highlight_data": highlight_data, + "confidence_data": confidence_data, + "word_confidence_data": word_confidence_data, + }, + ) ) if success: diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py index 4e0a746f3f..bc41b34594 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py @@ -1,10 +1,11 @@ import logging +import uuid from typing import Any from django.db.models import QuerySet from django.http import HttpRequest from rest_framework import status, viewsets -from rest_framework.exceptions import ValidationError +from rest_framework.exceptions import APIException, ValidationError from rest_framework.response import Response from rest_framework.versioning import URLPathVersioning from utils.common_utils import CommonUtils @@ -28,6 +29,42 @@ logger = logging.getLogger(__name__) +def _validated_tool_id(raw: Any) -> uuid.UUID: + """A query-string ``tool_id`` as a UUID, or a 400. + + ``CustomTool.tool_id`` is a UUID primary key, so a non-UUID value makes + ``filter()`` raise Django's ``ValidationError`` while the query is being + *built*. drf_standardized_errors maps only ``Http404`` and Django's + ``PermissionDenied``, so that surfaced as a 500 rather than a 400. + """ + try: + return uuid.UUID(str(raw)) + except (ValueError, AttributeError, TypeError): + raise ValidationError(detail="'tool_id' must be a valid UUID.") + + +def _required_organization(tool_id: Any) -> Any: + """The request's organization, refusing to proceed without one. + + ``UserContext.get_organization()`` returns None on both + ``Organization.DoesNotExist`` and ``ProgrammingError``, neither logged. A + None here compiles to ``organization_id IS NULL``, which matches nothing + whatever the tool id — downstream every output renders as ``""`` and the + user sees a blank project that has real persisted outputs, with nothing to + correlate in logs. These endpoints are only routed under + ``/api/v1/unstract//``, so a null org is a bug, not a state to serve. + """ + organization = UserContext.get_organization() + if organization is None: + logger.error( + "No organization in context while reading prompt-studio outputs " + "(tool %s); refusing to serve an unscoped empty result.", + tool_id, + ) + raise APIException(detail="Organization context is unavailable.") + return organization + + class PromptStudioOutputView(viewsets.ModelViewSet): versioning_class = URLPathVersioning serializer_class = PromptStudioOutputSerializer @@ -77,9 +114,11 @@ def latest_outputs_by_keys(self, request: HttpRequest) -> Response: if not prompt_keys: return Response({}, status=status.HTTP_200_OK) - # Custom actions skip filter_queryset(), so OrganizationFilterBackend - # never runs — scope explicitly to prevent cross-tenant reads. - organization = UserContext.get_organization() + tool_id = _validated_tool_id(tool_id) + + # A raw .objects query is not routed through filter_queryset(), so + # OrganizationFilterBackend does not see it — scope explicitly. + organization = _required_organization(tool_id) prompt_id_to_key = dict( ToolStudioPrompt.objects.filter( tool_id=tool_id, @@ -121,18 +160,21 @@ def get_output_for_tool_default(self, request: HttpRequest) -> Response: if not tool_id: raise ValidationError(detail=tool_validation_message) + tool_id = _validated_tool_id(tool_id) + organization = _required_organization(tool_id) + # Fetch ToolStudioPrompt records based on tool_id. - # Custom actions skip filter_queryset(), so OrganizationFilterBackend - # never runs — scope explicitly to prevent cross-tenant reads. + # A raw .objects query is not routed through filter_queryset(), so + # OrganizationFilterBackend does not see it — scope explicitly. # - # No exception handling here: filter() does not raise for a missing or - # out-of-org tool, it returns empty. Empty is also the correct result - # for a tool that simply has no prompts yet, which is the normal state - # of a newly created project — so this stays a 200 with an empty body - # rather than a validation error. + # No exception handling below: for a valid UUID that matches no row, or + # a tool in another organization, filter() returns empty rather than + # raising. Empty is also the correct result for a tool that simply has + # no prompts yet, the normal state of a newly created project — so that + # case stays a 200 with an empty body. tool_studio_prompts = ToolStudioPrompt.objects.filter( tool_id=tool_id, - tool_id__organization=UserContext.get_organization(), + tool_id__organization=organization, ).order_by("sequence_number") # Invoke helper method to frame and fetch default response. diff --git a/backend/prompt_studio/prompt_studio_v2/models.py b/backend/prompt_studio/prompt_studio_v2/models.py index 47aa284293..8f9f9bc316 100644 --- a/backend/prompt_studio/prompt_studio_v2/models.py +++ b/backend/prompt_studio/prompt_studio_v2/models.py @@ -16,8 +16,7 @@ class ToolStudioPrompt(BaseModel): It has Many to one relation with CustomTool for ToolStudio. """ - # See DocumentManager.objects — custom @action methods bypass the - # OrganizationFilterBackend, so scoping has to be at the manager. + # See DocumentManager.objects for why scoping lives at the manager. # tool_id is nullable, so prompts orphaned from their tool are excluded. objects = OrgAwareManager() diff --git a/backend/prompt_studio/tests/test_cross_org_isolation.py b/backend/prompt_studio/tests/test_cross_org_isolation.py index eac6a72d02..f602e53ffd 100644 --- a/backend/prompt_studio/tests/test_cross_org_isolation.py +++ b/backend/prompt_studio/tests/test_cross_org_isolation.py @@ -1,10 +1,11 @@ """Organization isolation for the prompt-studio child models. -Custom DRF ``@action`` methods never call ``filter_queryset()``, so -``OrganizationFilterBackend`` does not run on them and a raw -``.objects.get()/filter()`` inside one is not org-scoped. These tests pin the -controls that cover that gap: org scoping on the managers, plus explicit -scoping where an id arrives directly from the request. +``OrganizationFilterBackend`` only scopes querysets routed through +``filter_queryset()``. A raw ``.objects.get()/filter()`` written inside a view +bypasses it — including inside a custom DRF ``@action``, where +``self.get_object()`` *is* filtered but the hand-written queries beside it are +not. These tests pin the controls that cover that gap: org scoping on the +managers, plus explicit scoping where an id arrives directly from the request. Shape of each case: act as org A, pass an org B id, assert the call is refused and org B's row is untouched. @@ -15,12 +16,17 @@ import pytest from account_v2.models import Organization, User from adapter_processor_v2.models import AdapterInstance +from django.contrib.contenttypes.models import ContentType from django.test import TestCase from django.urls import NoReverseMatch, reverse +from permissions.roles import ResourceRole +from rest_framework.test import APIRequestFactory, force_authenticate +from tenant_account_v2.models import ResourceMembership from utils.user_context import UserContext from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.models import CustomTool +from prompt_studio.prompt_studio_core_v2.views import PromptStudioCoreView from prompt_studio.prompt_studio_document_manager_v2.models import DocumentManager from prompt_studio.prompt_studio_index_manager_v2.models import IndexManager from prompt_studio.prompt_studio_output_manager_v2.models import ( @@ -175,25 +181,6 @@ def test_delete_for_ide_lookup_is_tool_scoped(self): pk=self.a.document.document_id, tool=sibling ) - def test_rejected_default_leaves_the_existing_default_intact(self): - """A non-matching id must not clear the tool's current default. - - The de-dup update runs against every profile on the tool, so resolving - the target after it would leave the tool with no default at all when the - id turns out to be someone else's. - """ - assert ProfileManager.objects.get(pk=self.a.profile.profile_id).is_default - - with self.assertRaises(ProfileManager.DoesNotExist): - ProfileManager.objects.get( - pk=self.b.profile.profile_id, prompt_studio_tool=self.a.tool - ) - - self.a.profile.refresh_from_db() - assert self.a.profile.is_default, ( - "the tool lost its default profile while rejecting another org's id" - ) - def test_make_profile_default_lookup_is_tool_scoped(self): """This lookup runs after ``get_object()`` has already passed authz on the caller's own tool, so org scope alone does not constrain it.""" @@ -205,6 +192,90 @@ def test_make_profile_default_lookup_is_tool_scoped(self): UserContext.set_organization_identifier(self.b.org.organization_id) assert ProfileManager.objects.get(pk=self.b.profile.profile_id).is_default + # --- the ordering fix, driven through the view ------------------------ + + def _make_profile_default(self, tool, profile_id): + """PATCH make_profile_default as the owner of ``tool``. + + ``CustomTool.objects.for_user`` resolves visibility through + ResourceMembership, which the create *view* writes — the fixture builds + rows directly, so the OWNER row has to be added here or get_object() + 404s before the code under test runs. + """ + ResourceMembership.objects.get_or_create( + user=self.a.user, + role=ResourceRole.OWNER, + content_type=ContentType.objects.get_for_model(CustomTool), + object_id=str(tool.tool_id), + ) + view = PromptStudioCoreView.as_view({"patch": "make_profile_default"}) + request = APIRequestFactory().patch( + f"/prompt-studio/{tool.tool_id}/make_profile_default", + {"default_profile": str(profile_id)}, + format="json", + ) + force_authenticate(request, user=self.a.user) + return view(request, pk=str(tool.tool_id)) + + def _second_profile_on_tool_a(self): + return ProfileManager.objects.create( + profile_name="profile-a-second", + vector_store=self.a.profile.vector_store, + embedding_model=self.a.profile.embedding_model, + llm=self.a.profile.llm, + x2text=self.a.profile.x2text, + chunk_size=0, + chunk_overlap=0, + section="Default", + retrieval_strategy="simple", + similarity_top_k=3, + prompt_studio_tool=self.a.tool, + is_default=False, + created_by=self.a.user, + ) + + def test_make_profile_default_switches_the_default(self): + """The allow path: the old default is cleared and the new one set.""" + second = self._second_profile_on_tool_a() + + response = self._make_profile_default(self.a.tool, second.profile_id) + + assert response.status_code == 200, response.data + self.a.profile.refresh_from_db() + second.refresh_from_db() + assert second.is_default + assert not self.a.profile.is_default + + def test_rejected_default_leaves_the_existing_default_intact(self): + """A non-matching id must not clear the tool's current default. + + Driven through the view on purpose: the de-dup update runs against + every profile on the tool, and an ORM-only test never executes it, so + it cannot observe this property at all. + + What actually guards the invariant is resolving and clearing under one + of two conditions — resolve first, or clear first but inside the + transaction, where the 404 rolls the clear back. Mutation-tested: this + fails (0 defaults left) only on clear-first *without* the transaction, + which is what the code did before. Reverting just the ordering, with + ``transaction.atomic()`` still in place, is genuinely safe and does not + fail here. + """ + self._second_profile_on_tool_a() + assert ProfileManager.objects.get(pk=self.a.profile.profile_id).is_default + + response = self._make_profile_default(self.a.tool, self.b.profile.profile_id) + + assert response.status_code == 404, response.data + assert ( + ProfileManager.objects.filter( + prompt_studio_tool=self.a.tool, is_default=True + ).count() + == 1 + ), "the tool lost (or duplicated) its default while rejecting another org's id" + self.a.profile.refresh_from_db() + assert self.a.profile.is_default + # --- A-4: the dead, state-changing-over-GET route is gone ------------- def test_file_delete_route_removed(self): diff --git a/backend/tool_instance_v2/internal_views.py b/backend/tool_instance_v2/internal_views.py index 25165a9739..8bfd773986 100644 --- a/backend/tool_instance_v2/internal_views.py +++ b/backend/tool_instance_v2/internal_views.py @@ -23,7 +23,9 @@ class ToolExecutionInternalViewSet(viewsets.ModelViewSet): """Internal API for tool execution operations used by lightweight workers.""" serializer_class = ToolInstanceSerializer - # Backward compat: remove once all workers pass X-Organization-ID. + # OrganizationFilterBackend is off here; scoping runs through + # filter_queryset_by_organization, which fails closed, so a caller without + # X-Organization-ID gets zero rows. skip_org_filter = True def get_queryset(self): diff --git a/backend/utils/filters/organization_filter.py b/backend/utils/filters/organization_filter.py index a0eba72daf..072e9d03bd 100644 --- a/backend/utils/filters/organization_filter.py +++ b/backend/utils/filters/organization_filter.py @@ -40,6 +40,13 @@ class NotificationViewSet(viewsets.ModelViewSet): "pipeline__workflow__organization", "api__workflow__organization", ] + + Precedence: org_filter_paths wins over the model's pin in + ORG_PATH_OVERRIDES, and is checked before get_org_path is ever called. So + a viewset that sets both scopes through the paths here, not the pin, while + OrgAwareManager on the same model still uses the pin. Prefer the pin — + it applies at both layers. Reach for org_filter_paths only when the model + genuinely needs OR across several nullable paths. """ def filter_queryset(self, request, queryset, view): diff --git a/backend/utils/models/org_aware_manager.py b/backend/utils/models/org_aware_manager.py index ab40303dbb..33cdcd94d4 100644 --- a/backend/utils/models/org_aware_manager.py +++ b/backend/utils/models/org_aware_manager.py @@ -58,20 +58,35 @@ def get_queryset(self): try: org = UserContext.get_organization() - except (RuntimeError, OperationalError, ProgrammingError): + except (RuntimeError, OperationalError, ProgrammingError) as exc: # OperationalError: DB not reachable (startup, migrations) # ProgrammingError: schema not ready (during migrations) # RuntimeError: pytest-django blocks DB access outside - # @pytest.mark.django_db. Note: this is a broad catch — any - # RuntimeError (e.g. from StateStore/middleware) returns an - # unfiltered queryset (fail-open). This is acceptable because - # OrgAwareManager is defense-in-depth; OrganizationFilterBackend - # at the view layer is the primary security boundary and - # fails-closed independently. + # @pytest.mark.django_db. + # + # Deliberately fail open: these are all "the request context does + # not exist yet", not "this caller may not see these rows". + # OrganizationFilterBackend is the primary boundary and fails + # closed independently at the view layer. + # + # The RuntimeError arm is broader than its stated cause — + # StateStore.get raises it for any unrecognised CONCURRENCY_MODE + # — so log it. This path is rare (startup, migrations, tests), so + # a line here is signal rather than noise, and it is the only way + # an unexpected fail-open becomes visible. + logger.warning( + "OrgAwareManager: no organization context for %s (%s: %s); " + "returning an unfiltered queryset.", + self.model._meta.label, + type(exc).__name__, + exc, + ) return qs if org is None: - # No request context (Celery, management commands, shell) + # No request context: Celery, management commands, shell. Not + # logged — this is the normal state for every query those make, + # and a line per queryset would drown the case above. return qs path = get_org_path(self.model) diff --git a/backend/utils/models/org_path_discovery.py b/backend/utils/models/org_path_discovery.py index e057fedc29..79969b323f 100644 --- a/backend/utils/models/org_path_discovery.py +++ b/backend/utils/models/org_path_discovery.py @@ -26,9 +26,22 @@ # Reordering two fields can therefore swap in a different path of the same # length, and if that path runs through a nullable FK the resulting INNER JOIN # silently drops every row with a NULL — data loss that reads as "missing -# records", not as an error. Pinning freezes the path for both consumers -# (OrgAwareManager and OrganizationFilterBackend); test_org_path_discovery -# asserts each pin still matches BFS and traverses only non-nullable FKs. +# records", not as an error. Pinning freezes the path against that. +# +# What test_org_path_discovery actually asserts: each pin still matches what +# BFS would pick, and every hop on it is non-nullable *unless* listed in that +# module's KNOWN_NULLABLE_HOPS with a reason. Several pins are on that list — +# including every terminal `organization` FK, which DefaultOrganizationMixin +# declares null=True — so "pinned" does not mean "cannot drop rows", it means +# "the rows it drops are known and written down". +# +# Precedence, for the two consumers: +# - OrgAwareManager always uses the pin. +# - OrganizationFilterBackend checks a viewset's `org_filter_paths` FIRST and +# only falls back to the pin. A viewset that sets it therefore scopes that +# model through a different join than its pin. Prefer the pin; reach for +# `org_filter_paths` only when a model needs OR across several nullable +# paths, which is why notification_v2 has it. ORG_PATH_OVERRIDES: dict[str, str] = { "prompt_studio_document_manager_v2.DocumentManager": "tool__organization", "prompt_studio_index_manager_v2.IndexManager": ( diff --git a/backend/utils/organization_utils.py b/backend/utils/organization_utils.py index 9af586954d..1d7add8c0c 100644 --- a/backend/utils/organization_utils.py +++ b/backend/utils/organization_utils.py @@ -74,11 +74,18 @@ def get_organization_context(organization: Organization) -> dict[str, Any]: def filter_queryset_by_organization(queryset, request, organization_field="organization"): """Filter a Django queryset by the request's organization context. - Fails closed. Six internal viewsets set ``skip_org_filter = True``, which - disables OrganizationFilterBackend, leaving this function as their only - tenant boundary — so returning the queryset unfiltered when there is no - organization context hands back every organization's rows. A scoping helper - returns nothing when it cannot scope, never everything. + Fails closed. For every caller, this function is the only tenant boundary + in the request: the viewsets that reach it set ``skip_org_filter = True``, + which disables OrganizationFilterBackend, and the function-based + ``@api_view`` handlers that reach it have no filter backend at all. + Returning the queryset unfiltered when there is no organization context + would hand back every organization's rows. + + Scope note: this policy is local to this helper. ``OrgAwareManager`` + deliberately fails *open* when there is no request context, so Celery + tasks, management commands and the shell keep full access to the models it + scopes — see the comment on its ``get_queryset``. The two are not in + conflict; they guard different callers. The absent-header case is not exotic: ``InternalAPIAuthMiddleware`` logs a warning and continues when ``X-Organization-ID`` is missing, so any caller diff --git a/backend/utils/tests/test_org_path_discovery.py b/backend/utils/tests/test_org_path_discovery.py index 958c393139..384de242ba 100644 --- a/backend/utils/tests/test_org_path_discovery.py +++ b/backend/utils/tests/test_org_path_discovery.py @@ -21,7 +21,23 @@ # Nullable hops accepted as pre-existing behaviour, not introduced here. # Rows with a NULL value on these FKs are excluded from every org-scoped # query. Anything not listed must be non-nullable. -KNOWN_NULLABLE_HOPS = {("prompt_studio_v2.ToolStudioPrompt", "tool_id")} +# +# The terminal `organization` FK is on this list for every pin: +# DefaultOrganizationMixin declares it null=True, and save() backfills it from +# UserContext, which is None outside a request. So a CustomTool or +# AdapterInstance created by a management command, data migration, Celery task +# or shell persists with organization_id NULL. Those rows are already invisible +# to their own model's default manager, so the pins do not make them any less +# visible — but the hop is nullable and the assertion below must say so rather +# than skip it. +KNOWN_NULLABLE_HOPS = { + ("prompt_studio_v2.ToolStudioPrompt", "tool_id"), + ("prompt_studio_document_manager_v2.DocumentManager", "organization"), + ("prompt_studio_index_manager_v2.IndexManager", "organization"), + ("prompt_studio_output_manager_v2.PromptStudioOutputManager", "organization"), + ("prompt_studio_v2.ToolStudioPrompt", "organization"), + ("prompt_profile_manager_v2.ProfileManager", "organization"), +} @pytest.mark.parametrize("label,expected", PINS) @@ -43,15 +59,19 @@ def test_pin_matches_discovery(label, expected): @pytest.mark.parametrize("label,expected", PINS) def test_pin_traverses_only_non_nullable_fks(label, expected): - """Every hop before `organization` must be non-nullable. + """Every hop on the pin must be non-nullable, or listed as a known exception. Django turns a positive filter over a nullable FK into an INNER JOIN, which drops rows whose FK is NULL. On an org filter that is invisible data loss. + + The terminal `organization` hop is checked too, not skipped: it is the one + that is nullable on every pin, so excluding it would make this assertion + pass while proving nothing about the join that matters most. """ model = apps.get_model(label) hops = expected.split("__") - for hop in hops[:-1]: + for hop in hops: field = model._meta.get_field(hop) assert not field.null or (label, hop) in KNOWN_NULLABLE_HOPS, ( f"{model._meta.label}.{hop} is nullable: this pin drops every row " @@ -60,6 +80,5 @@ def test_pin_traverses_only_non_nullable_fks(label, expected): ) model = field.related_model - # Final hop must actually be the Organization FK. - org_field = model._meta.get_field(hops[-1]) - assert org_field.related_model._meta.label == "account_v2.Organization" + # The walk must have landed on Organization, not merely survived. + assert model._meta.label == "account_v2.Organization" diff --git a/backend/utils/tests/test_organization_scoping.py b/backend/utils/tests/test_organization_scoping.py index ce6e4b89f0..e277e9f4d0 100644 --- a/backend/utils/tests/test_organization_scoping.py +++ b/backend/utils/tests/test_organization_scoping.py @@ -1,11 +1,12 @@ """``filter_queryset_by_organization`` must fail closed. -Six internal viewsets set ``skip_org_filter = True``, which disables -OrganizationFilterBackend and leaves this helper as their only tenant -boundary. Returning the queryset unfiltered when there is no organization -context therefore returns every organization's rows, and the absent-header -case is reachable — the internal auth middleware warns and continues rather -than rejecting. +Every caller reaches it with OrganizationFilterBackend either disabled +(``skip_org_filter = True``) or absent — the function-based internal handlers +have no filter backend at all — so this helper is their only tenant boundary. +Returning the queryset unfiltered when there is no organization context +therefore returns every organization's rows, and the absent-header case is +reachable: the internal auth middleware warns and continues rather than +rejecting. """ import secrets diff --git a/backend/workflow_manager/file_execution/internal_views.py b/backend/workflow_manager/file_execution/internal_views.py index 8fed360740..256331a0d7 100644 --- a/backend/workflow_manager/file_execution/internal_views.py +++ b/backend/workflow_manager/file_execution/internal_views.py @@ -29,10 +29,10 @@ class FileExecutionInternalViewSet(viewsets.ModelViewSet): serializer_class = WorkflowFileExecutionSerializer lookup_field = "id" - # Backward compat: workers may call without X-Organization-ID during - # rolling deployments. Safe because internal APIs require service API key - # and get_queryset() applies org filtering when header is present. - # Remove once all workers reliably pass X-Organization-ID. + # OrganizationFilterBackend is off here; get_queryset() scopes instead, via + # filter_queryset_by_organization. That helper fails closed, so a worker + # calling without X-Organization-ID now gets zero rows rather than every + # organization's — the header is required in practice, not optional. skip_org_filter = True def get_object(self): diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py index 4f9de9f0aa..689a8e2b9f 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -49,10 +49,10 @@ class WorkflowExecutionInternalViewSet(viewsets.ReadOnlyModelViewSet): serializer_class = WorkflowExecutionSerializer lookup_field = "id" - # Backward compat: workers may call without X-Organization-ID during - # rolling deployments. Safe because internal APIs require service API key - # and get_queryset() applies org filtering when header is present. - # Remove once all workers reliably pass X-Organization-ID. + # OrganizationFilterBackend is off here; get_queryset() scopes instead, via + # filter_queryset_by_organization. That helper fails closed, so a worker + # calling without X-Organization-ID now gets zero rows rather than every + # organization's — the header is required in practice, not optional. skip_org_filter = True def get_queryset(self): diff --git a/backend/workflow_manager/workflow_v2/views.py b/backend/workflow_manager/workflow_v2/views.py index fefba8c21a..aa31f0f671 100644 --- a/backend/workflow_manager/workflow_v2/views.py +++ b/backend/workflow_manager/workflow_v2/views.py @@ -403,10 +403,10 @@ class WorkflowExecutionInternalViewSet(viewsets.ReadOnlyModelViewSet): serializer_class = WorkflowExecutionSerializer lookup_field = "id" - # Backward compat: workers may call without X-Organization-ID during - # rolling deployments. Safe because internal APIs require service API key - # and get_queryset() applies org filtering when header is present. - # Remove once all workers reliably pass X-Organization-ID. + # OrganizationFilterBackend is off here; get_queryset() scopes instead, via + # filter_queryset_by_organization. That helper fails closed, so a worker + # calling without X-Organization-ID now gets zero rows rather than every + # organization's — the header is required in practice, not optional. skip_org_filter = True def get_queryset(self): diff --git a/workers/ide_callback/tasks.py b/workers/ide_callback/tasks.py index e35b76f933..1a298dffe3 100644 --- a/workers/ide_callback/tasks.py +++ b/workers/ide_callback/tasks.py @@ -243,9 +243,14 @@ def ide_index_complete( organization_id=org_id, ) except Exception: - logger.warning( + # Non-fatal — primary indexing already succeeded — but not + # harmless: without the status, check_extraction_status stays + # False and every later Answer Prompt re-runs the full X2Text + # extraction. ERROR because nothing downstream reports it. + logger.error( "Failed to mark extraction_status for document %s " - "profile %s; primary indexing succeeded.", + "profile %s; extraction will be repeated on every " + "subsequent prompt run.", document_id, profile_manager_id, exc_info=True,