From 447d0c940b36d59b669fb137ebd007e263e41b44 Mon Sep 17 00:00:00 2001 From: Athul Date: Mon, 27 Jul 2026 10:02:04 +0530 Subject: [PATCH 01/10] 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 02/10] 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 03/10] 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 04/10] 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 05/10] 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 06/10] 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, From be459e960615251103fbb92b4e956ff2b80b7daa Mon Sep 17 00:00:00 2001 From: Athul Date: Fri, 28 Aug 2026 20:03:32 +0530 Subject: [PATCH 07/10] Address review on prompt-studio org scoping MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Manager and helper behaviour: - OrgAwareManager fails closed when an organization identifier is set but does not resolve. get_organization() flattens three states to None; two of them happen inside a request, where returning everything crosses tenants. No identifier at all stays fail-open for Celery, commands and the shell. - ProfileManager.for_user drops its own org join. get_queryset supplies org scoping now, and the second join was narrower — it dropped tool-less profiles and applied the opposite null policy. - PromptStudioOutputManager and IndexManager go back to the scoped manager on both halves of get_or_create; the unscoped create paired with a scoped update could return stale output as a success. - mark_extraction_status returns why it failed. The internal endpoint maps a missing document to 404 (outside the client's retry set) and a write failure to 500. - migration_utils drops the pre-transaction profile lookup that short-circuited the diagnostic below it. - delete_for_ide warns only when an unscoped probe finds index rows the scoped query missed, not on every never-indexed delete. - Output views validate document_manager and the list action's ids, not just tool_id; document_manager is required rather than silently matching NULL. Tests: - The two tool-scoping cases run through the view against a sibling tool in the same org, where org scope cannot mask a missing predicate. Both were vacuous before and pass with the predicate removed. - New cases for the fail-closed manager branch, the four request guards and the three extraction-status outcomes. Each verified against the unguarded code first. - Pin tests key nullable hops on the declaring model, assert the pin set against a literal, and prove the pin short-circuits discovery. - setUp registers the UserContext cleanup before anything can raise. Comments rewritten to describe the code rather than the change. --- backend/file_management/views.py | 5 +- .../prompt_profile_manager_v2/models.py | 26 +-- .../prompt_studio_core_v2/internal_views.py | 23 ++- .../prompt_studio_core_v2/migration_utils.py | 25 +-- .../prompt_studio_helper.py | 9 +- .../prompt_studio_core_v2/views.py | 46 +++--- .../prompt_studio_index_helper.py | 48 +++--- .../output_manager_helper.py | 40 +++-- .../prompt_studio_output_manager_v2/views.py | 45 ++++-- .../tests/test_cross_org_isolation.py | 152 ++++++++++++++---- .../tests/test_request_validation.py | 149 +++++++++++++++++ backend/utils/models/org_aware_manager.py | 31 +++- backend/utils/models/org_path_discovery.py | 3 +- backend/utils/organization_utils.py | 24 +-- .../utils/tests/test_org_path_discovery.py | 74 +++++++-- .../utils/tests/test_organization_scoping.py | 16 +- .../file_execution/internal_views.py | 5 +- backend/workflow_manager/internal_views.py | 5 +- backend/workflow_manager/workflow_v2/views.py | 5 +- 19 files changed, 536 insertions(+), 195 deletions(-) create mode 100644 backend/prompt_studio/tests/test_request_validation.py diff --git a/backend/file_management/views.py b/backend/file_management/views.py index 2a4d684763..ea39cb1d28 100644 --- a/backend/file_management/views.py +++ b/backend/file_management/views.py @@ -26,10 +26,7 @@ class FileManagementViewSet(viewsets.ModelViewSet): - """FileManagement view. - - Handles GET, POST, PUT and PATCH - """ + """FileManagement view.""" versioning_class = URLPathVersioning diff --git a/backend/prompt_studio/prompt_profile_manager_v2/models.py b/backend/prompt_studio/prompt_profile_manager_v2/models.py index fbd7437656..610618bd62 100644 --- a/backend/prompt_studio/prompt_profile_manager_v2/models.py +++ b/backend/prompt_studio/prompt_profile_manager_v2/models.py @@ -7,7 +7,6 @@ from tenant_account_v2.organization_member_service import OrganizationMemberService 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 @@ -17,31 +16,32 @@ class ProfileManagerModelManager(OrgAwareManager): def for_user(self, user): """Read visibility: profile's own share fields OR parent CustomTool sharing. - Org-scoped via the parent (no ``organization`` FK on this model). + Sharing only. Organization scoping comes from ``get_queryset``, which + every branch here builds on, and which scopes through + ``vector_store__organization``. This method used to AND on + ``prompt_studio_tool__organization`` as well; that is a second, + narrower org join that drops tool-less profiles and contradicts the + null policy in ``get_queryset``. + The parent-tool branch lets shared-project users see existing profiles so ``IsOwner`` on the viewset returns 403 on mutation instead of DRF raising 404 first. Mutation gating lives on the viewset; this method governs read visibility only. """ - org_scope = Q(prompt_studio_tool__organization=UserContext.get_organization()) - if getattr(user, "is_service_account", False): - return self.filter(org_scope) + return self.all() if OrganizationMemberService.is_user_organization_admin(user): - return self.filter(org_scope) + return self.all() # Union the legacy own-share branches with the parent-tool branch # so any row visible before the UN-2977 fix stays visible. accessible_tools = CustomTool.objects.for_user(user) return self.filter( - org_scope - & ( - Q(created_by=user) - | Q(shared_users=user) - | Q(shared_to_org=True) - | Q(prompt_studio_tool__in=accessible_tools) - ) + Q(created_by=user) + | Q(shared_users=user) + | Q(shared_to_org=True) + | Q(prompt_studio_tool__in=accessible_tools) ).distinct() 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 b4246e0cf7..7d916fb5ef 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/internal_views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/internal_views.py @@ -199,11 +199,12 @@ def extraction_status(request): try: from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_index_manager_v2.prompt_studio_index_helper import ( + ExtractionStatusResult, PromptStudioIndexHelper, ) profile_manager = ProfileManager.objects.get(pk=profile_manager_id) - success = PromptStudioIndexHelper.mark_extraction_status( + result = PromptStudioIndexHelper.mark_extraction_status( document_id=document_id, profile_manager=profile_manager, x2text_config_hash=x2text_config_hash, @@ -211,23 +212,35 @@ def extraction_status(request): extracted=extracted, error_message=error_message, ) - if not success: + if result is not ExtractionStatusResult.OK: # 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. + # + # The two failures need different statuses. The client retries + # {500, 502, 503, 504} three times with a 1s backoff factor, so a + # document that is gone would burn four round trips and ~7s of + # worker sleep on a condition no retry can change. 404 is outside + # that set. + missing = result is ExtractionStatusResult.DOCUMENT_MISSING logger.error( - "extraction_status not recorded for document %s profile %s", + "extraction_status not recorded for document %s profile %s (%s)", document_id, profile_manager_id, + result.value, ) return JsonResponse( { "success": False, - "error": "Extraction status could not be recorded", + "error": "Document not found" + if missing + else "Extraction status could not be recorded", }, - status=status.HTTP_500_INTERNAL_SERVER_ERROR, + status=status.HTTP_404_NOT_FOUND + if missing + else status.HTTP_500_INTERNAL_SERVER_ERROR, ) return JsonResponse({"success": True}) 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 ea838a06bf..500f4ff4f5 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py +++ b/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py @@ -38,17 +38,9 @@ def migrate_tool_to_adapter_based( ) return False - # Check if there's a summarize profile before entering transaction - try: - summarize_profile = ProfileManager.objects.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" - ) - return False - + # No pre-transaction lookup: the in-transaction fetch below repeats it + # exactly, and catching the miss up here short-circuited the diagnostic + # that tells the two miss reasons apart. try: with transaction.atomic(): # Re-fetch the instance within transaction to ensure fresh data @@ -66,13 +58,12 @@ def migrate_tool_to_adapter_based( of=("self",) ).get(prompt_studio_tool=tool_instance, is_summarize_llm=True) except ObjectDoesNotExist: - # 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. + # vector_store__organization, so a miss means either the + # profile does not exist or the org filter hides it. The + # second never self-heals — this lazy migration re-runs and + # re-skips on every invocation — so the two get different + # log levels. exists_unscoped = ProfileManager._base_manager.filter( prompt_studio_tool=tool_instance, is_summarize_llm=True ).exists() diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index 6d8763e05a..ba1c1fa989 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -70,6 +70,7 @@ ) from prompt_studio.prompt_studio_document_manager_v2.models import DocumentManager from prompt_studio.prompt_studio_index_manager_v2.prompt_studio_index_helper import ( # noqa: E501 + ExtractionStatusResult, PromptStudioIndexHelper, ) from prompt_studio.prompt_studio_output_manager_v2.output_manager_helper import ( @@ -2535,7 +2536,7 @@ def dynamic_extractor( result = dispatcher.dispatch(extract_context) if not result.success: msg = result.error or "Unknown extraction error" - success = PromptStudioIndexHelper.mark_extraction_status( + result = PromptStudioIndexHelper.mark_extraction_status( document_id=document_id, profile_manager=profile_manager, x2text_config_hash=x2text_config_hash, @@ -2543,7 +2544,7 @@ def dynamic_extractor( extracted=False, error_message=msg, ) - if not success: + if result is not ExtractionStatusResult.OK: logger.warning( f"Failed to mark extraction failure for document {document_id}. " f"Extraction failed but status not saved." @@ -2553,13 +2554,13 @@ def dynamic_extractor( ) extracted_text = result.data.get("extracted_text", "") - success = PromptStudioIndexHelper.mark_extraction_status( + result = PromptStudioIndexHelper.mark_extraction_status( document_id=document_id, profile_manager=profile_manager, x2text_config_hash=x2text_config_hash, enable_highlight=enable_highlight, ) - if not success: + if result is not ExtractionStatusResult.OK: logger.warning( f"Failed to mark extraction success for document {document_id}. " f"Extraction completed but status not saved." diff --git a/backend/prompt_studio/prompt_studio_core_v2/views.py b/backend/prompt_studio/prompt_studio_core_v2/views.py index 1bc941a8dd..6a258ec8e7 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/views.py @@ -391,10 +391,10 @@ 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. + # Validate before looking anything up. drf_standardized_errors maps + # neither the KeyError of a missing key nor the Django ValidationError + # a non-UUID raises while the query is built, so without this both are + # 500s next to the 404 a valid-but-unmatched id returns. default_profile = request.data.get("default_profile") if not default_profile: raise ValidationError(detail="'default_profile' is required.") @@ -403,10 +403,12 @@ def make_profile_default(self, request: HttpRequest, pk: Any = None) -> Response 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. + # Resolve the target before clearing anything. The transaction below + # is what actually guarantees the tool never ends up with zero + # defaults — a 404 raised inside it rolls the clear back — so this + # ordering is the second of two independent guards, not the only one. + # 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=default_profile, @@ -1171,17 +1173,23 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response: # 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, - ) + # Empty is almost always "never indexed", which is an ordinary + # delete and not worth a line. The case worth warning about is + # "the org filter hid the rows": there the Redis indexing flags + # outlive the document, and a re-upload of the same file is + # treated as already indexed. Only the unscoped probe can tell + # them apart, and it only runs on this already-empty path. + if IndexManager._base_manager.filter( + document_manager=document_id + ).exists(): + logger.warning( + "Index managers for document %s (tool %s) are not " + "visible in 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( 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 1c10bf5b1a..0a4530fea3 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,5 +1,6 @@ import json import logging +from enum import Enum from django.core.exceptions import ImproperlyConfigured from django.db import DatabaseError, transaction @@ -13,6 +14,19 @@ logger = logging.getLogger(__name__) +class ExtractionStatusResult(Enum): + """Why ``mark_extraction_status`` did or did not write. + + A plain bool collapsed "the document is gone" into "the write failed", + and the internal API turned both into a 500 — which the worker's client + retries. A missing document never becomes present on retry. + """ + + OK = "ok" + DOCUMENT_MISSING = "document_missing" + WRITE_FAILED = "write_failed" + + class PromptStudioIndexHelper: @staticmethod def handle_index_manager( @@ -76,7 +90,7 @@ def mark_extraction_status( enable_highlight: bool, extracted: bool = True, error_message: str | None = None, - ) -> bool: + ) -> ExtractionStatusResult: """Marks the extraction status for a given document. Uses x2text_config_hash (hash of X2Text config metadata) as the key. @@ -91,7 +105,10 @@ def mark_extraction_status( error_message (str | None): Error message if extraction failed. Returns: - bool: True if the status is successfully updated, False otherwise. + ExtractionStatusResult: OK on success; DOCUMENT_MISSING when the + document is gone or hidden by the org scope; WRITE_FAILED for + a genuine write error. Callers that only need "did it write" + compare against OK. """ try: @@ -111,15 +128,7 @@ 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",) 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( + index_manager, created = IndexManager.objects.select_for_update( of=("self",) ).get_or_create( document_manager=document, @@ -157,29 +166,28 @@ def mark_extraction_status( f"Error: {error_message}" ) - return True + return ExtractionStatusResult.OK except DocumentManager.DoesNotExist: - # 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. + # Two ways to get here: the row is gone, or the org-scoped manager + # hides it from this caller. Both mean the status was not written, + # and neither is fixed by trying again. logger.error( "Document %s not found or not visible in the current " "organization; extraction status not recorded.", document_id, ) - return False + return ExtractionStatusResult.DOCUMENT_MISSING 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". + # org path pin. Anything else propagates rather than being + # reported as a write failure it is not. logger.exception( "Failed to mark extraction status for document %s", document_id ) - return False + return ExtractionStatusResult.WRITE_FAILED @staticmethod def check_extraction_status( 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 54450201c7..29e73923bf 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,28 +76,24 @@ def update_or_create_prompt_output( the instance. """ try: - # _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, - }, - ) + # Scoped manager on both halves: the update below is scoped + # too, so a mismatch here would leave that update matching zero + # rows and silently returning stale output as a success. + 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, + }, ) 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 bc41b34594..e8cd1f9629 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py @@ -2,6 +2,7 @@ import uuid from typing import Any +from account_v2.models import Organization from django.db.models import QuerySet from django.http import HttpRequest from rest_framework import status, viewsets @@ -29,10 +30,10 @@ logger = logging.getLogger(__name__) -def _validated_tool_id(raw: Any) -> uuid.UUID: - """A query-string ``tool_id`` as a UUID, or a 400. +def _validated_uuid(raw: Any, field_name: str) -> uuid.UUID: + """A query-string id as a UUID, or a 400 naming the field. - ``CustomTool.tool_id`` is a UUID primary key, so a non-UUID value makes + These columns are UUID primary/foreign keys, 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. @@ -40,16 +41,17 @@ def _validated_tool_id(raw: Any) -> uuid.UUID: try: return uuid.UUID(str(raw)) except (ValueError, AttributeError, TypeError): - raise ValidationError(detail="'tool_id' must be a valid UUID.") + raise ValidationError(detail=f"'{field_name}' must be a valid UUID.") -def _required_organization(tool_id: Any) -> Any: +def _required_organization(tool_id: uuid.UUID) -> Organization: """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 + None here compiles to ``organization_id IS NULL``, which matches only rows + whose organization was never set and never the caller's tool — 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. @@ -84,6 +86,18 @@ def get_queryset(self) -> QuerySet | None: PromptStudioOutputManagerKeys.IS_SINGLE_PASS_EXTRACT, "false" ) + # Same 500-on-bad-UUID as the detail actions: build_filter_args copies + # query params straight through, and all four of these are UUID + # columns, so `?tool_id=abc` raises while the query is being built. + for key in ( + PromptStudioOutputManagerKeys.TOOL_ID, + PromptStudioOutputManagerKeys.PROMPT_ID, + PromptStudioOutputManagerKeys.PROFILE_MANAGER, + PromptStudioOutputManagerKeys.DOCUMENT_MANAGER, + ): + if key in filter_args: + filter_args[key] = _validated_uuid(filter_args[key], key) + # Convert the string representation to a boolean value is_single_pass_extract = CommonUtils.str_to_bool(is_single_pass_extract_param) @@ -114,7 +128,7 @@ def latest_outputs_by_keys(self, request: HttpRequest) -> Response: if not prompt_keys: return Response({}, status=status.HTTP_200_OK) - tool_id = _validated_tool_id(tool_id) + tool_id = _validated_uuid(tool_id, PromptStudioOutputManagerKeys.TOOL_ID) # A raw .objects query is not routed through filter_queryset(), so # OrganizationFilterBackend does not see it — scope explicitly. @@ -160,7 +174,20 @@ 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) + tool_id = _validated_uuid(tool_id, PromptStudioOutputManagerKeys.TOOL_ID) + # Same column type, same failure: this one reaches + # PromptStudioOutputManager.objects.filter(document_manager_id=...) in + # the helper below. Required, not optional — absent it stays None and + # compiles to `document_manager_id IS NULL`, which matches nothing and + # renders every prompt as "", so the project looks empty while holding + # real outputs. The only caller always sends it. + if not document_manager_id: + raise ValidationError( + detail="'document_manager' is required and must be a valid UUID." + ) + document_manager_id = _validated_uuid( + document_manager_id, PromptStudioOutputManagerKeys.DOCUMENT_MANAGER + ) organization = _required_organization(tool_id) # Fetch ToolStudioPrompt records based on tool_id. diff --git a/backend/prompt_studio/tests/test_cross_org_isolation.py b/backend/prompt_studio/tests/test_cross_org_isolation.py index f602e53ffd..44082611fd 100644 --- a/backend/prompt_studio/tests/test_cross_org_isolation.py +++ b/backend/prompt_studio/tests/test_cross_org_isolation.py @@ -104,18 +104,20 @@ class CrossOrgIsolationTest(TestCase): """Org A must not reach org B's prompt-studio rows through any manager.""" def setUp(self) -> None: + # Registered before anything can raise, and not in tearDown: UserContext + # is thread-local, so TestCase's transaction rollback does not clear it, + # and unittest skips tearDown when setUp fails. OrgFixture sets the + # identifier as its second statement and then makes eight create() + # calls, any of which can raise — without this the worker process would + # keep an identifier pointing at a rolled-back organization and quietly + # change manager behaviour for every later test. + self.addCleanup(UserContext.set_organization_identifier, None) self.a = OrgFixture(f"org-a-{secrets.token_hex(3)}") self.b = OrgFixture(f"org-b-{secrets.token_hex(3)}") # End state: acting as org A, as a request would. UserContext.set_organization_identifier(self.a.org.organization_id) - def tearDown(self) -> None: - # UserContext 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) -------------------- + # --- manager scoping: the default-deny layer --------------------------- def test_document_of_other_org_is_not_gettable(self): """A document id from another org must not resolve.""" @@ -151,46 +153,136 @@ def test_own_org_rows_remain_visible(self): 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).""" + """Management commands and shell keep full access (fail-open). + + Pinned deliberately: no identifier means no request to take one from, + which is a different state from the one below. + """ UserContext.set_organization_identifier(None) assert DocumentManager.objects.filter( pk=self.b.document.document_id ).exists() + def test_unresolvable_org_context_is_empty(self): + """An identifier that resolves to no row must not fall back to open. + + ``UserContext.get_organization()`` flattens three states to ``None``: + no identifier, ``Organization.DoesNotExist`` and ``ProgrammingError``. + The last two happen inside a request, so treating them like the first + one would serve every organization's rows to a caller whose own + organization could not be looked up. + """ + UserContext.set_organization_identifier("org-that-does-not-exist") + assert not DocumentManager.objects.filter( + pk=self.b.document.document_id + ).exists() + assert not DocumentManager.objects.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.""" + """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) ----------- + # --- explicit scoping in delete_for_ide and make_profile_default ------ - 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", + def _sibling_tool_in_org_a(self): + """A second tool in org A, with its own default profile and document. + + Org scope cannot distinguish these from ``self.a.tool``'s own rows, so + they are what makes the per-tool predicates observable. A cross-org id + would 404 on org scope alone and pass even with the predicate removed. + """ + tool = CustomTool.objects.create( + tool_name=f"sibling-{secrets.token_hex(3)}", 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 - ) + profile = ProfileManager.objects.create( + profile_name=f"sibling-profile-{secrets.token_hex(3)}", + 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=tool, + is_default=True, + created_by=self.a.user, + ) + document = DocumentManager.objects.create( + document_name=f"sibling-{secrets.token_hex(3)}.pdf", + tool=tool, + created_by=self.a.user, + ) + return tool, profile, document - 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 + def _delete_for_ide(self, tool, document_id): + """DELETE prompt-studio/file/ as the owner of ``tool``.""" + 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({"delete": "delete_for_ide"}) + request = APIRequestFactory().delete( + f"/prompt-studio/file/{tool.tool_id}", + {"document_id": str(document_id)}, + format="json", + ) + # The view reads the org off the session; the factory builds a bare + # request, so nothing else populates it. + request.session = {"organization": self.a.org.organization_id} + force_authenticate(request, user=self.a.user) + return view(request, pk=str(tool.tool_id)) + + def test_make_profile_default_refuses_a_sibling_tool_profile(self): + """Same org, different tool: org scope cannot catch this one. + + Driven through the view because the control is the ``prompt_studio_tool`` + predicate on the lookup, not anything the ORM does on its own. Removing + that predicate lets the owner of one tool flip another tool's default + profile within the same organization. + """ + sibling_tool, sibling_profile, _ = self._sibling_tool_in_org_a() + + response = self._make_profile_default(self.a.tool, sibling_profile.profile_id) + + assert response.status_code == 404, response.data + sibling_profile.refresh_from_db() + assert sibling_profile.is_default, "sibling tool's default was altered" + assert ( + ProfileManager.objects.filter( + prompt_studio_tool=sibling_tool, is_default=True + ).count() + == 1 + ) + + def test_delete_for_ide_refuses_a_sibling_tool_document(self): + """Same org, different tool: the ``tool=`` predicate is the only guard. + + Without it the lookup is a bare pk fetch, which both deletes another + tool's document and raises an unhandled ``DoesNotExist`` (500) when the + id is unknown. + """ + _, _, sibling_document = self._sibling_tool_in_org_a() + + response = self._delete_for_ide(self.a.tool, sibling_document.document_id) + + # Row first: the delete lands before the file-store call, so without + # the predicate this is what actually goes missing. + assert DocumentManager.objects.filter( + pk=sibling_document.document_id + ).exists(), "sibling tool's document was deleted" + assert response.status_code == 404, getattr(response, "data", response) # --- the ordering fix, driven through the view ------------------------ @@ -276,7 +368,7 @@ def test_rejected_default_leaves_the_existing_default_intact(self): self.a.profile.refresh_from_db() assert self.a.profile.is_default - # --- A-4: the dead, state-changing-over-GET route is gone ------------- + # --- 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.""" diff --git a/backend/prompt_studio/tests/test_request_validation.py b/backend/prompt_studio/tests/test_request_validation.py new file mode 100644 index 0000000000..52865f85fe --- /dev/null +++ b/backend/prompt_studio/tests/test_request_validation.py @@ -0,0 +1,149 @@ +"""The 400/500 guards on the prompt-studio output and callback handlers. + +Each of these replaces a 500: a non-UUID id raises Django's ``ValidationError`` +while the query is being *built*, which drf_standardized_errors does not map; +a missing organization compiles to ``IS NULL`` and serves a blank project that +has real outputs; and a failed extraction-status write used to return 200. + +Every case here was checked against the unguarded code first — strip the guard +and the case fails. +""" + +import json +import secrets +import uuid +from unittest.mock import patch + +from account_v2.models import Organization, User +from django.test import TestCase +from prompt_studio.prompt_studio_index_manager_v2.prompt_studio_index_helper import ( + ExtractionStatusResult, +) +from prompt_studio.prompt_studio_output_manager_v2.views import PromptStudioOutputView +from rest_framework.test import APIRequestFactory, force_authenticate +from utils.user_context import UserContext + +_LIST = PromptStudioOutputView.as_view({"get": "list"}) +_DEFAULT_PROFILE = PromptStudioOutputView.as_view({"get": "get_output_for_tool_default"}) +_LATEST_BY_KEYS = PromptStudioOutputView.as_view({"get": "latest_outputs_by_keys"}) + + +class OutputViewValidationTest(TestCase): + def setUp(self): + self.addCleanup(UserContext.set_organization_identifier, None) + slug = f"val-{secrets.token_hex(3)}" + 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(), + ) + + def _get(self, view, path, params): + request = APIRequestFactory().get(path, params) + force_authenticate(request, user=self.user) + return view(request) + + def test_non_uuid_tool_id_on_default_profile_is_400(self): + response = self._get( + _DEFAULT_PROFILE, + "/prompt-output/prompt-default-profile/", + {"tool_id": "abc", "document_manager": str(uuid.uuid4())}, + ) + assert response.status_code == 400, response.data + + def test_non_uuid_document_manager_on_default_profile_is_400(self): + response = self._get( + _DEFAULT_PROFILE, + "/prompt-output/prompt-default-profile/", + {"tool_id": str(uuid.uuid4()), "document_manager": "abc"}, + ) + assert response.status_code == 400, response.data + + def test_absent_document_manager_on_default_profile_is_400(self): + """Absent used to filter on NULL, which renders every prompt as "". + + A 200 with a blank body is indistinguishable from a project that has + no outputs yet, so the caller has nothing to act on. + """ + response = self._get( + _DEFAULT_PROFILE, + "/prompt-output/prompt-default-profile/", + {"tool_id": str(uuid.uuid4())}, + ) + assert response.status_code == 400, response.data + + def test_non_uuid_tool_id_on_list_is_400(self): + """The list action is the highest-traffic one on this viewset.""" + response = self._get(_LIST, "/prompt-output/", {"tool_id": "abc"}) + assert response.status_code == 400, getattr(response, "data", response) + + def test_non_uuid_tool_id_on_latest_by_keys_is_400(self): + response = self._get( + _LATEST_BY_KEYS, + "/prompt-output/latest-by-keys/", + {"tool_id": "abc", "prompt_keys": "a"}, + ) + assert response.status_code == 400, response.data + + def test_unresolvable_organization_is_not_served_as_empty(self): + """Refuse rather than return a blank project that has real outputs.""" + UserContext.set_organization_identifier("org-that-does-not-exist") + response = self._get( + _DEFAULT_PROFILE, + "/prompt-output/prompt-default-profile/", + {"tool_id": str(uuid.uuid4()), "document_manager": str(uuid.uuid4())}, + ) + assert response.status_code >= 500, response.data + + +class ExtractionStatusEndpointTest(TestCase): + """The internal callback endpoint must not report a failed write as 200. + + The worker never reads the body, so a 200 drops the status silently and + every later Answer Prompt re-runs the full extraction. + """ + + URL = "/internal/v1/prompt-studio/extraction-status/" + + def _post(self, result): + from prompt_studio.prompt_studio_core_v2 import internal_views + + payload = { + "document_id": str(uuid.uuid4()), + "profile_manager_id": str(uuid.uuid4()), + "x2text_config_hash": "hash", + "enable_highlight": False, + "extracted": True, + } + request = APIRequestFactory().post( + self.URL, json.dumps(payload), content_type="application/json" + ) + with patch.object( + internal_views, "_parse_json_body", return_value=(payload, None) + ), patch( + "prompt_studio.prompt_profile_manager_v2.models.ProfileManager.objects" + ) as profiles, patch( + "prompt_studio.prompt_studio_index_manager_v2.prompt_studio_index_helper" + ".PromptStudioIndexHelper.mark_extraction_status", + return_value=result, + ): + profiles.get.return_value = object() + return internal_views.extraction_status(request) + + def test_write_failure_is_a_retryable_500(self): + assert self._post(ExtractionStatusResult.WRITE_FAILED).status_code == 500 + + def test_missing_document_is_a_non_retryable_404(self): + """500 is in the client's retry set; a gone document never comes back. + + Three retries with a 1s backoff factor would burn ~7s of worker sleep + on a condition no retry can change. + """ + assert self._post(ExtractionStatusResult.DOCUMENT_MISSING).status_code == 404 + + def test_success_is_200(self): + assert self._post(ExtractionStatusResult.OK).status_code == 200 diff --git a/backend/utils/models/org_aware_manager.py b/backend/utils/models/org_aware_manager.py index 33cdcd94d4..8329095617 100644 --- a/backend/utils/models/org_aware_manager.py +++ b/backend/utils/models/org_aware_manager.py @@ -69,11 +69,13 @@ def get_queryset(self): # 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. + # The RuntimeError arm is broader than its stated cause: + # StateStore compares an env string against a ConcurrencyMode + # member, which never matches, so it raises whenever + # CONCURRENCY_MODE is set at all — including to the documented + # "thread". Hence the log line. This path is rare (startup, + # migrations, tests), so it 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.", @@ -84,9 +86,22 @@ def get_queryset(self): return qs if org is None: - # 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. + if UserContext.get_organization_identifier(): + # An identifier is set but did not resolve to a row — + # Organization.DoesNotExist or ProgrammingError inside + # get_organization(), both of which it flattens to None. That + # happens *inside* a request, so returning everything here + # would cross tenants. Only the no-identifier case below is + # the "no request context" one. + logger.warning( + "OrgAwareManager: organization identifier is set but did " + "not resolve for %s; returning an empty queryset.", + self.model._meta.label, + ) + return qs.none() + # No request context at all: Celery, management commands, shell. + # Not logged — this is the normal state for every query those + # make, and a line per queryset would drown the cases 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 79969b323f..d2b3921231 100644 --- a/backend/utils/models/org_path_discovery.py +++ b/backend/utils/models/org_path_discovery.py @@ -51,8 +51,7 @@ "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. + # are excluded. This is the path already in force. "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 diff --git a/backend/utils/organization_utils.py b/backend/utils/organization_utils.py index 1d7add8c0c..63f0eca21b 100644 --- a/backend/utils/organization_utils.py +++ b/backend/utils/organization_utils.py @@ -75,17 +75,19 @@ def filter_queryset_by_organization(queryset, request, organization_field="organ """Filter a Django queryset by the request's organization context. 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. + in the request: each one reaches it with ``OrganizationFilterBackend`` + inert, either by opting out with ``skip_org_filter = True`` or by being on + a view class that declares no filter backends at all. Returning the + queryset unfiltered when there is no organization context would hand back + every organization's rows. + + Scope note: ``OrgAwareManager`` draws the same line, on the same + condition. It fails closed whenever an organization identifier is set but + does not resolve, and stays open only when no identifier is set at all — + Celery tasks, management commands and the shell, which have no request to + take one from. Callers overlap: the ``@csrf_exempt`` internal views reach + both. Fail-open there is the absence of a request, not the absence of a + header. 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 384de242ba..88534b23cf 100644 --- a/backend/utils/tests/test_org_path_discovery.py +++ b/backend/utils/tests/test_org_path_discovery.py @@ -13,39 +13,75 @@ from utils.models.org_path_discovery import ( ORG_PATH_OVERRIDES, _discover_org_path, + _org_path_cache, get_org_path, ) PINS = sorted(ORG_PATH_OVERRIDES.items()) +# The pin set itself, as a literal. PINS is derived from the dict under test, +# so every parametrized case below disappears along with a deleted pin and the +# file stays green while losing coverage. This is the only assertion here that +# a deletion cannot take with it. +EXPECTED_PINNED_LABELS = { + "prompt_studio_document_manager_v2.DocumentManager", + "prompt_studio_index_manager_v2.IndexManager", + "prompt_studio_output_manager_v2.PromptStudioOutputManager", + "prompt_studio_v2.ToolStudioPrompt", + "prompt_profile_manager_v2.ProfileManager", +} + # 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. # -# 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. +# Keyed by the model that *declares* the field, not by the pin it appears on — +# a hop belongs to the model the walk is standing on when it reads that name, +# and several pins share the same hop. The two terminal `organization` FKs are +# here because DefaultOrganizationMixin declares them null=True and save() +# backfills 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"), + ("prompt_studio_core_v2.CustomTool", "organization"), + ("adapter_processor_v2.AdapterInstance", "organization"), } +def test_pin_set_is_exactly_what_is_expected(): + """A pin removed or added has to be a deliberate edit here.""" + assert set(ORG_PATH_OVERRIDES) == EXPECTED_PINNED_LABELS + + @pytest.mark.parametrize("label,expected", PINS) def test_pin_is_returned(label, expected): - """get_org_path serves the pin, bypassing BFS.""" + """get_org_path serves the pin.""" assert get_org_path(apps.get_model(label)) == expected +@pytest.mark.parametrize("label,expected", PINS) +def test_pin_short_circuits_discovery(label, expected, monkeypatch): + """...and serves it *instead of* running BFS, not merely in agreement. + + BFS independently arrives at all five pins today, so the assertion above + holds with the short-circuit removed. Stubbing discovery to a value no pin + has is what separates the two. + """ + model = apps.get_model(label) + monkeypatch.setattr( + "utils.models.org_path_discovery._discover_org_path", + lambda _model: "sentinel__organization", + ) + # The memo would answer from an earlier test's real lookup and hide the + # stub, so drop this model's entry for the duration. + monkeypatch.delitem(_org_path_cache, model, raising=False) + assert get_org_path(model) == expected + + @pytest.mark.parametrize("label,expected", PINS) def test_pin_matches_discovery(label, expected): """The pin still agrees with what BFS would pick. @@ -73,9 +109,13 @@ def test_pin_traverses_only_non_nullable_fks(label, expected): 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 " - f"with a NULL {hop}. Pick a non-nullable path or add it to " + # Owner read before the walk advances: the key has to name the model + # that declares the field, which is what the failure message prints + # and what the reader has to add to the set. + owner = model._meta.label + assert not field.null or (owner, hop) in KNOWN_NULLABLE_HOPS, ( + f"{owner}.{hop} is nullable: this pin drops every row with a NULL " + f"{hop}. Pick a non-nullable path or add ({owner!r}, {hop!r}) to " f"KNOWN_NULLABLE_HOPS with a reason." ) model = field.related_model diff --git a/backend/utils/tests/test_organization_scoping.py b/backend/utils/tests/test_organization_scoping.py index e277e9f4d0..f0b924c29a 100644 --- a/backend/utils/tests/test_organization_scoping.py +++ b/backend/utils/tests/test_organization_scoping.py @@ -1,8 +1,8 @@ """``filter_queryset_by_organization`` must fail closed. -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. +Every caller reaches it with OrganizationFilterBackend inert, either by opting +out with ``skip_org_filter = True`` or by being on a view class that declares +no filter backends 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 @@ -18,11 +18,17 @@ from workflow_manager.workflow_v2.models.workflow import Workflow +_ABSENT = object() + + 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: + def __init__(self, organization_id=_ABSENT, path="/internal/test/"): + # Two distinct shapes reach production, which reads the attribute with + # getattr(..., None): the attribute missing entirely, and the attribute + # present and None. Defaulting to None would collapse them. + if organization_id is not _ABSENT: self.organization_id = organization_id self.path = path diff --git a/backend/workflow_manager/file_execution/internal_views.py b/backend/workflow_manager/file_execution/internal_views.py index 256331a0d7..83688cd26e 100644 --- a/backend/workflow_manager/file_execution/internal_views.py +++ b/backend/workflow_manager/file_execution/internal_views.py @@ -30,9 +30,8 @@ class FileExecutionInternalViewSet(viewsets.ModelViewSet): serializer_class = WorkflowFileExecutionSerializer lookup_field = "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. + # filter_queryset_by_organization, which fails closed. X-Organization-ID is + # therefore required in practice: a worker that omits it gets zero rows. 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 689a8e2b9f..58d1693c50 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -50,9 +50,8 @@ class WorkflowExecutionInternalViewSet(viewsets.ReadOnlyModelViewSet): serializer_class = WorkflowExecutionSerializer lookup_field = "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. + # filter_queryset_by_organization, which fails closed. X-Organization-ID is + # therefore required in practice: a worker that omits it gets zero rows. 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 aa31f0f671..c999b7f9f5 100644 --- a/backend/workflow_manager/workflow_v2/views.py +++ b/backend/workflow_manager/workflow_v2/views.py @@ -404,9 +404,8 @@ class WorkflowExecutionInternalViewSet(viewsets.ReadOnlyModelViewSet): serializer_class = WorkflowExecutionSerializer lookup_field = "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. + # filter_queryset_by_organization, which fails closed. X-Organization-ID is + # therefore required in practice: a worker that omits it gets zero rows. skip_org_filter = True def get_queryset(self): From bdd9ce8ce27562a35e3c73de6994c72eaa41990b Mon Sep 17 00:00:00 2001 From: Athul Date: Mon, 31 Aug 2026 17:32:44 +0530 Subject: [PATCH 08/10] UN-3815 [FIX] Rewrap a docstring line in the output-manager view Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01BGjvz6qwgMhQ7cfxT1Sj9W --- .../prompt_studio_output_manager_v2/views.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 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 e8cd1f9629..083de4f9b7 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py @@ -51,10 +51,10 @@ def _required_organization(tool_id: uuid.UUID) -> Organization: ``Organization.DoesNotExist`` and ``ProgrammingError``, neither logged. A None here compiles to ``organization_id IS NULL``, which matches only rows whose organization was never set and never the caller's tool — 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. + 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: From 3031e996f33ce103f7eae73fd20638176d39cf08 Mon Sep 17 00:00:00 2001 From: Athul Date: Mon, 31 Aug 2026 17:35:09 +0530 Subject: [PATCH 09/10] UN-3815 [FIX] Drop the last PR-relative phrasings from comments Three comments still described the code by contrast with what it replaced, which reads as archaeology once the change is history. Rewritten in the present tense to describe the code as it stands. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01BGjvz6qwgMhQ7cfxT1Sj9W --- .../prompt_profile_manager_v2/models.py | 8 +++---- .../tests/test_request_validation.py | 24 +++++++++++-------- 2 files changed, 18 insertions(+), 14 deletions(-) diff --git a/backend/prompt_studio/prompt_profile_manager_v2/models.py b/backend/prompt_studio/prompt_profile_manager_v2/models.py index 610618bd62..46d81ccbab 100644 --- a/backend/prompt_studio/prompt_profile_manager_v2/models.py +++ b/backend/prompt_studio/prompt_profile_manager_v2/models.py @@ -18,10 +18,10 @@ def for_user(self, user): Sharing only. Organization scoping comes from ``get_queryset``, which every branch here builds on, and which scopes through - ``vector_store__organization``. This method used to AND on - ``prompt_studio_tool__organization`` as well; that is a second, - narrower org join that drops tool-less profiles and contradicts the - null policy in ``get_queryset``. + ``vector_store__organization``. Do not AND on + ``prompt_studio_tool__organization`` here: that FK is nullable, so a + second org join drops tool-less profiles and contradicts the null + policy in ``get_queryset``. The parent-tool branch lets shared-project users see existing profiles so ``IsOwner`` on the viewset returns 403 on mutation diff --git a/backend/prompt_studio/tests/test_request_validation.py b/backend/prompt_studio/tests/test_request_validation.py index 52865f85fe..f1bff80a63 100644 --- a/backend/prompt_studio/tests/test_request_validation.py +++ b/backend/prompt_studio/tests/test_request_validation.py @@ -3,7 +3,7 @@ Each of these replaces a 500: a non-UUID id raises Django's ``ValidationError`` while the query is being *built*, which drf_standardized_errors does not map; a missing organization compiles to ``IS NULL`` and serves a blank project that -has real outputs; and a failed extraction-status write used to return 200. +has real outputs; and a failed extraction-status write reports success. Every case here was checked against the unguarded code first — strip the guard and the case fails. @@ -64,7 +64,7 @@ def test_non_uuid_document_manager_on_default_profile_is_400(self): assert response.status_code == 400, response.data def test_absent_document_manager_on_default_profile_is_400(self): - """Absent used to filter on NULL, which renders every prompt as "". + """An absent id would filter on NULL, rendering every prompt as "". A 200 with a blank body is indistinguishable from a project that has no outputs yet, so the caller has nothing to act on. @@ -122,14 +122,18 @@ def _post(self, result): request = APIRequestFactory().post( self.URL, json.dumps(payload), content_type="application/json" ) - with patch.object( - internal_views, "_parse_json_body", return_value=(payload, None) - ), patch( - "prompt_studio.prompt_profile_manager_v2.models.ProfileManager.objects" - ) as profiles, patch( - "prompt_studio.prompt_studio_index_manager_v2.prompt_studio_index_helper" - ".PromptStudioIndexHelper.mark_extraction_status", - return_value=result, + with ( + patch.object( + internal_views, "_parse_json_body", return_value=(payload, None) + ), + patch( + "prompt_studio.prompt_profile_manager_v2.models.ProfileManager.objects" + ) as profiles, + patch( + "prompt_studio.prompt_studio_index_manager_v2.prompt_studio_index_helper" + ".PromptStudioIndexHelper.mark_extraction_status", + return_value=result, + ), ): profiles.get.return_value = object() return internal_views.extraction_status(request) From 09aec4e6c62ec7c3d80d50a4be40486c9bec1895 Mon Sep 17 00:00:00 2001 From: Athul Date: Fri, 4 Sep 2026 19:18:00 +0530 Subject: [PATCH 10/10] UN-3815 [FIX] Address review threads on prompt-studio org scoping MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four handlers reported a scoped lookup that resolved nothing as if it had worked. prompt_output filtered ToolStudioPrompt through the nullable tool_id__organization pin and handed the result to a helper that early-exits on an empty list, so a run whose prompts the scope hid answered 200 with an empty body — and the worker never reads the body on success. It now compares resolved to requested and 404s on a mismatch. ProfileManager.objects is scoped the same way, but every profile_manager_id lookup sat outside a narrowed except and fell to a bare 500 — inside the worker client's {500,502,503,504} retry set, ~7s of sleep for a condition no retry changes. _resolve_profile maps DoesNotExist and Django's ValidationError (malformed id) to the same 404 at all four call sites. mark_extraction_status folds ValidationError into DOCUMENT_MISSING for the same reason, and extraction_status now dispatches ExtractionStatusResult member by member so a new one cannot be absorbed into the 500 branch. delete_for_ide deleted the row before the object store with no shared transaction, so a file-store failure returned "File deletion failed." for a document already gone. The file goes first; now the error is accurate and the retry is real. Tests the reviewers asked for: the six internal viewsets' only tenant boundary driven through WorkflowExecutionInternalViewSet and FileExecutionInternalViewSet (including the get_object() path that bypasses get_queryset()); ProfileManager. for_user across all three branches, including shared_to_org=True; the output read endpoints against real cross-org ids. Verified by removing the guards. Also: share _validated_uuid as utils.uuid_validation.validated_uuid, freeze ORG_PATH_OVERRIDES behind MappingProxyType, stop rebinding `result` across two types, drop FileInfoIdeSerializer and FileManagerHelper.delete_file (dead since the /file/delete route went, no consumer in OSS or cloud), correct the get_object() comment on the one viewset that overrides it, and say plainly that the explicit tool_id__organization kwargs are defence in depth. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01QKgDApDSSdFvGwuk7k7695 --- .../file_management/file_management_helper.py | 8 - backend/file_management/serializer.py | 6 - .../prompt_studio_core_v2/internal_views.py | 137 ++++++++++--- .../prompt_studio_helper.py | 11 +- .../test_prompt_output_outputs_validation.py | 15 +- .../prompt_studio_core_v2/views.py | 20 +- .../prompt_studio_index_helper.py | 19 +- .../prompt_studio_output_manager_v2/views.py | 41 ++-- .../tests/test_cross_org_isolation.py | 193 +++++++++++++----- .../tests/test_request_validation.py | 155 ++++++++++++++ backend/utils/models/org_path_discovery.py | 38 ++-- .../tests/test_organization_scoping_views.py | 134 ++++++++++++ backend/utils/uuid_validation.py | 23 +++ .../file_execution/internal_views.py | 9 +- 14 files changed, 644 insertions(+), 165 deletions(-) create mode 100644 backend/utils/tests/test_organization_scoping_views.py create mode 100644 backend/utils/uuid_validation.py diff --git a/backend/file_management/file_management_helper.py b/backend/file_management/file_management_helper.py index 008368d09d..51cc141809 100644 --- a/backend/file_management/file_management_helper.py +++ b/backend/file_management/file_management_helper.py @@ -225,14 +225,6 @@ def _get_base_path(file_system: UnstractFileSystem, path: str): base_path = base_path.rstrip("/") + "/" return base_path - @staticmethod - def delete_file(file_system: UnstractFileSystem, path: str, file_name: str) -> bool: - fs = file_system.get_fsspec_fs() - base_path = FileManagerHelper._get_base_path(file_system, path) - file_path = str(Path(base_path) / file_name) - FileManagerHelper._delete_file(fs, file_path) - return True - @staticmethod def delete_related_files( file_system: UnstractFileSystem, diff --git a/backend/file_management/serializer.py b/backend/file_management/serializer.py index 4d2890230c..6cc40ad1a7 100644 --- a/backend/file_management/serializer.py +++ b/backend/file_management/serializer.py @@ -50,11 +50,5 @@ class FileUploadIdeSerializer(serializers.Serializer): ) -class FileInfoIdeSerializer(serializers.Serializer): - document_id = serializers.CharField() - tool_id = serializers.CharField() - view_type = serializers.CharField(required=False) - - class FileListRequestIdeSerializer(serializers.Serializer): tool_id = serializers.CharField() 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 b5fdab2f75..63fca1ad7a 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/internal_views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/internal_views.py @@ -24,6 +24,36 @@ _ERR_INVALID_JSON = "Invalid JSON" +def _resolve_profile(profile_manager_id): + """Return ``(profile, None)``, or ``(None, JsonResponse)`` when it is absent. + + ``ProfileManager.objects`` is org-scoped on ``vector_store__organization``, + so a miss here has the same two causes as a missing document — the row is + gone, or the org scope hides it from this caller — and neither is fixed by + trying again. A malformed id raises Django's ``ValidationError`` while the + query is built, which is just as permanent. + + Left to the generic ``except Exception`` these all became a 500, which the + worker's client retries three times with a 1s backoff factor: ~7s of worker + sleep on a condition no retry can change. 404 is outside that retry set. + """ + from django.core.exceptions import ValidationError + + from prompt_studio.prompt_profile_manager_v2.models import ProfileManager + + try: + return ProfileManager.objects.get(pk=profile_manager_id), None + except (ProfileManager.DoesNotExist, ValidationError): + logger.error( + "Profile manager %s not found or not visible in the current " "organization.", + profile_manager_id, + ) + return None, JsonResponse( + {"success": False, "error": "Profile manager not found"}, + status=status.HTTP_404_NOT_FOUND, + ) + + def _parse_json_body(request): """Parse JSON from request body, returning (data, None) or (None, JsonResponse).""" try: @@ -107,6 +137,36 @@ def prompt_output(request): ) ) + # ``ToolStudioPrompt.objects`` is org-scoped on the nullable + # ``tool_id__organization``, so this filter can return fewer prompts + # than were asked for — or none. handle_prompt_output_update early-exits + # on an empty list, so the endpoint would answer 200 with an empty body + # and the worker, which never reads the body on success, would discard + # every prompt output for the run. Same reasoning as extraction_status + # below: fail loudly, and outside the client's {500,502,503,504} retry + # set, because no retry resolves a prompt the scope hides. + requested = set(prompt_ids) + if len(prompts) != len(requested): + resolved = {str(p.prompt_id) for p in prompts} + logger.error( + "prompt_output: %d of %d prompts resolved for document %s; " + "unresolved=%s. Refusing to persist a partial run.", + len(prompts), + len(requested), + document_id, + sorted(str(p) for p in requested - resolved), + ) + return JsonResponse( + { + "success": False, + "error": ( + "One or more prompts were not found or are not visible " + "in the current organization" + ), + }, + status=status.HTTP_404_NOT_FOUND, + ) + response = OutputManagerHelper.handle_prompt_output_update( run_id=run_id, prompts=prompts, @@ -158,12 +218,13 @@ def index_update(request): ) try: - from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_index_manager_v2.prompt_studio_index_helper import ( PromptStudioIndexHelper, ) - profile_manager = ProfileManager.objects.get(pk=profile_manager_id) + profile_manager, err = _resolve_profile(profile_manager_id) + if err: + return err PromptStudioIndexHelper.handle_index_manager( document_id=document_id, profile_manager=profile_manager, @@ -223,13 +284,14 @@ def extraction_status(request): ) try: - from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_index_manager_v2.prompt_studio_index_helper import ( ExtractionStatusResult, PromptStudioIndexHelper, ) - profile_manager = ProfileManager.objects.get(pk=profile_manager_id) + profile_manager, err = _resolve_profile(profile_manager_id) + if err: + return err result = PromptStudioIndexHelper.mark_extraction_status( document_id=document_id, profile_manager=profile_manager, @@ -238,38 +300,44 @@ def extraction_status(request): extracted=extracted, error_message=error_message, ) + # A 200 on anything but OK 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. + # + # The two failures need different statuses. The client retries + # {500, 502, 503, 504} three times with a 1s backoff factor, so a + # document that is gone would burn four round trips and ~7s of worker + # sleep on a condition no retry can change. 404 is outside that set. + # + # Matched member by member with no wildcard branch: a fourth member + # added later raises here instead of being absorbed into the 500 case. if result is not ExtractionStatusResult.OK: - # 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. - # - # The two failures need different statuses. The client retries - # {500, 502, 503, 504} three times with a 1s backoff factor, so a - # document that is gone would burn four round trips and ~7s of - # worker sleep on a condition no retry can change. 404 is outside - # that set. - missing = result is ExtractionStatusResult.DOCUMENT_MISSING logger.error( "extraction_status not recorded for document %s profile %s (%s)", document_id, profile_manager_id, result.value, ) - return JsonResponse( - { - "success": False, - "error": "Document not found" - if missing - else "Extraction status could not be recorded", - }, - status=status.HTTP_404_NOT_FOUND - if missing - else status.HTTP_500_INTERNAL_SERVER_ERROR, - ) - - return JsonResponse({"success": True}) + match result: + case ExtractionStatusResult.OK: + return JsonResponse({"success": True}) + case ExtractionStatusResult.DOCUMENT_MISSING: + return JsonResponse( + {"success": False, "error": "Document not found"}, + status=status.HTTP_404_NOT_FOUND, + ) + case ExtractionStatusResult.WRITE_FAILED: + return JsonResponse( + { + "success": False, + "error": "Extraction status could not be recorded", + }, + status=status.HTTP_500_INTERNAL_SERVER_ERROR, + ) + case unhandled: + raise AssertionError(f"Unhandled ExtractionStatusResult: {unhandled!r}") except Exception as e: logger.exception("extraction_status internal API failed") @@ -353,9 +421,9 @@ def profile_detail(request, profile_id): Returns vector_store, embedding_model, x2text adapter IDs and chunk_overlap. """ try: - from prompt_studio.prompt_profile_manager_v2.models import ProfileManager - - profile = ProfileManager.objects.get(pk=profile_id) + profile, err = _resolve_profile(profile_id) + if err: + return err return JsonResponse( { "success": True, @@ -463,7 +531,6 @@ def summary_index_key(request): try: from utils.file_storage.constants import FileStorageKeys - from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.prompt_ide_base_tool import ( PromptIdeBaseTool, ) @@ -472,7 +539,9 @@ def summary_index_key(request): from unstract.sdk1.file_storage.env_helper import EnvHelper from unstract.sdk1.utils.indexing import IndexingUtils - profile = ProfileManager.objects.get(pk=summary_profile_id) + profile, err = _resolve_profile(summary_profile_id) + if err: + return err fs_instance = EnvHelper.get_storage( storage_type=StorageType.PERMANENT, env_name=FileStorageKeys.PERMANENT_REMOTE_STORAGE, diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index d209764b0d..774d6e19e3 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -2575,7 +2575,7 @@ def dynamic_extractor( result = dispatcher.dispatch(extract_context) if not result.success: msg = result.error or "Unknown extraction error" - result = PromptStudioIndexHelper.mark_extraction_status( + status_result = PromptStudioIndexHelper.mark_extraction_status( document_id=document_id, profile_manager=profile_manager, x2text_config_hash=x2text_config_hash, @@ -2583,7 +2583,7 @@ def dynamic_extractor( extracted=False, error_message=msg, ) - if result is not ExtractionStatusResult.OK: + if status_result is not ExtractionStatusResult.OK: logger.warning( f"Failed to mark extraction failure for document {document_id}. " f"Extraction failed but status not saved." @@ -2593,13 +2593,16 @@ def dynamic_extractor( ) extracted_text = result.data.get("extracted_text", "") - result = PromptStudioIndexHelper.mark_extraction_status( + # Distinct name: ``result`` is the dispatcher's ExecutionResult and is + # still read above. Rebinding it to an ExtractionStatusResult made + # ``result.data`` correct only by branch ordering. + status_result = PromptStudioIndexHelper.mark_extraction_status( document_id=document_id, profile_manager=profile_manager, x2text_config_hash=x2text_config_hash, enable_highlight=enable_highlight, ) - if result is not ExtractionStatusResult.OK: + if status_result is not ExtractionStatusResult.OK: logger.warning( f"Failed to mark extraction success for document {document_id}. " f"Extraction completed but status not saved." diff --git a/backend/prompt_studio/prompt_studio_core_v2/tests/test_prompt_output_outputs_validation.py b/backend/prompt_studio/prompt_studio_core_v2/tests/test_prompt_output_outputs_validation.py index 34f7dbea7a..31f35eb843 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/tests/test_prompt_output_outputs_validation.py +++ b/backend/prompt_studio/prompt_studio_core_v2/tests/test_prompt_output_outputs_validation.py @@ -13,6 +13,7 @@ """ import json +from types import SimpleNamespace from unittest.mock import MagicMock, patch from prompt_studio.prompt_studio_core_v2.internal_views import prompt_output @@ -41,6 +42,16 @@ def _body(response): return json.loads(response.content) +def _resolved(*prompt_ids): + """What the ORM returns on the happy path. + + prompt_output now compares resolved count to requested count — an org scope + that drops rows must not answer 200 — so the accept-path stubs have to + resolve every id they ask for. Only ``prompt_id`` is read on that path. + """ + return [SimpleNamespace(prompt_id=pid) for pid in prompt_ids] + + def test_list_outputs_rejected_with_400_and_a_reason(): """The regression: this used to reach the helper and raise AttributeError.""" response = prompt_output(_request([{"invoice_number": "INV-001"}, {"b": 2}])) @@ -105,7 +116,7 @@ def test_dict_outputs_still_reach_the_helper(): "prompt_studio.prompt_studio_output_manager_v2." "output_manager_helper.OutputManagerHelper.handle_prompt_output_update" ) as handler: - prompts.filter.return_value.order_by.return_value = [] + prompts.filter.return_value.order_by.return_value = _resolved("p1") handler.return_value = [] response = prompt_output(_request({"invoice_number": "INV-001"})) handler.assert_called_once() @@ -123,7 +134,7 @@ def test_missing_outputs_defaults_to_empty_dict_and_is_accepted(): "prompt_studio.prompt_studio_output_manager_v2." "output_manager_helper.OutputManagerHelper.handle_prompt_output_update" ) as handler: - prompts.filter.return_value.order_by.return_value = [] + prompts.filter.return_value.order_by.return_value = _resolved("p1") handler.return_value = [] response = prompt_output(request) assert response.status_code == 200 diff --git a/backend/prompt_studio/prompt_studio_core_v2/views.py b/backend/prompt_studio/prompt_studio_core_v2/views.py index 72318e4e3e..f889789a26 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/views.py @@ -32,6 +32,7 @@ from utils.pagination import OptionalPagination from utils.user_context import UserContext from utils.user_session import UserSessionUtils +from utils.uuid_validation import validated_uuid from prompt_studio.lookup_utils import ( get_latest_lookup_mutation_for_tool, @@ -394,10 +395,7 @@ def make_profile_default(self, request: HttpRequest, pk: Any = None) -> Response 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.") + default_profile = validated_uuid(default_profile, "default_profile") # Resolve the target before clearing anything. The transaction below # is what actually guarantees the tool never ends up with zero @@ -1168,9 +1166,15 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response: DocumentIndexingService.remove_document_indexing( org_id=org_id, user_id=user_id, doc_id_key=raw_index_id ) - # Delete the document record - document.delete() - # Delete the files + # Object store first, then the row. The two share no transaction, + # so the order decides which partial failure the except block below + # can report honestly. Deleting the row first and failing on the + # file returned "File deletion failed." (400, reads as "nothing + # happened") for a document already permanently gone, leaving the + # file orphaned with nothing left to retry against. This way a + # failed file delete leaves the row intact and the retry is real; + # the residue in the other direction is a row whose file is already + # gone, which the next delete clears. file_name: str = document.document_name PromptStudioFileHelper.delete_for_ide( org_id=org_id, @@ -1178,6 +1182,8 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response: tool_id=str(custom_tool.tool_id), file_name=file_name, ) + # Delete the document record + document.delete() return Response( {"data": "File deleted succesfully."}, status=status.HTTP_200_OK, 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 0a4530fea3..9125da9693 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 @@ -2,7 +2,7 @@ import logging from enum import Enum -from django.core.exceptions import ImproperlyConfigured +from django.core.exceptions import ImproperlyConfigured, ValidationError from django.db import DatabaseError, transaction from prompt_studio.prompt_profile_manager_v2.models import ProfileManager @@ -106,9 +106,9 @@ def mark_extraction_status( Returns: ExtractionStatusResult: OK on success; DOCUMENT_MISSING when the - document is gone or hidden by the org scope; WRITE_FAILED for - a genuine write error. Callers that only need "did it write" - compare against OK. + document is gone, hidden by the org scope, or named by a + malformed id; WRITE_FAILED for a genuine write error. Callers + that only need "did it write" compare against OK. """ try: @@ -168,10 +168,13 @@ def mark_extraction_status( return ExtractionStatusResult.OK - except DocumentManager.DoesNotExist: - # Two ways to get here: the row is gone, or the org-scoped manager - # hides it from this caller. Both mean the status was not written, - # and neither is fixed by trying again. + except (DocumentManager.DoesNotExist, ValidationError): + # Three ways to get here: the row is gone, the org-scoped manager + # hides it from this caller, or ``document_id`` is not a UUID and + # ``get(pk=...)`` raised while building the query. All three mean + # the status was not written and none is fixed by trying again, so + # they share the 404 branch rather than the retryable 500 an + # unmapped ValidationError used to produce. logger.error( "Document %s not found or not visible in the current " "organization; extraction status not recorded.", 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 083de4f9b7..3ef714fa44 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py @@ -12,6 +12,7 @@ from utils.common_utils import CommonUtils from utils.filtering import FilterHelper from utils.user_context import UserContext +from utils.uuid_validation import validated_uuid from prompt_studio.prompt_studio_output_manager_v2.constants import ( PromptOutputManagerErrorMessage, @@ -30,20 +31,6 @@ logger = logging.getLogger(__name__) -def _validated_uuid(raw: Any, field_name: str) -> uuid.UUID: - """A query-string id as a UUID, or a 400 naming the field. - - These columns are UUID primary/foreign keys, 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=f"'{field_name}' must be a valid UUID.") - - def _required_organization(tool_id: uuid.UUID) -> Organization: """The request's organization, refusing to proceed without one. @@ -96,7 +83,7 @@ def get_queryset(self) -> QuerySet | None: PromptStudioOutputManagerKeys.DOCUMENT_MANAGER, ): if key in filter_args: - filter_args[key] = _validated_uuid(filter_args[key], key) + filter_args[key] = validated_uuid(filter_args[key], key) # Convert the string representation to a boolean value is_single_pass_extract = CommonUtils.str_to_bool(is_single_pass_extract_param) @@ -128,10 +115,17 @@ def latest_outputs_by_keys(self, request: HttpRequest) -> Response: if not prompt_keys: return Response({}, status=status.HTTP_200_OK) - tool_id = _validated_uuid(tool_id, PromptStudioOutputManagerKeys.TOOL_ID) - - # A raw .objects query is not routed through filter_queryset(), so - # OrganizationFilterBackend does not see it — scope explicitly. + tool_id = validated_uuid(tool_id, PromptStudioOutputManagerKeys.TOOL_ID) + + # Defence in depth, not the only scope. Since these models moved to + # OrgAwareManager (pinned to tool_id__organization in + # ORG_PATH_OVERRIDES) the manager appends this same predicate on its + # own, resolving the org through the same UserContext call + # _required_organization() makes. Kept explicit because a raw .objects + # query is not routed through filter_queryset(), so the view layer + # contributes nothing here and the manager pin would be the single + # point of failure. test_cross_org_isolation pins the manager + # independently, so removing these kwargs stays a safe follow-up. organization = _required_organization(tool_id) prompt_id_to_key = dict( ToolStudioPrompt.objects.filter( @@ -174,7 +168,7 @@ def get_output_for_tool_default(self, request: HttpRequest) -> Response: if not tool_id: raise ValidationError(detail=tool_validation_message) - tool_id = _validated_uuid(tool_id, PromptStudioOutputManagerKeys.TOOL_ID) + tool_id = validated_uuid(tool_id, PromptStudioOutputManagerKeys.TOOL_ID) # Same column type, same failure: this one reaches # PromptStudioOutputManager.objects.filter(document_manager_id=...) in # the helper below. Required, not optional — absent it stays None and @@ -185,14 +179,15 @@ def get_output_for_tool_default(self, request: HttpRequest) -> Response: raise ValidationError( detail="'document_manager' is required and must be a valid UUID." ) - document_manager_id = _validated_uuid( + document_manager_id = validated_uuid( document_manager_id, PromptStudioOutputManagerKeys.DOCUMENT_MANAGER ) organization = _required_organization(tool_id) # Fetch ToolStudioPrompt records based on tool_id. - # A raw .objects query is not routed through filter_queryset(), so - # OrganizationFilterBackend does not see it — scope explicitly. + # Defence in depth, as above: OrgAwareManager already pins this model + # to tool_id__organization, and a raw .objects query gets nothing from + # the view layer. # # No exception handling below: for a valid UUID that matches no row, or # a tool in another organization, filter() returns empty rather than diff --git a/backend/prompt_studio/tests/test_cross_org_isolation.py b/backend/prompt_studio/tests/test_cross_org_isolation.py index 44082611fd..5519aedec8 100644 --- a/backend/prompt_studio/tests/test_cross_org_isolation.py +++ b/backend/prompt_studio/tests/test_cross_org_isolation.py @@ -12,6 +12,7 @@ """ import secrets +from unittest.mock import patch import pytest from account_v2.models import Organization, User @@ -32,9 +33,33 @@ from prompt_studio.prompt_studio_output_manager_v2.models import ( PromptStudioOutputManager, ) +from prompt_studio.prompt_studio_output_manager_v2.views import PromptStudioOutputView from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt +def _make_profile(name, tool, adapter, *, is_default, created_by) -> ProfileManager: + """A ProfileManager with the fields none of these tests vary. + + Only the name, tool, adapter and default flag ever differ between call + sites; spelling out the other nine each time hid that. + """ + return ProfileManager.objects.create( + profile_name=name, + 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=tool, + is_default=is_default, + created_by=created_by, + ) + + class OrgFixture: """One organization with a fully populated prompt-studio object graph.""" @@ -55,19 +80,11 @@ def __init__(self, slug: str): 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, + self.adapter = self._adapter(slug) + self.profile = _make_profile( + f"profile-{slug}", + self.tool, + self.adapter, is_default=True, created_by=self.user, ) @@ -202,18 +219,10 @@ def _sibling_tool_in_org_a(self): organization=self.a.org, created_by=self.a.user, ) - profile = ProfileManager.objects.create( - profile_name=f"sibling-profile-{secrets.token_hex(3)}", - 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=tool, + profile = _make_profile( + f"sibling-profile-{secrets.token_hex(3)}", + tool, + self.a.adapter, is_default=True, created_by=self.a.user, ) @@ -224,14 +233,23 @@ def _sibling_tool_in_org_a(self): ) return tool, profile, document - def _delete_for_ide(self, tool, document_id): - """DELETE prompt-studio/file/ as the owner of ``tool``.""" + def _grant_owner(self, tool) -> None: + """The OWNER row the create *view* writes. + + ``CustomTool.objects.for_user`` resolves visibility through + ResourceMembership; the fixtures build rows directly, so without this + ``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), ) + + def _delete_for_ide(self, tool, document_id): + """DELETE prompt-studio/file/ as the owner of ``tool``.""" + self._grant_owner(tool) view = PromptStudioCoreView.as_view({"delete": "delete_for_ide"}) request = APIRequestFactory().delete( f"/prompt-studio/file/{tool.tool_id}", @@ -277,8 +295,9 @@ def test_delete_for_ide_refuses_a_sibling_tool_document(self): response = self._delete_for_ide(self.a.tool, sibling_document.document_id) - # Row first: the delete lands before the file-store call, so without - # the predicate this is what actually goes missing. + # The row is what the predicate protects: get_object_or_404 above runs + # before either delete, so a missing predicate is what would let this + # row be removed at all. assert DocumentManager.objects.filter( pk=sibling_document.document_id ).exists(), "sibling tool's document was deleted" @@ -287,19 +306,8 @@ def test_delete_for_ide_refuses_a_sibling_tool_document(self): # --- 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), - ) + """PATCH make_profile_default as the owner of ``tool``.""" + self._grant_owner(tool) view = PromptStudioCoreView.as_view({"patch": "make_profile_default"}) request = APIRequestFactory().patch( f"/prompt-studio/{tool.tool_id}/make_profile_default", @@ -310,18 +318,10 @@ def _make_profile_default(self, tool, profile_id): 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, + return _make_profile( + "profile-a-second", + self.a.tool, + self.a.adapter, is_default=False, created_by=self.a.user, ) @@ -368,6 +368,91 @@ def test_rejected_default_leaves_the_existing_default_intact(self): self.a.profile.refresh_from_db() assert self.a.profile.is_default + # --- ProfileManager.for_user: sharing must not widen the org scope ----- + + def _profile_ids_for_user(self, user): + return {str(p.profile_id) for p in ProfileManager.objects.for_user(user)} + + def test_for_user_excludes_another_org_even_when_shared_to_org(self): + """``for_user`` dropped its explicit tool-org filter and now leans + entirely on ``get_queryset()``. ``shared_to_org=True`` is the branch + that would otherwise match every organization's rows at once. + """ + self.b.profile.shared_to_org = True + self.b.profile.save(update_fields=["shared_to_org"]) + + visible = self._profile_ids_for_user(self.a.user) + + assert str(self.a.profile.profile_id) in visible + assert str(self.b.profile.profile_id) not in visible + + def test_for_user_service_account_branch_is_still_org_scoped(self): + """``self.all()`` is not ``_base_manager.all()`` — it inherits the + scope. A branch rewritten to bypass ``self`` would fail here.""" + self.a.user.is_service_account = True + + visible = self._profile_ids_for_user(self.a.user) + + assert str(self.a.profile.profile_id) in visible + assert str(self.b.profile.profile_id) not in visible + + def test_for_user_org_admin_branch_is_still_org_scoped(self): + """Admin of org A is not admin of every org.""" + with patch( + "prompt_studio.prompt_profile_manager_v2.models." + "OrganizationMemberService.is_user_organization_admin", + return_value=True, + ): + visible = self._profile_ids_for_user(self.a.user) + + assert str(self.a.profile.profile_id) in visible + assert str(self.b.profile.profile_id) not in visible + + # --- the output read endpoints, against real cross-org ids ------------- + + def _output_view(self, action: str, params: dict): + view = PromptStudioOutputView.as_view({"get": action}) + request = APIRequestFactory().get("/prompt-studio/output", params) + force_authenticate(request, user=self.a.user) + return view(request) + + def test_latest_outputs_by_keys_refuses_another_orgs_tool(self): + """A real org-B tool id, not just an unmatched UUID. + + The 400 cases already covered only prove the id is parsed. This is what + pins the scoping: org B's prompt key exists and holds a real output, so + an unscoped query would return it. + """ + response = self._output_view( + "latest_outputs_by_keys", + {"tool_id": str(self.b.tool.tool_id), "prompt_keys": self.b.prompt.prompt_key}, + ) + + assert response.status_code == 200, response.data + assert response.data == {} + + def test_latest_outputs_by_keys_still_returns_own_output(self): + """Without this, a query scoped to nothing would pass above.""" + response = self._output_view( + "latest_outputs_by_keys", + {"tool_id": str(self.a.tool.tool_id), "prompt_keys": self.a.prompt.prompt_key}, + ) + + assert response.status_code == 200, response.data + assert response.data == {self.a.prompt.prompt_key: "secret"} + + def test_get_output_for_tool_default_refuses_another_orgs_ids(self): + response = self._output_view( + "get_output_for_tool_default", + { + "tool_id": str(self.b.tool.tool_id), + "document_manager": str(self.b.document.document_id), + }, + ) + + assert response.status_code == 200, response.data + assert response.data == {} + # --- the dead, state-changing-over-GET route is gone ------------------- def test_file_delete_route_removed(self): diff --git a/backend/prompt_studio/tests/test_request_validation.py b/backend/prompt_studio/tests/test_request_validation.py index f1bff80a63..b4cd843a8d 100644 --- a/backend/prompt_studio/tests/test_request_validation.py +++ b/backend/prompt_studio/tests/test_request_validation.py @@ -12,6 +12,7 @@ import json import secrets import uuid +from types import SimpleNamespace from unittest.mock import patch from account_v2.models import Organization, User @@ -151,3 +152,157 @@ def test_missing_document_is_a_non_retryable_404(self): def test_success_is_200(self): assert self._post(ExtractionStatusResult.OK).status_code == 200 + + +class InternalCallbackResolutionTest(TestCase): + """An id the org-scoped managers cannot resolve must not become a 500. + + Both handlers resolve ids through managers that are now org-scoped, so a + lookup can miss for two reasons — the row is gone, or the scope hides it — + and neither changes on retry. The worker's client retries {500,502,503,504} + three times with a 1s backoff factor, so a 500 here costs ~7s of worker + sleep and still fails. 404 is outside that set. + """ + + def _payload(self, **overrides): + payload = { + "run_id": str(uuid.uuid4()), + "prompt_ids": [str(uuid.uuid4()), str(uuid.uuid4())], + "outputs": {}, + "document_id": str(uuid.uuid4()), + "is_single_pass_extract": False, + "profile_manager_id": str(uuid.uuid4()), + "metadata": {}, + } + payload.update(overrides) + return payload + + # --- prompt_output: a partial resolve is not a success ------------------ + + @staticmethod + def _prompt(prompt_id=None): + """Only the attribute the mismatch log reads is needed here.""" + return SimpleNamespace(prompt_id=prompt_id or uuid.uuid4()) + + def _prompt_output(self, resolved): + """Drive prompt_output with ``resolved`` prompts coming back.""" + from prompt_studio.prompt_studio_core_v2 import internal_views + + payload = self._payload() + request = APIRequestFactory().post( + "/internal/v1/prompt-studio/prompt-output/", + json.dumps(payload), + content_type="application/json", + ) + with ( + patch.object( + internal_views, "_parse_json_body", return_value=(payload, None) + ), + patch( + "prompt_studio.prompt_studio_v2.models.ToolStudioPrompt.objects" + ) as prompts, + patch( + "prompt_studio.prompt_studio_output_manager_v2.output_manager_helper" + ".OutputManagerHelper.handle_prompt_output_update", + return_value={}, + ) as handler, + ): + prompts.filter.return_value.order_by.return_value = resolved + return internal_views.prompt_output(request), handler + + def test_prompt_output_refuses_when_the_scope_drops_every_prompt(self): + """handle_prompt_output_update early-exits on an empty list, so this + used to answer 200 with an empty body and discard the whole run.""" + response, handler = self._prompt_output([]) + + assert response.status_code == 404, response.content + handler.assert_not_called() + + def test_prompt_output_refuses_a_partial_resolve(self): + """One of two prompts resolved is still a run that would be recorded + wrong; the count comparison is what catches it.""" + response, handler = self._prompt_output([self._prompt()]) + + assert response.status_code == 404, response.content + handler.assert_not_called() + + def test_prompt_output_proceeds_when_every_prompt_resolves(self): + """Without this, a handler that refused everything would pass above.""" + response, handler = self._prompt_output([self._prompt(), self._prompt()]) + + assert response.status_code == 200, response.content + handler.assert_called_once() + + # --- the profile lookups: 404, not the generic 500 ---------------------- + + def _with_missing_profile(self, handler_name, payload): + from prompt_studio.prompt_profile_manager_v2.models import ProfileManager + from prompt_studio.prompt_studio_core_v2 import internal_views + + request = APIRequestFactory().post( + "/internal/v1/prompt-studio/", json.dumps(payload), + content_type="application/json", + ) + with ( + patch.object( + internal_views, "_parse_json_body", return_value=(payload, None) + ), + patch( + "prompt_studio.prompt_profile_manager_v2.models.ProfileManager.objects" + ) as profiles, + ): + profiles.get.side_effect = ProfileManager.DoesNotExist + return getattr(internal_views, handler_name)(request) + + def test_extraction_status_with_an_unresolvable_profile_is_404(self): + response = self._with_missing_profile( + "extraction_status", + { + "document_id": str(uuid.uuid4()), + "profile_manager_id": str(uuid.uuid4()), + "x2text_config_hash": "hash", + "enable_highlight": False, + "extracted": True, + }, + ) + assert response.status_code == 404, response.content + + def test_index_update_with_an_unresolvable_profile_is_404(self): + response = self._with_missing_profile( + "index_update", + { + "document_id": str(uuid.uuid4()), + "profile_manager_id": str(uuid.uuid4()), + "doc_id": "doc-1", + }, + ) + assert response.status_code == 404, response.content + + def test_a_malformed_profile_id_is_404_not_a_retryable_500(self): + """``get(pk=...)`` raises Django's ValidationError, not DoesNotExist, + for an id that is not a UUID — just as permanent.""" + from django.core.exceptions import ValidationError as DjangoValidationError + + from prompt_studio.prompt_studio_core_v2 import internal_views + + payload = { + "document_id": str(uuid.uuid4()), + "profile_manager_id": "not-a-uuid", + "doc_id": "doc-1", + } + request = APIRequestFactory().post( + "/internal/v1/prompt-studio/", json.dumps(payload), + content_type="application/json", + ) + with ( + patch.object( + internal_views, "_parse_json_body", return_value=(payload, None) + ), + patch( + "prompt_studio.prompt_profile_manager_v2.models.ProfileManager.objects" + ) as profiles, + ): + profiles.get.side_effect = DjangoValidationError("bad uuid") + response = internal_views.index_update(request) + + assert response.status_code == 404, response.content diff --git a/backend/utils/models/org_path_discovery.py b/backend/utils/models/org_path_discovery.py index d2b3921231..4ce2d08dc2 100644 --- a/backend/utils/models/org_path_discovery.py +++ b/backend/utils/models/org_path_discovery.py @@ -9,6 +9,8 @@ import logging from collections import deque +from collections.abc import Mapping +from types import MappingProxyType from django.db import models @@ -42,22 +44,26 @@ # 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": ( - "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. - "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", -} +# Read-only: a wrong entry here is a cross-tenant leak, so the table is not +# something an importer should be able to reach in and change. +ORG_PATH_OVERRIDES: Mapping[str, str] = MappingProxyType( + { + "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. + "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: diff --git a/backend/utils/tests/test_organization_scoping_views.py b/backend/utils/tests/test_organization_scoping_views.py new file mode 100644 index 0000000000..2076ee3465 --- /dev/null +++ b/backend/utils/tests/test_organization_scoping_views.py @@ -0,0 +1,134 @@ +"""``filter_queryset_by_organization`` driven through the viewsets that use it. + +test_organization_scoping pins the helper against a synthetic request. That +proves the helper fails closed, not that the ~15 production call sites reach it +with the context it needs, nor that a viewset opting out of +``OrganizationFilterBackend`` with ``skip_org_filter = True`` has anything left +scoping it. + +These two viewsets are chosen because between them they cover both shapes the +call sites take: ``get_queryset()`` (five of the six internal viewsets) and a +``get_object()`` that deliberately bypasses ``get_queryset()`` for single-object +lookups (only FileExecutionInternalViewSet). + +``X-Organization-ID`` is what ``InternalAPIAuthMiddleware`` turns into +``request.organization_id``. The middleware warns and continues when the header +is absent, so "no organization_id attribute" is a reachable request state and is +tested as one. +""" + +import secrets + +import pytest +from account_v2.models import Organization +from django.test import TestCase +from rest_framework.test import APIRequestFactory +from workflow_manager.file_execution.internal_views import FileExecutionInternalViewSet +from workflow_manager.file_execution.models import WorkflowFileExecution +from workflow_manager.internal_views import WorkflowExecutionInternalViewSet +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models.execution import WorkflowExecution +from workflow_manager.workflow_v2.models.workflow import Workflow + +_ABSENT = object() + + +@pytest.mark.django_db +class InternalViewSetOrgScopingTest(TestCase): + def setUp(self) -> None: + self.factory = APIRequestFactory() + self.a = self._org_with_execution("a") + self.b = self._org_with_execution("b") + + def _org_with_execution(self, tag: str): + slug = f"org-{tag}-{secrets.token_hex(3)}" + org = Organization.objects.create( + name=slug, display_name=slug, organization_id=slug + ) + workflow = Workflow._base_manager.create( + workflow_name=f"wf-{slug}", organization=org + ) + execution = WorkflowExecution._base_manager.create( + workflow=workflow, status=ExecutionStatus.COMPLETED.value + ) + file_execution = WorkflowFileExecution._base_manager.create( + workflow_execution=execution, + file_name=f"{slug}.pdf", + status=ExecutionStatus.COMPLETED.value, + ) + return org, workflow, execution, file_execution + + def _get(self, path: str, organization_id=_ABSENT): + request = self.factory.get(path) + # Set exactly what the middleware sets. Absent is a distinct state from + # present-and-None, and production reads it with getattr(..., None). + if organization_id is not _ABSENT: + request.organization_id = organization_id + return request + + # --- get_queryset() path: five of the six internal viewsets ------------- + + def _list_executions(self, organization_id=_ABSENT): + view = WorkflowExecutionInternalViewSet.as_view({"get": "list"}) + return view(self._get("/internal/workflow-execution/", organization_id)) + + def _execution_ids(self, response) -> set[str]: + results = response.data + if isinstance(results, dict): + results = results.get("results", results.get("data", [])) + return {str(row["id"]) for row in results} + + def test_list_without_the_header_serves_no_rows(self): + """The header is optional at the middleware, so this is reachable.""" + response = self._list_executions() + assert response.status_code == 200, response.data + assert self._execution_ids(response) == set() + + def test_list_with_an_unresolvable_org_serves_no_rows(self): + response = self._list_executions("org-that-does-not-exist") + assert response.status_code == 200, response.data + assert self._execution_ids(response) == set() + + def test_list_serves_only_the_callers_own_org(self): + _, _, execution_a, _ = self.a + _, _, execution_b, _ = self.b + + response = self._list_executions(self.a[0].organization_id) + + assert response.status_code == 200, response.data + ids = self._execution_ids(response) + assert str(execution_a.id) in ids + assert str(execution_b.id) not in ids + + # --- get_object() path: bypasses get_queryset(), scopes independently --- + + def _retrieve_file_execution(self, pk, organization_id=_ABSENT): + view = FileExecutionInternalViewSet.as_view({"get": "retrieve"}) + request = self._get(f"/internal/file-execution/{pk}/", organization_id) + return view(request, id=str(pk)) + + def test_retrieve_without_the_header_is_refused(self): + _, _, _, file_execution = self.a + response = self._retrieve_file_execution(file_execution.id) + assert response.status_code == 404, response.data + + def test_retrieve_of_another_orgs_pk_is_refused(self): + """Org A's context against an org B id: the whole point of the helper.""" + _, _, _, file_execution_b = self.b + + response = self._retrieve_file_execution( + file_execution_b.id, self.a[0].organization_id + ) + + assert response.status_code == 404, response.data + + def test_retrieve_of_own_row_still_works(self): + """Without this, a viewset that refused everything would pass above.""" + _, _, _, file_execution_a = self.a + + response = self._retrieve_file_execution( + file_execution_a.id, self.a[0].organization_id + ) + + assert response.status_code == 200, response.data + assert str(response.data["id"]) == str(file_execution_a.id) diff --git a/backend/utils/uuid_validation.py b/backend/utils/uuid_validation.py new file mode 100644 index 0000000000..daf66cab98 --- /dev/null +++ b/backend/utils/uuid_validation.py @@ -0,0 +1,23 @@ +"""Parse a request-supplied id as a UUID, or fail as a 400. + +These columns are UUID primary/foreign keys, 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. + +Shared rather than copied per view: this is a security-relevant parse, and two +copies are two chances for one of them to widen its except tuple or lose a case. +""" + +import uuid +from typing import Any + +from rest_framework.exceptions import ValidationError + + +def validated_uuid(raw: Any, field_name: str) -> uuid.UUID: + """``raw`` as a UUID, or a 400 naming ``field_name``.""" + try: + return uuid.UUID(str(raw)) + except (ValueError, AttributeError, TypeError): + raise ValidationError(detail=f"'{field_name}' must be a valid UUID.") diff --git a/backend/workflow_manager/file_execution/internal_views.py b/backend/workflow_manager/file_execution/internal_views.py index 83688cd26e..7a3d67333e 100644 --- a/backend/workflow_manager/file_execution/internal_views.py +++ b/backend/workflow_manager/file_execution/internal_views.py @@ -29,9 +29,12 @@ class FileExecutionInternalViewSet(viewsets.ModelViewSet): serializer_class = WorkflowFileExecutionSerializer lookup_field = "id" - # OrganizationFilterBackend is off here; get_queryset() scopes instead, via - # filter_queryset_by_organization, which fails closed. X-Organization-ID is - # therefore required in practice: a worker that omits it gets zero rows. + # OrganizationFilterBackend is off here. Unlike the other viewsets with + # this pattern, get_object() below does not go through get_queryset() — it + # builds its own queryset to skip the 3-table JOIN — so the two scope + # independently, each via filter_queryset_by_organization, which fails + # closed. X-Organization-ID is therefore required in practice on both + # paths: a worker that omits it gets zero rows. skip_org_filter = True def get_object(self):