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

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,21 +16,19 @@
 )
 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
+    Handles GET, POST, PUT and PATCH
     """
 
     versioning_class = URLPathVersioning
@@ -99,36 +95,3 @@ def upload(self, request: HttpRequest) -> Response:
             logger.info(f"Uploading file: {file_name}" if file_name else "Uploading file")
             FileManagerHelper.upload_file(file_system, path, uploaded_file, file_name)
         return Response({"message": "Files are uploaded successfully!"})
-
-    @action(detail=True, methods=["get"])
-    def delete(self, request: HttpRequest) -> Response:
-        serializer = FileInfoIdeSerializer(data=request.GET)
-        serializer.is_valid(raise_exception=True)
-        document_id: str = serializer.validated_data.get("document_id")
-        document: DocumentManager = DocumentManager.objects.get(pk=document_id)
-        file_name: str = document.document_name
-        tool_id: str = serializer.validated_data.get("tool_id")
-        file_path = FileManagerHelper.handle_sub_directory_for_tenants(
-            UserSessionUtils.get_organization_id(request),
-            is_create=False,
-            user_id=request.user.user_id,
-            tool_id=tool_id,
-        )
-        path = file_path
-        file_system = LocalStorageFS(settings={"path": path})
-        try:
-            # Delete the document record
-            document.delete()
-
-            # Delete the file
-            FileManagerHelper.delete_file(file_system, path, file_name)
-            return Response(
-                {"data": "File deleted succesfully."},
-                status=status.HTTP_200_OK,
-            )
-        except Exception as exc:
-            logger.error(f"Exception thrown from file deletion, error {exc}")
-            return Response(
-                {"data": "File deletion failed."},
-                status=status.HTTP_400_BAD_REQUEST,
-            )

backend/prompt_studio/prompt_profile_manager_v2/models.py

@@ -5,14 +5,15 @@
 from django.db import models
 from django.db.models import Q
 from tenant_account_v2.organization_member_service import OrganizationMemberService
-from utils.models.base_model import BaseModel, BaseModelManager
+from utils.models.base_model import BaseModel
+from utils.models.org_aware_manager import OrgAwareManager
 from utils.user_context import UserContext
 
 from prompt_studio.prompt_studio_core_v2.exceptions import DefaultProfileError
 from prompt_studio.prompt_studio_core_v2.models import CustomTool
 
 
-class ProfileManagerModelManager(BaseModelManager):
+class ProfileManagerModelManager(OrgAwareManager):
     def for_user(self, user):
         """Read visibility: profile's own share fields OR parent CustomTool sharing.
 

backend/prompt_studio/prompt_studio_core_v2/migration_utils.py

