From 4992cc620fced5b276a13b513a100f1eec9a91a7 Mon Sep 17 00:00:00 2001 From: Alexandre Teixeira Date: Thu, 29 Oct 2026 02:09:37 +0000 Subject: [PATCH] test: protect chat document context owner boundaries --- tests/test_chat_document_owner_context.py | 119 ++++++++++++++++++++++ 1 file changed, 119 insertions(+) create mode 100644 tests/test_chat_document_owner_context.py diff --git a/tests/test_chat_document_owner_context.py b/tests/test_chat_document_owner_context.py new file mode 100644 index 000000000..fbb311313 --- /dev/null +++ b/tests/test_chat_document_owner_context.py @@ -0,0 +1,119 @@ +"""The chat stream must never pass another owner's document to the model.""" + +import json +from types import SimpleNamespace + +import pytest +from fastapi import FastAPI +from fastapi.testclient import TestClient + +from core.auth import AuthManager +from routes.auth_routes import SESSION_COOKIE +from tests.helpers.database import isolated_session_database + + +@pytest.fixture +def document_chat(tmp_path, monkeypatch): + from routes import chat_routes, chat_helpers, session_routes + from src.agent_tools import document_tools + from src.auth_helpers import effective_user + + monkeypatch.setenv("AUTH_ENABLED", "true") + auth = AuthManager(str(tmp_path / "auth.json")) + for owner in ("alice", "bob"): + assert auth.create_user(owner, "disposable-test-password") + captured = [] + + async def context(sess, request, *args, **kwargs): + # Preprocessing and external model I/O are outside the owner lookup. + # The actual route, session-owner gate and document ORM queries run. + return chat_helpers.ChatContext( + preface=[], rag_sources=[], web_sources=[], used_memories=[], + messages=[{"role": "user", "content": "Summarize my writing"}], + context_length=4096, was_compacted=False, user=effective_user(request), uprefs={}, + preset=chat_helpers.PresetInfo(0.2, 128, None, None), + preprocessed=chat_helpers.PreprocessedMessage("", "", "", [], []), + ) + + async def model(endpoint, model, messages, **kwargs): + doc = kwargs.get("active_document") + captured.append(None if doc is None else (doc.id, doc.owner, doc.current_content)) + text = doc.current_content if doc is not None else "no document context" + yield f'data: {json.dumps({"delta": text})}\n\n' + yield 'data: [DONE]\n\n' + + monkeypatch.setattr(chat_routes, "build_chat_context", context) + monkeypatch.setattr(chat_routes, "stream_agent_loop", model) + monkeypatch.setattr(chat_routes, "save_assistant_response", lambda *a, **kw: None) + monkeypatch.setattr(chat_routes, "run_post_response_tasks", lambda *a, **kw: None) + monkeypatch.setattr(document_tools, "_active_document_id", None) + with isolated_session_database(tmp_path) as (manager, database): + factory = database.SessionLocal + for module in (chat_routes, chat_helpers, session_routes): + monkeypatch.setattr(module, "SessionLocal", factory) + # Existing route tests may re-import session_routes after chat imported + # its owner check. Bind that real function's database too, without + # replacing the authorization logic or reloading production modules. + monkeypatch.setitem(chat_routes._verify_session_owner.__globals__, "SessionLocal", factory) + with factory() as db: + db.add_all([ + database.Session(id="alice-chat", name="Alice", owner="alice", + endpoint_url="http://model.invalid/v1", model="fixture-model"), + database.Session(id="bob-chat", name="Bob", owner="bob", + endpoint_url="http://model.invalid/v1", model="fixture-model"), + database.ModelEndpoint(id="fixture-endpoint", name="Fixture", + base_url="http://model.invalid/v1", owner="alice"), + ]) + db.commit() + app = FastAPI() + app.state.auth_manager = auth + + @app.middleware("http") + async def fixture_session(request, call_next): + request.state.current_user = auth.get_username_for_token(request.cookies.get(SESSION_COOKIE)) + return await call_next(request) + + app.include_router(chat_routes.setup_chat_routes( + manager, SimpleNamespace(), SimpleNamespace(), SimpleNamespace(), + SimpleNamespace(), SimpleNamespace(), + )) + token = auth.create_session("alice", "disposable-test-password") + assert token + with TestClient(app) as client: + client.cookies.set(SESSION_COOKIE, token) + yield SimpleNamespace(client=client, factory=factory, database=database, + captured=captured, document_tools=document_tools) + + +@pytest.mark.parametrize("source", ["explicit", "session", "email-session", "memory"]) +@pytest.mark.parametrize("document_owner", ["alice", "bob"]) +def test_chat_document_context_is_owner_scoped(document_chat, source, document_owner): + case = document_chat + form = {"message": "Summarize my writing", "session": "alice-chat", "mode": "agent", + "compare_mode": "true", "no_memory": "true", "use_rag": "false"} + with case.factory() as db: + for owner in ("alice", "bob"): + # A stale session link is not evidence that the caller owns a doc. + linked = "alice-chat" if source in {"session", "email-session"} and owner == document_owner else None + db.add(case.database.Document(id=f"{owner}-doc", owner=owner, session_id=linked, + title=f"{owner} draft", + language="email" if source == "email-session" else "markdown", + current_content=f"{owner}-private-context", is_active=True)) + db.commit() + if source == "explicit": + form.update(active_doc_id=f"{document_owner}-doc", active_doc_state="visible") + elif source == "memory": + case.document_tools.set_active_document(f"{document_owner}-doc") + response = case.client.post("/api/chat_stream", data=form) + assert response.status_code == 200, response.text + if document_owner == "alice": + assert case.captured == [("alice-doc", "alice", "alice-private-context")] + assert "alice-private-context" in response.text + else: + assert case.captured == [None] + assert "bob-private-context" not in response.text + with case.factory() as db: + assert db.get(case.database.Document, "bob-doc").owner == "bob" + assert db.get(case.database.Document, "bob-doc").session_id == ( + "alice-chat" if source in {"session", "email-session"} and document_owner == "bob" else None + )