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/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..ea39cb1d28 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,22 +16,17 @@ ) 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__) class FileManagementViewSet(viewsets.ModelViewSet): - """FileManagement view. - - Handles GET,POST,PUT,PATCH and DELETE - """ + """FileManagement view.""" versioning_class = URLPathVersioning @@ -99,36 +92,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/notification_v2/internal_views.py b/backend/notification_v2/internal_views.py index 596945d01f..565f9427e1 100644 --- a/backend/notification_v2/internal_views.py +++ b/backend/notification_v2/internal_views.py @@ -33,7 +33,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_profile_manager_v2/models.py b/backend/prompt_studio/prompt_profile_manager_v2/models.py index 10a234f462..46d81ccbab 100644 --- a/backend/prompt_studio/prompt_profile_manager_v2/models.py +++ b/backend/prompt_studio/prompt_profile_manager_v2/models.py @@ -5,42 +5,43 @@ 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.user_context import UserContext +from utils.models.base_model import BaseModel +from utils.models.org_aware_manager import OrgAwareManager 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. - 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``. 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 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 acd2859b72..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,15 @@ 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( + 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, x2text_config_hash=x2text_config_hash, @@ -237,7 +300,44 @@ def extraction_status(request): extracted=extracted, error_message=error_message, ) - return JsonResponse({"success": success}) + # 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: + logger.error( + "extraction_status not recorded for document %s profile %s (%s)", + document_id, + profile_manager_id, + result.value, + ) + 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") @@ -321,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, @@ -431,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, ) @@ -440,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/migration_utils.py b/backend/prompt_studio/prompt_studio_core_v2/migration_utils.py index b9236d04f0..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 @@ -60,13 +52,34 @@ 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" - ) + # ProfileManager.objects is scoped through + # 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() + 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/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index 94f3850c62..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 @@ -69,6 +69,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 ( @@ -2574,7 +2575,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( + status_result = PromptStudioIndexHelper.mark_extraction_status( document_id=document_id, profile_manager=profile_manager, x2text_config_hash=x2text_config_hash, @@ -2582,7 +2583,7 @@ def dynamic_extractor( extracted=False, error_message=msg, ) - if not success: + 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." @@ -2592,13 +2593,16 @@ def dynamic_extractor( ) extracted_text = result.data.get("extracted_text", "") - success = 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 not success: + 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 8990e870eb..f889789a26 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/views.py @@ -9,9 +9,10 @@ import magic from account_v2.custom_exceptions import DuplicateData from celery import signature -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 from django.utils import timezone from file_management.constants import FileInformationKey as FileKey from file_management.exceptions import FileNotFound @@ -22,6 +23,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 @@ -30,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, @@ -385,13 +388,38 @@ 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 + # 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.") + 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 + # 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, + prompt_studio_tool=prompt_tool, ) - profile_manager = ProfileManager.objects.get(pk=request.data["default_profile"]) - 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. 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(update_fields=["is_default"]) return Response( status=status.HTTP_200_OK, @@ -1103,19 +1131,50 @@ 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. 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( + DocumentManager, pk=document_id, tool=custom_tool + ) try: # Delete indexed flags in redis index_managers = IndexManager.objects.filter(document_manager=document_id) + if not index_managers.exists(): + # 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( 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, @@ -1123,12 +1182,27 @@ 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, ) 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 15c76c5087..07ee681c4c 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,13 @@ class DocumentManager(BaseModel): """Model to store the document details.""" + # 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) 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..6d84ed832b 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,9 @@ class IndexManager(BaseModel): """Model to store the index details.""" + # See DocumentManager.objects for why scoping lives 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..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 @@ -1,7 +1,9 @@ import json import logging +from enum import Enum -from django.db import transaction +from django.core.exceptions import ImproperlyConfigured, ValidationError +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 @@ -12,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( @@ -75,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. @@ -90,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, 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: @@ -109,12 +127,13 @@ 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",) keeps the lock on index_manager rows only. + 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 @@ -147,17 +166,31 @@ def mark_extraction_status( f"Error: {error_message}" ) - return True - - except DocumentManager.DoesNotExist: - logger.error(f"Document with ID {document_id} does not exist.") - return False + return ExtractionStatusResult.OK + + 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.", + document_id, + ) + return ExtractionStatusResult.DOCUMENT_MISSING - except Exception as e: + except (DatabaseError, TypeError, ImproperlyConfigured): + # DatabaseError covers IntegrityError/OperationalError, TypeError a + # malformed extraction_status payload, ImproperlyConfigured a bad + # org path pin. Anything else propagates rather than being + # reported as a write failure it is not. 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 + return ExtractionStatusResult.WRITE_FAILED @staticmethod def check_extraction_status( 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..9420b87e9d 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,9 @@ class PromptStudioOutputManager(BaseModel): By default the tools will be added to private tool hub. """ + # See DocumentManager.objects for why scoping lives 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/output_manager_helper.py b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py index 699cacb749..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,6 +76,9 @@ def update_or_create_prompt_output( the instance. """ try: + # 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, 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..3ef714fa44 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py @@ -1,16 +1,18 @@ import logging +import uuid from typing import Any -from django.core.exceptions import ObjectDoesNotExist +from account_v2.models import Organization 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 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, @@ -29,6 +31,29 @@ logger = logging.getLogger(__name__) +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 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. + """ + 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 @@ -48,6 +73,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) @@ -78,9 +115,18 @@ 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_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( tool_id=tool_id, @@ -119,17 +165,39 @@ 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 - tool_studio_prompts = ToolStudioPrompt.objects.filter( - tool_id=tool_id - ).order_by("sequence_number") - except ObjectDoesNotExist: - raise ValidationError(detail=tool_not_found) + 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. + # 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 + # 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=organization, + ).order_by("sequence_number") # Invoke helper method to frame and fetch default response. result: dict[str, Any] = OutputManagerHelper.fetch_default_output_response( diff --git a/backend/prompt_studio/prompt_studio_v2/models.py b/backend/prompt_studio/prompt_studio_v2/models.py index faaf7b0313..8f9f9bc316 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,10 @@ class ToolStudioPrompt(BaseModel): It has Many to one relation with CustomTool for ToolStudio. """ + # See DocumentManager.objects for why scoping lives 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..5519aedec8 --- /dev/null +++ b/backend/prompt_studio/tests/test_cross_org_isolation.py @@ -0,0 +1,463 @@ +"""Organization isolation for the prompt-studio child models. + +``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. +""" + +import secrets +from unittest.mock import patch + +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 ( + 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.""" + + 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, + ) + self.adapter = self._adapter(slug) + self.profile = _make_profile( + f"profile-{slug}", + self.tool, + self.adapter, + 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: + # 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) + + # --- 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.""" + 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). + + 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): + """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 in delete_for_ide and make_profile_default ------ + + 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, + ) + profile = _make_profile( + f"sibling-profile-{secrets.token_hex(3)}", + tool, + self.a.adapter, + 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 _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}", + {"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) + + # 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" + assert response.status_code == 404, getattr(response, "data", response) + + # --- the ordering fix, driven through the view ------------------------ + + def _make_profile_default(self, tool, profile_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", + {"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 _make_profile( + "profile-a-second", + self.a.tool, + self.a.adapter, + 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 + + # --- 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): + """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") 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..b4cd843a8d --- /dev/null +++ b/backend/prompt_studio/tests/test_request_validation.py @@ -0,0 +1,308 @@ +"""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 reports success. + +Every case here was checked against the unguarded code first — strip the guard +and the case fails. +""" + +import json +import secrets +import uuid +from types import SimpleNamespace +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): + """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. + """ + 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 + + +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/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..8329095617 100644 --- a/backend/utils/models/org_aware_manager.py +++ b/backend/utils/models/org_aware_manager.py @@ -58,20 +58,50 @@ 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 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.", + self.model._meta.label, + type(exc).__name__, + exc, + ) return qs if org is None: - # No request context (Celery, management commands, shell) + 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 9a8011c716..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 @@ -19,6 +21,50 @@ _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 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. +# 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: """Get the cached FK path from a model to Organization. @@ -26,6 +72,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/organization_utils.py b/backend/utils/organization_utils.py index 15053684bf..63f0eca21b 100644 --- a/backend/utils/organization_utils.py +++ b/backend/utils/organization_utils.py @@ -72,7 +72,32 @@ 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. For every caller, this function is the only tenant boundary + 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 + 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 +105,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_org_path_discovery.py b/backend/utils/tests/test_org_path_discovery.py new file mode 100644 index 0000000000..88534b23cf --- /dev/null +++ b/backend/utils/tests/test_org_path_discovery.py @@ -0,0 +1,124 @@ +"""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, + _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. +# +# 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_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.""" + 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. + + 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 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: + field = model._meta.get_field(hop) + # 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 + + # 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 new file mode 100644 index 0000000000..f0b924c29a --- /dev/null +++ b/backend/utils/tests/test_organization_scoping.py @@ -0,0 +1,87 @@ +"""``filter_queryset_by_organization`` must fail closed. + +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 +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 + + +_ABSENT = object() + + +class _Request: + """Stands in for the request object the helper reads context off.""" + + 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 + + +@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 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 8fed360740..7a3d67333e 100644 --- a/backend/workflow_manager/file_execution/internal_views.py +++ b/backend/workflow_manager/file_execution/internal_views.py @@ -29,10 +29,12 @@ 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. 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): diff --git a/backend/workflow_manager/internal_views.py b/backend/workflow_manager/internal_views.py index 0aa0fb1f3a..8c919551cd 100644 --- a/backend/workflow_manager/internal_views.py +++ b/backend/workflow_manager/internal_views.py @@ -49,10 +49,9 @@ 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, 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 fefba8c21a..c999b7f9f5 100644 --- a/backend/workflow_manager/workflow_v2/views.py +++ b/backend/workflow_manager/workflow_v2/views.py @@ -403,10 +403,9 @@ 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, 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/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,