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

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,

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()

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"),
     ]
 )

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,
-            )

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):

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):

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()
 
 

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 @@
             )
         )
 
+        # ``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 @@
         )
 
     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,21 +284,60 @@
         )
 
     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,
             enable_highlight=enable_highlight,
             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 @@
     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 @@
     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 @@
         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,