import asyncio import sys import types from src.agent_tools import TOOL_HANDLERS from src.agent_tools.document_tools import ( _owned_document_query, set_active_document, ) class _Column: def __init__(self, name): self.name = name def __eq__(self, value): return (self.name, "eq", value) def desc(self): return (self.name, "desc") def ilike(self, value): return (self.name, "ilike", value) class _Document: id = _Column("id") owner = _Column("owner") is_active = _Column("is_active") title = _Column("title") language = _Column("language") updated_at = _Column("updated_at") class _Query: def __init__(self, docs=None, first_doc=None): self.filters = [] self.docs = docs or [] self.first_doc = first_doc def filter(self, *clauses): self.filters.extend(clauses) return self def order_by(self, *args): return self def limit(self, *args): return self def all(self): return self.docs def first(self): return self.first_doc class _Db: def __init__(self, query): self.query_obj = query def query(self, *args): return self.query_obj def close(self): pass def _install_database_stub(monkeypatch, module_name, query): db = _Db(query) db_mod = types.ModuleType(module_name) db_mod.SessionLocal = lambda: db db_mod.Document = _Document db_mod.DocumentVersion = object db_mod.Session = object monkeypatch.setitem(sys.modules, module_name, db_mod) return db def test_owned_document_query_rejects_missing_owner(): query = _Query() assert _owned_document_query(query, _Document, None) is query assert False in query.filters def test_owned_document_query_filters_to_owner(): query = _Query() assert _owned_document_query(query, _Document, "alice") is query assert ("owner", "eq", "alice") in query.filters def test_manage_documents_list_filters_to_calling_owner(monkeypatch): query = _Query() _install_database_stub(monkeypatch, "core.database", query) result = asyncio.run( TOOL_HANDLERS["manage_documents"]('{"action":"list"}', {"owner": "alice"}) ) assert result["documents"] == [] assert ("owner", "eq", "alice") in query.filters def test_manage_documents_read_filters_to_calling_owner(monkeypatch): query = _Query() _install_database_stub(monkeypatch, "core.database", query) result = asyncio.run( TOOL_HANDLERS["manage_documents"]( '{"action":"read","document_id":"doc-bob"}', {"owner": "alice"} ) ) assert result["exit_code"] == 1 assert ("id", "eq", "doc-bob") in query.filters assert ("owner", "eq", "alice") in query.filters def test_update_document_active_id_filters_to_calling_owner(monkeypatch): query = _Query() _install_database_stub(monkeypatch, "src.database", query) set_active_document("doc-bob") try: result = asyncio.run( TOOL_HANDLERS["update_document"]("new content", {"owner": "alice"}) ) finally: set_active_document(None) assert result["exit_code"] == 1 assert "Requested document not found" in result["error"] assert ("id", "eq", "doc-bob") in query.filters assert ("owner", "eq", "alice") in query.filters def test_suggest_document_active_id_filters_to_calling_owner(monkeypatch): query = _Query() _install_database_stub(monkeypatch, "src.database", query) set_active_document("doc-bob") try: result = asyncio.run( TOOL_HANDLERS["suggest_document"]( "<<>>\nold\n<<>>\nnew\n<<>>\nbetter\n<<>>", {"owner": "alice"}, ) ) finally: set_active_document(None) assert result["error"] == "Document doc-bob not found" assert ("id", "eq", "doc-bob") in query.filters assert ("owner", "eq", "alice") in query.filters def test_suggest_document_accepts_reason_that_says_clearer(): from src.agent_tools.document_tools import parse_suggest_blocks parsed = parse_suggest_blocks( "<<>>\nvery good\n<<>>\nclear and actionable\n" "<<>>\nThe wording is more specific and clearer.\n<<>>" ) assert len(parsed) == 1 assert parsed[0]["replace"] == "clear and actionable" def test_suggestion_ids_are_stable_for_content_but_distinct_across_suggestions(): from src.agent_tools.document_tools import _stable_suggestion_id first = {"find": "wordy", "replace": "brief", "reason": "clarity"} second = {"find": "slow", "replace": "quick", "reason": "pace"} assert _stable_suggestion_id("doc-1", first) == _stable_suggestion_id("doc-1", first) assert _stable_suggestion_id("doc-1", first) != _stable_suggestion_id("doc-1", second) assert _stable_suggestion_id("doc-1", first) != _stable_suggestion_id("doc-2", first) def test_suggestion_matches_only_requested_email_sentence(monkeypatch): doc = types.SimpleNamespace(id='draft', current_content=( 'To: test@example.com\nIn-Reply-To: \n' 'X-Source-UID: 123\n---\n8am works for me.\nPrevious message stays.')) _install_database_stub(monkeypatch, 'src.database', _Query(first_doc=doc)) set_active_document('draft') try: result = asyncio.run(TOOL_HANDLERS['suggest_document']( '<<>>\n8am works for me.\n<<>>\nFriday at 8am works.\n<<>>\nAdd day\n<<>>', {'owner': 'alice'})) finally: set_active_document(None) assert result['suggestions'][0]['find'] == '8am works for me.' def test_document_tool_dispatch_forwards_owner(): source = open("src/tool_execution.py", encoding="utf-8").read() assert "_document_tool_dispatch(tool, content, session_id, owner)" in source # Also verify TOOL_HANDLERS has the expected entries for key in ("create_document", "update_document", "edit_document", "suggest_document", "manage_documents"): assert key in TOOL_HANDLERS, f"TOOL_HANDLERS missing key: {key}" assert callable(TOOL_HANDLERS[key]), f"TOOL_HANDLERS[{key!r}] is not callable"