Stove-pipe RAG retrieval behind ChatCompanyScope (#40).
CI / test (pull_request) Successful in 10s
Unit Tests / test (pull_request) Successful in 9s

Validate conversation ownership before workspace resolution, fail closed
when Chroma filters are missing, and close ConversationDetailView IDOR so
tenants cannot pull another company's document chunks.
This commit is contained in:
2026-08-01 14:31:10 -05:00
parent eedc842b08
commit 7fae0175cf
9 changed files with 507 additions and 90 deletions
+95 -36
View File
@@ -23,6 +23,14 @@ from .models import Conversation, Prompt, PromptMetric, DocumentWorkspace, Docum
from .serializers import PromptSerializer from .serializers import PromptSerializer
from .services.llm_service import AsyncLLMService from .services.llm_service import AsyncLLMService
from .services.rag_services import AsyncRAGService from .services.rag_services import AsyncRAGService
from .services.chat_tenant_scope import (
ChatTenantScopeError,
asgi_user_or_none,
create_conversation_for_user,
get_workspace_for_scope,
resolve_chat_company_scope,
resolve_chat_user as resolve_chat_user_sync,
)
from .services.title_generator import title_generator from .services.title_generator import title_generator
from .services.moderation_classifier import moderation_classifier, ModerationLabel from .services.moderation_classifier import moderation_classifier, ModerationLabel
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
@@ -47,33 +55,27 @@ CHANNEL_NAME: str = "llm_messages"
MODEL_NAME: str = "llama3.2" MODEL_NAME: str = "llama3.2"
PROMPT_CLASSIFIER = PromptClassifier() PROMPT_CLASSIFIER = PromptClassifier()
@database_sync_to_async
def create_conversation(prompt, email, title):
# return the conversation id
conversation = Conversation.objects.create(title=title)
conversation.save()
user = CustomUser.objects.get(email=email) @database_sync_to_async
conversation.user_id = user.id def create_conversation(prompt, email, title, user=None):
conversation.save() """Create a conversation for ``user`` (preferred) or legacy ``email``."""
return conversation.id if user is None:
user = CustomUser.objects.get(email=email)
return create_conversation_for_user(user, title)
@database_sync_to_async @database_sync_to_async
def resolve_chat_user(email=None, conversation_id=None): def resolve_chat_user(
if email: email=None, conversation_id=None, token=None, authenticated_user=None
user = CustomUser.objects.filter(email__iexact=email).first() ):
if user: # conversation_id intentionally unused for identity — ownership is checked
return user # via resolve_chat_company_scope after the principal is known.
if conversation_id: return resolve_chat_user_sync(
conversation = ( email=email,
Conversation.objects.select_related("user") token=token,
.filter(id=conversation_id) authenticated_user=authenticated_user,
.first() conversation_id=conversation_id,
) )
if conversation and conversation.user_id:
return conversation.user
return None
@database_sync_to_async @database_sync_to_async
@@ -89,9 +91,20 @@ def enforce_feature_gate(user, feature):
@database_sync_to_async @database_sync_to_async
def get_workspace(conversation_id): def get_workspace(conversation_id, user=None):
conversation = Conversation.objects.get(id=conversation_id) """Resolve workspace only after conversation ownership is validated."""
return DocumentWorkspace.objects.get(company=conversation.user.company) if user is None:
raise ChatTenantScopeError(
"Authenticated chat user is required.",
code="user_not_found",
)
scope = resolve_chat_company_scope(user, conversation_id)
return get_workspace_for_scope(scope)
@database_sync_to_async
def resolve_tenant_scope(user, conversation_id=None):
return resolve_chat_company_scope(user, conversation_id)
@database_sync_to_async @database_sync_to_async
@@ -223,12 +236,17 @@ def finish_prompt_metric(prompt_metric, response_length, tokens_in=None, tokens_
@database_sync_to_async @database_sync_to_async
def get_retriever(conversation_id): def get_retriever(conversation_id, user=None):
"""Legacy helper — always applies a workspace metadata filter (fail closed)."""
if user is None:
raise ChatTenantScopeError(
"Authenticated chat user is required.",
code="user_not_found",
)
logger.info(f"getting workspace from conversation: {conversation_id}") logger.info(f"getting workspace from conversation: {conversation_id}")
conversation = Conversation.objects.get(id=conversation_id) scope = resolve_chat_company_scope(user, conversation_id)
logger.info(f"Got conversation: {conversation}") workspace = get_workspace_for_scope(scope)
workspace = DocumentWorkspace.objects.get(company=conversation.user.company) logger.info(f"Got workspace: {workspace.id} company={scope.company_id}")
logger.info(f"Got workspace: {conversation}")
persist_directory = getattr( persist_directory = getattr(
django_settings, "CHROMA_PERSIST_DIRECTORY", "./chroma_db/" django_settings, "CHROMA_PERSIST_DIRECTORY", "./chroma_db/"
) )
@@ -236,7 +254,10 @@ def get_retriever(conversation_id):
persist_directory=persist_directory, persist_directory=persist_directory,
embedding=OllamaEmbeddings(**ollama_embeddings_kwargs()), embedding=OllamaEmbeddings(**ollama_embeddings_kwargs()),
) )
return vectorstore.as_retriever() return vectorstore.as_retriever(
search_type="similarity",
search_kwargs={"k": 4, "filter": {"workspace_id": workspace.id}},
)
async def get_conversation_file_async(conversation_id): async def get_conversation_file_async(conversation_id):
try: try:
@@ -289,6 +310,7 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
message = normalize_user_message(data.get("message", None)) message = normalize_user_message(data.get("message", None))
conversation_id = data.get("conversation_id", None) conversation_id = data.get("conversation_id", None)
email = data.get("email", None) email = data.get("email", None)
token = data.get("token") or data.get("access")
file = data.get("file", None) file = data.get("file", None)
file_type = data.get("fileType", "") file_type = data.get("fileType", "")
model = data.get("modelName", "Turbo") model = data.get("modelName", "Turbo")
@@ -306,7 +328,10 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
return return
chat_user = await resolve_chat_user( chat_user = await resolve_chat_user(
email=email, conversation_id=conversation_id email=email,
conversation_id=conversation_id,
token=token,
authenticated_user=asgi_user_or_none(self.scope.get("user")),
) )
if chat_user is None: if chat_user is None:
await self.send_json_message( await self.send_json_message(
@@ -339,7 +364,37 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
# we need to create a new conversation # we need to create a new conversation
# we will generate a name for it too # we will generate a name for it too
title = await title_generator.generate_async(message) title = await title_generator.generate_async(message)
conversation_id = await create_conversation(message, email, title) conversation_id = await create_conversation(
message, email, title, user=chat_user
)
try:
tenant_scope = await resolve_tenant_scope(chat_user, conversation_id)
except ChatTenantScopeError as exc:
logger.warning(
"websocket tenant validation failed conversation_id=%s user_id=%s code=%s",
conversation_id,
chat_user.id,
exc.code,
)
await self.send_json_message(
json.dumps(
{
"type": "error",
"code": exc.code,
"content": exc.message,
}
)
)
return
logger.info(
"chat_scope_validated conversation_id=%s user_id=%s company_id=%s workspace_id=%s",
tenant_scope.conversation_id,
tenant_scope.user_id,
tenant_scope.company_id,
tenant_scope.workspace_id,
)
if conversation_id: if conversation_id:
decoded_file = None decoded_file = None
@@ -443,8 +498,12 @@ class ChatConsumerAgain(AsyncWebsocketConsumer):
if prompt_type == PromptType.RAG: if prompt_type == PromptType.RAG:
service = AsyncRAGService() service = AsyncRAGService()
workspace = await get_workspace(conversation_id) workspace = await get_workspace(
return service.generate_response(messages, prompt_instance.message, workspace) conversation_id, user=chat_user
)
return service.generate_response(
messages, prompt_instance.message, workspace
)
elif prompt_type == PromptType.DATA_ANALYSIS: elif prompt_type == PromptType.DATA_ANALYSIS:
service = AsyncDataAnalysisService() service = AsyncDataAnalysisService()
+72 -26
View File
@@ -18,6 +18,14 @@ from .models import Conversation, Prompt, PromptMetric, DocumentWorkspace, Custo
from .serializers import PromptSerializer from .serializers import PromptSerializer
from .services.llm_service import AsyncLLMService from .services.llm_service import AsyncLLMService
from .services.rag_services import AsyncRAGService from .services.rag_services import AsyncRAGService
from .services.chat_tenant_scope import (
ChatTenantScopeError,
asgi_user_or_none,
create_conversation_for_user,
get_workspace_for_scope,
resolve_chat_company_scope,
resolve_chat_user as resolve_chat_user_sync,
)
from .services.title_generator import title_generator from .services.title_generator import title_generator
from .services.moderation_classifier import moderation_classifier, ModerationLabel from .services.moderation_classifier import moderation_classifier, ModerationLabel
from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType from .services.prompt_classifier.prompt_classifier import PromptClassifier, PromptType
@@ -41,29 +49,22 @@ PROMPT_CLASSIFIER = PromptClassifier()
# --- Database Helpers (Reused) --- # --- Database Helpers (Reused) ---
@database_sync_to_async @database_sync_to_async
def create_conversation(prompt, email, title): def create_conversation(prompt, email, title, user=None):
conversation = Conversation.objects.create(title=title) if user is None:
user = CustomUser.objects.get(email=email) user = CustomUser.objects.get(email=email)
conversation.user_id = user.id return create_conversation_for_user(user, title)
conversation.save()
return conversation.id
@database_sync_to_async @database_sync_to_async
def resolve_chat_user(email=None, conversation_id=None): def resolve_chat_user(
if email: email=None, conversation_id=None, token=None, authenticated_user=None
user = CustomUser.objects.filter(email__iexact=email).first() ):
if user: return resolve_chat_user_sync(
return user email=email,
if conversation_id: token=token,
conversation = ( authenticated_user=authenticated_user,
Conversation.objects.select_related("user") conversation_id=conversation_id,
.filter(id=conversation_id) )
.first()
)
if conversation and conversation.user_id:
return conversation.user
return None
@database_sync_to_async @database_sync_to_async
@@ -78,9 +79,19 @@ def enforce_feature_gate(user, feature):
assert_feature_allowed(user, feature) assert_feature_allowed(user, feature)
@database_sync_to_async @database_sync_to_async
def get_workspace(conversation_id): def get_workspace(conversation_id, user=None):
conversation = Conversation.objects.get(id=conversation_id) if user is None:
return DocumentWorkspace.objects.get(company=conversation.user.company) raise ChatTenantScopeError(
"Authenticated chat user is required.",
code="user_not_found",
)
scope = resolve_chat_company_scope(user, conversation_id)
return get_workspace_for_scope(scope)
@database_sync_to_async
def resolve_tenant_scope(user, conversation_id=None):
return resolve_chat_company_scope(user, conversation_id)
@database_sync_to_async @database_sync_to_async
def get_messages(conversation_id, prompt, file_string: str = None, file_type: str = ""): def get_messages(conversation_id, prompt, file_string: str = None, file_type: str = ""):
@@ -284,7 +295,8 @@ async def generation_node(state: ChatState) -> ChatState:
if prompt_type == PromptType.RAG: if prompt_type == PromptType.RAG:
service = AsyncRAGService() service = AsyncRAGService()
workspace = await get_workspace(conversation_id) chat_user = state.get("chat_user")
workspace = await get_workspace(conversation_id, user=chat_user)
generator = service.generate_response(messages, prompt_instance.message, workspace) generator = service.generate_response(messages, prompt_instance.message, workspace)
return {"response_generator": generator} return {"response_generator": generator}
@@ -349,6 +361,7 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
message = normalize_user_message(data.get("message", None)) message = normalize_user_message(data.get("message", None))
conversation_id = data.get("conversation_id", None) conversation_id = data.get("conversation_id", None)
email = data.get("email", None) email = data.get("email", None)
token = data.get("token") or data.get("access")
file = data.get("file", None) file = data.get("file", None)
file_type = data.get("fileType", "") file_type = data.get("fileType", "")
@@ -365,7 +378,10 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
return return
chat_user = await resolve_chat_user( chat_user = await resolve_chat_user(
email=email, conversation_id=conversation_id email=email,
conversation_id=conversation_id,
token=token,
authenticated_user=asgi_user_or_none(self.scope.get("user")),
) )
if chat_user is None: if chat_user is None:
await self.send_json_message( await self.send_json_message(
@@ -396,7 +412,37 @@ class ChatConsumerGraph(AsyncWebsocketConsumer):
if not conversation_id: if not conversation_id:
title = await title_generator.generate_async(message) title = await title_generator.generate_async(message)
conversation_id = await create_conversation(message, email, title) conversation_id = await create_conversation(
message, email, title, user=chat_user
)
try:
tenant_scope = await resolve_tenant_scope(chat_user, conversation_id)
except ChatTenantScopeError as exc:
logger.warning(
"websocket tenant validation failed conversation_id=%s user_id=%s code=%s",
conversation_id,
chat_user.id,
exc.code,
)
await self.send_json_message(
json.dumps(
{
"type": "error",
"code": exc.code,
"content": exc.message,
}
)
)
return
logger.info(
"chat_scope_validated conversation_id=%s user_id=%s company_id=%s workspace_id=%s",
tenant_scope.conversation_id,
tenant_scope.user_id,
tenant_scope.company_id,
tenant_scope.workspace_id,
)
if conversation_id: if conversation_id:
print("Conversation ID: ", conversation_id) print("Conversation ID: ", conversation_id)
@@ -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
+37 -16
View File
@@ -118,6 +118,7 @@ class RAGService(BaseService):
metadata={ metadata={
"source": doc.file.name, "source": doc.file.name,
"workspace_id": doc.workspace_id, "workspace_id": doc.workspace_id,
"company_id": doc.workspace.company_id,
"document_id": doc.id, "document_id": doc.id,
}, },
) )
@@ -132,10 +133,11 @@ class RAGService(BaseService):
def ingest_documents(self, workspace: DocumentWorkspace | None = None) -> None: def ingest_documents(self, workspace: DocumentWorkspace | None = None) -> None:
"""Ingest documents from a workspace into the vector store.""" """Ingest documents from a workspace into the vector store."""
print(f"Getting the Document via the workspace: {workspace}") print(f"Getting the Document via the workspace: {workspace}")
qs = Document.objects.select_related("workspace")
if workspace: if workspace:
documents = [doc for doc in Document.objects.filter(workspace=workspace)] documents = list(qs.filter(workspace=workspace))
else: else:
documents = [doc for doc in Document.objects.all()] documents = list(qs.all())
print(f"Processing the documents : {documents}") print(f"Processing the documents : {documents}")
self._prepare_documents(documents) self._prepare_documents(documents)
@@ -192,9 +194,17 @@ class RAGService(BaseService):
tmp_created = self._materialize_file_field(file_ref) tmp_created = self._materialize_file_field(file_ref)
file_path = tmp_created 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 = { metadata = {
"source": original_name, "source": original_name,
"workspace_id": ws_id, "workspace_id": ws_id,
"company_id": company_id,
"original_filename": original_name, "original_filename": original_name,
"file_path": original_name, "file_path": original_name,
} }
@@ -220,6 +230,17 @@ class RAGService(BaseService):
self.vector_store.persist() self.vector_store.persist()
return results 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): class SyncRAGService(RAGService):
"""Synchronous RAG service implementation.""" """Synchronous RAG service implementation."""
@@ -265,9 +286,12 @@ class SyncRAGService(RAGService):
def _retriever_with_history(self, input_dict: Dict[str, Any]) -> str: def _retriever_with_history(self, input_dict: Dict[str, Any]) -> str:
"""Retrieve documents considering conversation history.""" """Retrieve documents considering conversation history."""
query = input_dict["query"] 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: if not relevant_docs:
print("didn't find any relevant docs") print("didn't find any relevant docs")
return relevant_docs return relevant_docs
@@ -277,11 +301,9 @@ class SyncRAGService(RAGService):
def search_documents( def search_documents(
self, query: str, workspace: Optional[DocumentWorkspace] = None, k: int = 4 self, query: str, workspace: Optional[DocumentWorkspace] = None, k: int = 4
) -> List[Document]: ) -> List[Document]:
"""Search relevant documents from the vector store.""" """Search relevant documents from the vector store (workspace required)."""
filter_dict = {} filter_dict = self._workspace_filter(workspace)
if workspace: search_kwargs = {"k": k, "filter": filter_dict}
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: {search_kwargs}")
retriever = self.vector_store.as_retriever( retriever = self.vector_store.as_retriever(
search_type="similarity", search_type="similarity",
@@ -358,17 +380,14 @@ class AsyncRAGService(RAGService):
async def search_documents( async def search_documents(
self, query: str, workspace: Optional[DocumentWorkspace] = None, k: int = 4 self, query: str, workspace: Optional[DocumentWorkspace] = None, k: int = 4
) -> List[Document]: ) -> List[Document]:
"""Search relevant documents from the vector store.""" """Search relevant documents from the vector store (workspace required)."""
filter_dict = {} filter_dict = self._workspace_filter(workspace)
print(f"Do we have a workspace: {workspace}") print(f"Do we have a workspace: {workspace}")
if workspace: print(f"search_kwargs: {{'k': {k}, 'filter': {filter_dict}}}")
filter_dict["workspace_id"] = workspace.id
search_kwargs = {"k": k, "filter": filter_dict if filter_dict else None}
print(f"search_kwargs: {search_kwargs}")
retriever = self.vector_store.as_retriever( retriever = self.vector_store.as_retriever(
search_type="mmr", 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) return await retriever.aget_relevant_documents(query)
@@ -380,6 +399,8 @@ class AsyncRAGService(RAGService):
**kwargs, **kwargs,
) -> AsyncGenerator[str, None]: ) -> AsyncGenerator[str, None]:
"""Generate response with streaming support.""" """Generate response with streaming support."""
if workspace is None:
raise ValueError("workspace is required for RAG generation")
chain_input = { chain_input = {
"query": query, "query": query,
"conversation": conversation, "conversation": conversation,
@@ -0,0 +1,64 @@
from django.test import TestCase
from rest_framework_simplejwt.tokens import RefreshToken
from chat_backend.services.chat_tenant_scope import (
ChatCompanyScope,
ChatTenantScopeError,
resolve_chat_company_scope,
resolve_chat_user,
user_from_access_token,
)
from .factories import make_company, make_conversation, make_user, make_workspace
class ChatTenantScopeTestCase(TestCase):
def setUp(self):
self.company = make_company()
self.user = make_user(company=self.company)
self.workspace = make_workspace(self.company)
self.conversation = make_conversation(user=self.user)
def test_resolve_scope_for_owned_conversation(self):
scope = resolve_chat_company_scope(self.user, self.conversation.id)
self.assertEqual(
scope,
ChatCompanyScope(
user_id=self.user.id,
company_id=self.company.id,
workspace_id=self.workspace.id,
conversation_id=self.conversation.id,
),
)
def test_resolve_scope_rejects_cross_user_conversation(self):
other = make_user(
email="other-tenant@example.com", company=make_company("OtherCo")
)
foreign = make_conversation(user=other)
make_workspace(other.company)
with self.assertRaises(ChatTenantScopeError) as ctx:
resolve_chat_company_scope(self.user, foreign.id)
self.assertEqual(ctx.exception.code, "conversation_forbidden")
def test_resolve_chat_user_does_not_bind_identity_to_conversation(self):
other = make_user(company=make_company("VictimCo"), email="victim@example.com")
foreign = make_conversation(user=other)
resolved = resolve_chat_user(conversation_id=foreign.id)
self.assertIsNone(resolved)
def test_resolve_chat_user_prefers_jwt_over_email(self):
token = str(RefreshToken.for_user(self.user).access_token)
spoof = make_user(company=make_company("Spoof"), email="spoof@example.com")
resolved = resolve_chat_user(email=spoof.email, token=token)
self.assertEqual(resolved.id, self.user.id)
def test_user_from_access_token_rejects_garbage(self):
self.assertIsNone(user_from_access_token("not-a-jwt"))
+27 -1
View File
@@ -47,10 +47,35 @@ class DatabaseHelperTestCase(TransactionTestCase):
@parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)]) @parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)])
async def test_get_workspace(self, _name, module): async def test_get_workspace(self, _name, module):
workspace = await module.get_workspace(self.conversation.id) workspace = await module.get_workspace(
self.conversation.id, user=self.user
)
self.assertEqual(workspace.id, self.workspace.id) self.assertEqual(workspace.id, self.workspace.id)
@parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)])
async def test_get_workspace_rejects_other_users_conversation(self, _name, module):
other_company = await sync_to_async(make_company)("OtherCo")
other = await sync_to_async(make_user)(
email="other-tenant@example.com", company=other_company
)
other_conversation = await sync_to_async(make_conversation)(user=other)
await sync_to_async(make_workspace)(other_company)
with self.assertRaises(consumers.ChatTenantScopeError):
await module.get_workspace(other_conversation.id, user=self.user)
@parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)])
async def test_resolve_tenant_scope_binds_company_and_workspace(
self, _name, module
):
scope = await module.resolve_tenant_scope(self.user, self.conversation.id)
self.assertEqual(scope.user_id, self.user.id)
self.assertEqual(scope.company_id, self.company.id)
self.assertEqual(scope.workspace_id, self.workspace.id)
self.assertEqual(scope.conversation_id, self.conversation.id)
@parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)]) @parameterized.expand([("websocket", consumers), ("langgraph", consumers_graph)])
async def test_get_messages_stores_prompt_and_returns_history(self, _name, module): async def test_get_messages_stores_prompt_and_returns_history(self, _name, module):
messages, prompt_instance = await module.get_messages( messages, prompt_instance = await module.get_messages(
@@ -210,6 +235,7 @@ class GraphNodeTestCase(TransactionTestCase):
"response_generator": None, "response_generator": None,
"error": None, "error": None,
"model_name": "Turbo", "model_name": "Turbo",
"chat_user": self.user,
} }
state.update(overrides) state.update(overrides)
return state return state
@@ -62,7 +62,8 @@ class RAGServiceTestCase(TransactionTestCase):
self.addCleanup(reset_singletons) self.addCleanup(reset_singletons)
self.service = AsyncRAGService() self.service = AsyncRAGService()
self.workspace = make_workspace(make_company()) self.company = make_company()
self.workspace = make_workspace(self.company)
def _patch(self, target): def _patch(self, target):
patcher = mock.patch(target) patcher = mock.patch(target)
@@ -137,6 +138,7 @@ class RAGServiceTestCase(TransactionTestCase):
added = self.service.vector_store.add_documents.call_args[0][0] added = self.service.vector_store.add_documents.call_args[0][0]
self.assertIn("ingest me", added[0].page_content) self.assertIn("ingest me", added[0].page_content)
self.assertEqual(added[0].metadata["workspace_id"], self.workspace.id) self.assertEqual(added[0].metadata["workspace_id"], self.workspace.id)
self.assertEqual(added[0].metadata["company_id"], self.company.id)
self.assertEqual(added[0].metadata["document_id"], document.id) self.assertEqual(added[0].metadata["document_id"], document.id)
def test_ingest_documents_deletes_the_materialized_temp_file(self): def test_ingest_documents_deletes_the_materialized_temp_file(self):
@@ -229,15 +231,11 @@ class RAGServiceTestCase(TransactionTestCase):
search_kwargs={"k": 2, "filter": {"workspace_id": self.workspace.id}}, search_kwargs={"k": 2, "filter": {"workspace_id": self.workspace.id}},
) )
async def test_search_documents_without_workspace_has_no_filter(self): async def test_search_documents_without_workspace_fails_closed(self):
retriever = self.service.vector_store.as_retriever.return_value with self.assertRaises(ValueError):
retriever.aget_relevant_documents = mock.AsyncMock(return_value=[]) await self.service.search_documents("revenue")
await self.service.search_documents("revenue") self.service.vector_store.as_retriever.assert_not_called()
self.service.vector_store.as_retriever.assert_called_with(
search_type="mmr", search_kwargs={"k": 4, "filter": None}
)
async def test_format_history_labels_speakers(self): async def test_format_history_labels_speakers(self):
history = await self.service._format_history( history = await self.service._format_history(
@@ -87,6 +87,16 @@ class ConversationDetailViewTestCase(APITestCase):
[("hello", True), ("hi there", False)], [("hello", True), ("hi there", False)],
) )
def test_get_rejects_other_users_conversation(self):
other = make_user(email="other-tenant@example.com", company=make_company("Other"))
foreign = make_conversation(user=other)
make_prompt(foreign, message="secret")
response = self.client.get(self.url, {"conversation_id": foreign.id})
self.assertEqual(response.status_code, status.HTTP_404_NOT_FOUND)
self.assertEqual(Prompt.objects.filter(conversation=foreign).count(), 1)
def test_post_stores_assistant_prompt(self): def test_post_stores_assistant_prompt(self):
response = self.client.post( response = self.client.post(
self.url, self.url,
+13 -2
View File
@@ -502,7 +502,16 @@ class ConversationPreferences(APIView):
class ConversationDetailView(APIView): class ConversationDetailView(APIView):
def get(self, request, format="json"): def get(self, request, format="json"):
conversation_id = request.query_params.get("conversation_id") conversation_id = request.query_params.get("conversation_id")
prompts = Prompt.objects.filter(conversation__id=conversation_id) if not Conversation.objects.filter(
id=conversation_id, user=request.user, deleted=False
).exists():
return Response(
{"detail": "Conversation not found."},
status=status.HTTP_404_NOT_FOUND,
)
prompts = Prompt.objects.filter(
conversation__id=conversation_id, conversation__user=request.user
)
serailzer = PromptSerializer(prompts, many=True) serailzer = PromptSerializer(prompts, many=True)
return Response(serailzer.data, status=status.HTTP_200_OK) return Response(serailzer.data, status=status.HTTP_200_OK)
@@ -525,7 +534,9 @@ class ConversationDetailView(APIView):
is_user = bool(request.data.get("is_user")) is_user = bool(request.data.get("is_user"))
try: try:
conversation = Conversation.objects.get(id=conversation_id) conversation = Conversation.objects.get(
id=conversation_id, user=request.user, deleted=False
)
# add the prompt to the conversation # add the prompt to the conversation
serializer = PromptSerializer( serializer = PromptSerializer(