@@ -60,9 +60,11 @@ def migrate_tool_to_adapter_based(
 
                 # Re-fetch the summarize profile with lock within transaction
                 try:
-                    summarize_profile = ProfileManager.objects.select_for_update().get(
-                        prompt_studio_tool=tool_instance, is_summarize_llm=True
-                    )
+                    # of=("self",): the org-scoped manager joins through
+                    # AdapterInstance, which would otherwise be locked too.
+                    summarize_profile = ProfileManager.objects.select_for_update(
+                        of=("self",)
+                    ).get(prompt_studio_tool=tool_instance, is_summarize_llm=True)
                 except ObjectDoesNotExist:
                     logger.info(
                         f"No summarize profile found for tool {tool_instance.tool_id}, skipping migration"

backend/prompt_studio/prompt_studio_core_v2/views.py

@@ -11,9 +11,10 @@
 from api_v2.models import APIDeployment
 from celery import signature
 from celery.result import AsyncResult
-from django.db import IntegrityError
+from django.db import IntegrityError, transaction
 from django.db.models import Count, OuterRef, QuerySet, Subquery
 from django.http import HttpRequest, HttpResponse
+from django.shortcuts import get_object_or_404
 from django.utils import timezone
 from file_management.constants import FileInformationKey as FileKey
 from file_management.exceptions import FileNotFound
@@ -445,13 +446,24 @@ def make_profile_default(self, request: HttpRequest, pk: Any = None) -> Response
             self.get_object()
         )  # Assuming you have a get_object method in your viewset
 
-        ProfileManager.objects.filter(prompt_studio_tool=prompt_tool).update(
-            is_default=False
+        # Resolve the target before clearing anything: the id comes straight
+        # from the request body, and clearing first would leave the tool with no
+        # default at all when it does not match. Scoped to the same tool the
+        # caller already passed authz on, so another tool's id is a 404.
+        profile_manager = get_object_or_404(
+            ProfileManager,
+            pk=request.data["default_profile"],
+            prompt_studio_tool=prompt_tool,
         )
 
-        profile_manager = 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.
+        with transaction.atomic():
+            ProfileManager.objects.filter(prompt_studio_tool=prompt_tool).update(
+                is_default=False
+            )
+            profile_manager.is_default = True
+            profile_manager.save()
 
         return Response(
             status=status.HTTP_200_OK,
@@ -1186,7 +1198,13 @@ def delete_for_ide(self, request: HttpRequest, pk: uuid) -> Response:
         document_id: str = serializer.validated_data.get(ToolStudioPromptKeys.DOCUMENT_ID)
         org_id = UserSessionUtils.get_organization_id(request)
         user_id = custom_tool.created_by.user_id
-        document: DocumentManager = DocumentManager.objects.get(pk=document_id)
+        # Scope to the tool the caller already passed authz on — tighter than
+        # org scope, and this action never runs filter_queryset().
+        # get_object_or_404 keeps a non-matching id a 404 rather than an
+        # unhandled DoesNotExist, which the DRF handler turns into a 500.
+        document: DocumentManager = get_object_or_404(
+            DocumentManager, pk=document_id, tool=custom_tool
+        )
 
         try:
             # Delete indexed flags in redis

backend/prompt_studio/prompt_studio_document_manager_v2/models.py

@@ -3,13 +3,18 @@
 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
 
 
 class DocumentManager(BaseModel):
     """Model to store the document details."""
 
+    # Org scoping lives here because custom @action methods never call
+    # filter_queryset(), so OrganizationFilterBackend does not run on them.
+    objects = OrgAwareManager()
+
     document_id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
 
     document_name = models.CharField(

backend/prompt_studio/prompt_studio_index_manager_v2/models.py

@@ -7,6 +7,7 @@
 from django.db.models.signals import pre_delete
 from django.dispatch import receiver
 from utils.models.base_model import BaseModel
+from utils.models.org_aware_manager import OrgAwareManager
 from utils.user_context import UserContext
 
 from prompt_studio.prompt_profile_manager_v2.models import ProfileManager
@@ -21,6 +22,10 @@
 class IndexManager(BaseModel):
     """Model to store the index details."""
 
+    # See DocumentManager.objects — custom @action methods bypass the
+    # OrganizationFilterBackend, so scoping has to be at the manager.
+    objects = OrgAwareManager()
+
     index_manager_id = models.UUIDField(
         primary_key=True, default=uuid.uuid4, editable=False
     )

backend/prompt_studio/prompt_studio_index_manager_v2/prompt_studio_index_helper.py

@@ -109,12 +109,15 @@ def mark_extraction_status(
 
                 # Lock the row (or create an empty one) so concurrent callers
                 # merge into the same dict rather than clobbering each other.
-                index_manager, created = (
-                    IndexManager.objects.select_for_update().get_or_create(
-                        document_manager=document,
-                        profile_manager=profile_manager,
-                        defaults={"extraction_status": {}},
-                    )
+                # of=("self",) because the org-scoped manager joins through
+                # DocumentManager and CustomTool; without it Postgres locks
+                # rows in those tables too.
+                index_manager, created = IndexManager.objects.select_for_update(
+                    of=("self",)
+                ).get_or_create(
+                    document_manager=document,
+                    profile_manager=profile_manager,
+                    defaults={"extraction_status": {}},
                 )
 
                 # Merge in place — update_or_create(defaults=...) would replace

backend/prompt_studio/prompt_studio_output_manager_v2/models.py

@@ -3,6 +3,7 @@
 from account_v2.models import User
 from django.db import models
 from utils.models.base_model import BaseModel
+from utils.models.org_aware_manager import OrgAwareManager
 
 from prompt_studio.prompt_profile_manager_v2.models import ProfileManager
 from prompt_studio.prompt_studio_core_v2.models import CustomTool
@@ -16,6 +17,10 @@ class PromptStudioOutputManager(BaseModel):
     By default the tools will be added to private tool hub.
     """
 
+    # See DocumentManager.objects — custom @action methods bypass the
+    # OrganizationFilterBackend, so scoping has to be at the manager.
+    objects = OrgAwareManager()
+
     prompt_output_id = models.UUIDField(
         primary_key=True, default=uuid.uuid4, editable=False
     )

backend/prompt_studio/prompt_studio_output_manager_v2/views.py

@@ -1,7 +1,6 @@
 import logging
 from typing import Any
 
-from django.core.exceptions import ObjectDoesNotExist
 from django.db.models import QuerySet
 from django.http import HttpRequest
 from rest_framework import status, viewsets
@@ -119,17 +118,22 @@ def get_output_for_tool_default(self, request: HttpRequest) -> Response:
         tool_id = request.GET.get("tool_id")
         document_manager_id = request.GET.get("document_manager")
         tool_validation_message = PromptOutputManagerErrorMessage.TOOL_VALIDATION
-        tool_not_found = PromptOutputManagerErrorMessage.TOOL_NOT_FOUND
         if not tool_id:
             raise ValidationError(detail=tool_validation_message)
 
-        try:
-            # Fetch ToolStudioPrompt records based on tool_id
-            tool_studio_prompts = ToolStudioPrompt.objects.filter(
-                tool_id=tool_id
-            ).order_by("sequence_number")
-        except ObjectDoesNotExist:
-            raise ValidationError(detail=tool_not_found)
+        # Fetch ToolStudioPrompt records based on tool_id.
+        # Custom actions skip filter_queryset(), so OrganizationFilterBackend
+        # never runs — scope explicitly to prevent cross-tenant reads.
+        #
+        # No exception handling here: filter() does not raise for a missing or
+        # out-of-org tool, it returns empty. Empty is also the correct result
+        # for a tool that simply has no prompts yet, which is the normal state
+        # of a newly created project — so this stays a 200 with an empty body
+        # rather than a validation error.
+        tool_studio_prompts = ToolStudioPrompt.objects.filter(
+            tool_id=tool_id,
+            tool_id__organization=UserContext.get_organization(),
+        ).order_by("sequence_number")
 
         # Invoke helper method to frame and fetch default response.
         result: dict[str, Any] = OutputManagerHelper.fetch_default_output_response(

backend/prompt_studio/prompt_studio_v2/models.py

@@ -4,6 +4,7 @@
 from django.db import models
 from django.utils import timezone
 from utils.models.base_model import BaseModel
+from utils.models.org_aware_manager import OrgAwareManager
 
 from prompt_studio.prompt_profile_manager_v2.models import ProfileManager
 from prompt_studio.prompt_studio_core_v2.models import CustomTool
@@ -15,6 +16,11 @@ class ToolStudioPrompt(BaseModel):
     It has Many to one relation with CustomTool for ToolStudio.
     """
 
+    # See DocumentManager.objects — custom @action methods bypass the
+    # OrganizationFilterBackend, so scoping has to be at the manager.
+    # tool_id is nullable, so prompts orphaned from their tool are excluded.
+    objects = OrgAwareManager()
+
     class EnforceType(models.TextChoices):
         TEXT = "text", "Response sent as Text"
         NUMBER = "number", "Response sent as number"