Stove-pipe RAG retrieval to prevent cross-tenant leakage (#40) (#41)

## Summary

- Closes [#40](#40)
- Aligns chat/RAG with the abc_worker stove-pipe pattern ([b13cec8](b13cec88f9)): immutable `ChatCompanyScope` per turn, conversation ownership validation, fail-closed Chroma filters
- Prefer ASGI/JWT identity over client email; never bind identity from `conversation_id` alone
- Close `ConversationDetailView` IDOR (prompts only for `request.user`)

## Changes

- New `services/chat_tenant_scope.py` with frozen `ChatCompanyScope` + ownership checks
- WebSocket consumers (`consumers.py` / `consumers_graph.py`) validate scope before `get_messages` / RAG
- `search_documents` requires a workspace (no more `filter: None` over the shared collection)
- Ingest writes `company_id` metadata (retrieval still keys on `workspace_id` for back-compat)
- Legacy `get_retriever` always applies a workspace filter

## Test plan

- [x] `manage.py test chat_backend.tests.test_chat_tenant_scope chat_backend.tests.test_consumers chat_backend.tests.test_services_rag chat_backend.tests.test_views_conversations`
- [ ] Manual: user A cannot stream RAG context from user B `conversation_id`
- [ ] Manual: RAG still returns own-company docs after deploy (existing vectors with `workspace_id` only)
- [ ] Follow-up: FE can send JWT `token`/`access` on WS payloads for stronger identity bindingReviewed-on: #41

llm_be/chat_backend/services/chat_tenant_scope.py

@@ -0,0 +1,182 @@
+"""Immutable per-turn company/workspace scope for chat + RAG.
+
+Mirrors the abc_worker ChatTenantScope stove-pipe: resolve identity once,
+validate conversation ownership, never derive tenant from an untrusted
+conversation_id alone.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+from django.contrib.auth.models import AnonymousUser
+from rest_framework_simplejwt.exceptions import TokenError
+from rest_framework_simplejwt.tokens import AccessToken
+
+from chat_backend.models import Conversation, CustomUser, DocumentWorkspace
+
+
+class ChatTenantScopeError(Exception):
+    """Raised when chat tenant resolution or ownership validation fails."""
+
+    def __init__(self, message: str, *, code: str = "tenant_scope_denied"):
+        super().__init__(message)
+        self.message = message
+        self.code = code
+
+
+@dataclass(frozen=True)
+class ChatCompanyScope:
+    """Frozen tenant identity for one websocket turn / RAG retrieval."""
+
+    user_id: int
+    company_id: int
+    workspace_id: int
+    conversation_id: Optional[int] = None
+
+
+def user_from_access_token(token: str) -> Optional[CustomUser]:
+    """Resolve an active user from a SimpleJWT access token string."""
+    if not token or not isinstance(token, str):
+        return None
+    try:
+        access = AccessToken(token)
+        user_id = access.get("user_id")
+        if not user_id:
+            return None
+        return CustomUser.objects.filter(id=user_id, is_active=True).first()
+    except TokenError:
+        return None
+
+
+def resolve_chat_user(
+    *,
+    email: Optional[str] = None,
+    token: Optional[str] = None,
+    authenticated_user=None,
+    conversation_id: Optional[int] = None,
+) -> Optional[CustomUser]:
+    """
+    Resolve the chat principal for a websocket turn.
+
+    Preference order:
+    1. Authenticated ASGI/session user
+    2. JWT access token (payload or query)
+    3. Client email (legacy FE path)
+
+    Does not fall back to conversation.user — that would bind identity to an
+    attacker-chosen conversation_id. ``conversation_id`` is accepted for API
+    compatibility but ignored for identity resolution.
+    """
+    del conversation_id
+    if authenticated_user is not None and getattr(
+        authenticated_user, "is_authenticated", False
+    ):
+        if isinstance(authenticated_user, CustomUser):
+            return authenticated_user
+        user = CustomUser.objects.filter(
+            id=authenticated_user.pk, is_active=True
+        ).first()
+        if user:
+            return user
+
+    token_user = user_from_access_token(token) if token else None
+    if token_user:
+        return token_user
+
+    if email:
+        return CustomUser.objects.filter(email__iexact=email, is_active=True).first()
+
+    return None
+
+
+def resolve_chat_company_scope(
+    user: CustomUser,
+    conversation_id: Optional[int] = None,
+) -> ChatCompanyScope:
+    """
+    Build an immutable company/workspace scope for ``user``.
+
+    When ``conversation_id`` is set, require ``conversation.user_id == user.id``
+    and that the conversation owner's company matches the user's company.
+    """
+    if user is None or not getattr(user, "id", None):
+        raise ChatTenantScopeError(
+            "Authenticated chat user is required.",
+            code="user_not_found",
+        )
+    if not getattr(user, "company_id", None):
+        raise ChatTenantScopeError(
+            "User is not attached to a company workspace.",
+            code="company_missing",
+        )
+
+    if conversation_id is not None:
+        conversation = (
+            Conversation.objects.select_related("user")
+            .filter(id=conversation_id, deleted=False)
+            .first()
+        )
+        if conversation is None:
+            raise ChatTenantScopeError(
+                "Conversation was not found.",
+                code="conversation_not_found",
+            )
+        if conversation.user_id != user.id:
+            raise ChatTenantScopeError(
+                "Conversation does not belong to the authenticated user.",
+                code="conversation_forbidden",
+            )
+        owner_company_id = getattr(conversation.user, "company_id", None)
+        if owner_company_id != user.company_id:
+            raise ChatTenantScopeError(
+                "Conversation company does not match the authenticated user.",
+                code="conversation_forbidden",
+            )
+
+    workspace = (
+        DocumentWorkspace.objects.filter(company_id=user.company_id)
+        .order_by("id")
+        .first()
+    )
+    if workspace is None:
+        raise ChatTenantScopeError(
+            "No document workspace exists for this company.",
+            code="workspace_missing",
+        )
+
+    return ChatCompanyScope(
+        user_id=user.id,
+        company_id=user.company_id,
+        workspace_id=workspace.id,
+        conversation_id=conversation_id,
+    )
+
+
+def create_conversation_for_user(user: CustomUser, title: str) -> int:
+    """Create a conversation owned by ``user`` and return its id."""
+    conversation = Conversation.objects.create(title=title, user=user)
+    return conversation.id
+
+
+def get_workspace_for_scope(scope: ChatCompanyScope) -> DocumentWorkspace:
+    """Load workspace rows only when they match the frozen scope keys."""
+    try:
+        return DocumentWorkspace.objects.get(
+            id=scope.workspace_id, company_id=scope.company_id
+        )
+    except DocumentWorkspace.DoesNotExist as exc:
+        raise ChatTenantScopeError(
+            "Scoped document workspace was not found.",
+            code="workspace_missing",
+        ) from exc
+
+
+def asgi_user_or_none(scope_user):
+    """Return an authenticated user from Channels scope, else None."""
+    if scope_user is None or isinstance(scope_user, AnonymousUser):
+        return None
+    if getattr(scope_user, "is_authenticated", False):
+        return scope_user
+    return None

llm_be/chat_backend/services/rag_services.py

@@ -118,6 +118,7 @@ class RAGService(BaseService):
                     metadata={
                         "source": doc.file.name,
                         "workspace_id": doc.workspace_id,
+                        "company_id": doc.workspace.company_id,
                         "document_id": doc.id,
                     },
                 )
@@ -132,10 +133,11 @@ class RAGService(BaseService):
     def ingest_documents(self, workspace: DocumentWorkspace | None = None) -> None:
         """Ingest documents from a workspace into the vector store."""
         print(f"Getting the Document via the workspace: {workspace}")
+        qs = Document.objects.select_related("workspace")
         if workspace:
-            documents = [doc for doc in Document.objects.filter(workspace=workspace)]
+            documents = list(qs.filter(workspace=workspace))
         else:
-            documents = [doc for doc in Document.objects.all()]
+            documents = list(qs.all())
 
         print(f"Processing the documents : {documents}")
         self._prepare_documents(documents)
@@ -192,9 +194,17 @@ class RAGService(BaseService):
                     tmp_created = self._materialize_file_field(file_ref)
                     file_path = tmp_created
 
+                company_id = None
+                if ws_id is not None:
+                    company_id = (
+                        DocumentWorkspace.objects.filter(id=ws_id)
+                        .values_list("company_id", flat=True)
+                        .first()
+                    )
                 metadata = {
                     "source": original_name,
                     "workspace_id": ws_id,
+                    "company_id": company_id,
                     "original_filename": original_name,
                     "file_path": original_name,
                 }
@@ -220,6 +230,17 @@ class RAGService(BaseService):
         self.vector_store.persist()
         return results
 
+    def _workspace_filter(self, workspace: DocumentWorkspace) -> Dict[str, Any]:
+        """Build a fail-closed Chroma metadata filter for one workspace.
+
+        ``company_id`` is written on ingest for defense-in-depth / future dual
+        filters, but retrieval keys on ``workspace_id`` so older vectors without
+        ``company_id`` metadata still match after deploy.
+        """
+        if workspace is None or getattr(workspace, "id", None) is None:
+            raise ValueError("workspace is required for RAG retrieval")
+        return {"workspace_id": workspace.id}
+
 
 class SyncRAGService(RAGService):
     """Synchronous RAG service implementation."""
@@ -265,9 +286,12 @@ class SyncRAGService(RAGService):
     def _retriever_with_history(self, input_dict: Dict[str, Any]) -> str:
         """Retrieve documents considering conversation history."""
         query = input_dict["query"]
-        conversation = input_dict["conversation"]
+        workspace = input_dict.get("workspace")
+        if workspace is None:
+            conversation = input_dict.get("conversation")
+            workspace = getattr(conversation, "workspace", None)
 
-        relevant_docs = self.search_documents(query, conversation.workspace)
+        relevant_docs = self.search_documents(query, workspace)
         if not relevant_docs:
             print("didn't find any relevant docs")
             return relevant_docs
@@ -277,11 +301,9 @@ class SyncRAGService(RAGService):
     def search_documents(
         self, query: str, workspace: Optional[DocumentWorkspace] = None, k: int = 4
     ) -> List[Document]:
-        """Search relevant documents from the vector store."""
-        filter_dict = {}
-        if workspace:
-            filter_dict["workspace_id"] = workspace.id
-        search_kwargs = {"k": k, "filter": filter_dict if filter_dict else None}
+        """Search relevant documents from the vector store (workspace required)."""
+        filter_dict = self._workspace_filter(workspace)
+        search_kwargs = {"k": k, "filter": filter_dict}
         print(f"search_kwargs: {search_kwargs}")
         retriever = self.vector_store.as_retriever(
             search_type="similarity",
@@ -358,17 +380,14 @@ class AsyncRAGService(RAGService):
     async def search_documents(
         self, query: str, workspace: Optional[DocumentWorkspace] = None, k: int = 4
     ) -> List[Document]:
-        """Search relevant documents from the vector store."""
-        filter_dict = {}
+        """Search relevant documents from the vector store (workspace required)."""
+        filter_dict = self._workspace_filter(workspace)
         print(f"Do we have a workspace: {workspace}")
-        if workspace:
-            filter_dict["workspace_id"] = workspace.id
-        search_kwargs = {"k": k, "filter": filter_dict if filter_dict else None}
-        print(f"search_kwargs: {search_kwargs}")
+        print(f"search_kwargs: {{'k': {k}, 'filter': {filter_dict}}}")
 
         retriever = self.vector_store.as_retriever(
             search_type="mmr",
-            search_kwargs={"k": k, "filter": filter_dict if filter_dict else None},
+            search_kwargs={"k": k, "filter": filter_dict},
         )
         return await retriever.aget_relevant_documents(query)
 
@@ -380,6 +399,8 @@ class AsyncRAGService(RAGService):
         **kwargs,
     ) -> AsyncGenerator[str, None]:
         """Generate response with streaming support."""
+        if workspace is None:
+            raise ValueError("workspace is required for RAG generation")
         chain_input = {
             "query": query,
             "conversation": conversation,