diff --git a/README.md b/README.md index dfc1b4c8..c8a79240 100644 --- a/README.md +++ b/README.md @@ -78,6 +78,15 @@ nothing, and answers 404 with the chunks left in place. Because a 404 is indistinguishable from "already deleted", a caller that treats it as success will orphan those chunks silently. +**Presence checks name one owner.** `POST /exists` takes `file_ids` and an +optional `entity_id`, and answers `{"file_ids": [...]}` with the ids that hold +chunks owned by exactly that entity, or by the caller when no entity is named. +The retrieval routes read the caller and the entity together, so a copy the +caller embedded under their own id would answer for the entity there; this +route lets a client tell whether a file is already embedded for an agent before +embedding it again. It only narrows what the same request could read through +`/query`, and an older build answers 404, which a client can treat as "unknown". + **Upgrade the client first.** Deploy order matters, in one direction only: - A client that sends `entity_id` against an *older* build is inert — the diff --git a/app/models.py b/app/models.py index 57c01130..f459d969 100644 --- a/app/models.py +++ b/app/models.py @@ -43,3 +43,8 @@ class QueryMultipleBody(BaseModel): file_ids: List[str] k: int = 4 entity_id: Optional[str] = None + + +class ExistsRequestBody(BaseModel): + file_ids: List[str] + entity_id: Optional[str] = None diff --git a/app/routes/document_routes.py b/app/routes/document_routes.py index c3cdb2ee..861f3bd3 100644 --- a/app/routes/document_routes.py +++ b/app/routes/document_routes.py @@ -45,7 +45,7 @@ PARALLEL_EXECUTION, RAG_DISTANCE_THRESHOLD, ) -from app.scope import file_clause, files_clause, resolve_scope +from app.scope import file_clause, files_clause, resolve_scope, resolve_owner_scope # Warn once at import time if the user set a threshold under Atlas, where # the score direction is inverted (Atlas vectorSearchScore: higher = better) @@ -88,6 +88,7 @@ def _apply_distance_threshold(documents): StoreDocument, QueryRequestBody, DocumentResponse, + ExistsRequestBody, QueryMultipleBody, ) from app.services.vector_store.async_pg_vector import AsyncPgVector @@ -419,6 +420,50 @@ async def get_all_ids(request: Request, entity_id: str = None): raise HTTPException(status_code=500, detail=str(e)) +@router.post("/exists") +async def get_existing_file_ids(request: Request, body: ExistsRequestBody): + """Which of ``file_ids`` hold chunks owned by exactly one namespace. + + The owner is ``entity_id`` when given, otherwise the caller, so a client can + tell whether a file is already embedded for an agent without a similarity + search and without a copy under the caller's own id answering for it. + """ + scope = resolve_owner_scope(request, body.entity_id) + try: + if isinstance(vector_store, AsyncPgVector): + found = await vector_store.get_filtered_ids( + body.file_ids, + owners=scope.owners, + executor=request.app.state.thread_pool, + ) + else: + found = await run_in_executor( + request.app.state.thread_pool, + vector_store.get_filtered_ids, + body.file_ids, + scope.owners, + ) + + present = set(found) + return { + "file_ids": [id for id in dict.fromkeys(body.file_ids) if id in present] + } + except HTTPException as http_exc: + logger.error( + "HTTP Exception in get_existing_file_ids | Status: %d | Detail: %s", + http_exc.status_code, + http_exc.detail, + ) + raise http_exc + except Exception as e: + logger.error( + "Failed to check file ids | Error: %s | Traceback: %s", + str(e), + traceback.format_exc(), + ) + raise HTTPException(status_code=500, detail=str(e)) + + @router.get("/health") async def health_check(): try: diff --git a/app/scope.py b/app/scope.py index 5f7237f1..5aedcdcc 100644 --- a/app/scope.py +++ b/app/scope.py @@ -79,5 +79,23 @@ def resolve_scope(request: Request, entity_id: Optional[str] = None) -> ScopeFil return ScopeFilter(owners=tuple(sorted(owners))) +def resolve_owner_scope( + request: Request, entity_id: Optional[str] = None +) -> ScopeFilter: + """Resolve the single owner a presence check asks about. + + :func:`resolve_scope` reads the union of the caller and the named entity, + which is right for retrieval but cannot say *whose* chunks a file has: a copy + the caller embedded under their own id answers for the entity too. A presence + check needs one namespace, so a named entity replaces the caller here rather + than widening the set. That only ever narrows what ``resolve_scope`` would + read for the same request, so it reaches nothing the caller could not already + query. + """ + if entity_id: + return ScopeFilter(owners=(entity_id,)) + return resolve_scope(request) + + def scope_owners(scope: ScopeFilter) -> List[str]: return list(scope.owners) diff --git a/tests/test_authorization.py b/tests/test_authorization.py index a7094d38..67884413 100644 --- a/tests/test_authorization.py +++ b/tests/test_authorization.py @@ -406,3 +406,44 @@ def test_unauthenticated_deployment_keeps_reading_its_own_chunks(monkeypatch): assert listed.json() == [public_file] finally: ROWS.pop() + + +def test_exists_answers_for_the_named_entity_alone(attacker_headers): + """A presence check names one namespace. The caller's own row under the + shared file id must not answer for the agent, which ``/query`` cannot tell + apart because it reads the caller and the entity together.""" + response = client.post( + "/exists", + json={"file_ids": [SHARED_FILE, AGENT_FILE], "entity_id": AGENT}, + headers=attacker_headers, + ) + + assert response.status_code == 200 + assert response.json() == {"file_ids": [AGENT_FILE]} + + +def test_exists_without_entity_id_answers_for_the_caller(attacker_headers): + """With no entity named the caller is the namespace, so a victim's file id + reads as absent and the route is not an existence oracle.""" + response = client.post( + "/exists", + json={"file_ids": [VICTIM_FILE, SHARED_FILE, ATTACKER_FILE, SHARED_FILE]}, + headers=attacker_headers, + ) + + assert response.status_code == 200 + assert response.json() == {"file_ids": [SHARED_FILE, ATTACKER_FILE]} + + +def test_exists_on_an_unauthenticated_deployment(monkeypatch): + """Without ``JWT_SECRET`` writes stamp the entity or ``public``, and the + presence check resolves to the same owner.""" + monkeypatch.delenv("JWT_SECRET", raising=False) + + named = client.post( + "/exists", json={"file_ids": [AGENT_FILE, VICTIM_FILE], "entity_id": AGENT} + ) + unnamed = client.post("/exists", json={"file_ids": [AGENT_FILE]}) + + assert named.json() == {"file_ids": [AGENT_FILE]} + assert unnamed.json() == {"file_ids": []} diff --git a/tests/test_query_execution.py b/tests/test_query_execution.py index 4e4d7518..c7118c3b 100644 --- a/tests/test_query_execution.py +++ b/tests/test_query_execution.py @@ -11,7 +11,7 @@ from langchain_core.documents import Document from app.config import VectorDBType -from app.models import QueryMultipleBody, QueryRequestBody +from app.models import ExistsRequestBody, QueryMultipleBody, QueryRequestBody from app.routes import document_routes from app.services.vector_store.async_pg_vector import AsyncPgVector from app.services.vector_store.extended_pg_vector import ExtendedPgVector @@ -253,3 +253,41 @@ async def test_empty_result_contract_unchanged(query_context, query_route): await query_context.query() assert caught.value.status_code == 404 assert caught.value.detail == "No documents found for the given query" + + +async def test_sync_store_presence_check_leaves_event_loop_responsive(monkeypatch): + """A synchronous store (Atlas) answers ``/exists`` with a blocking driver call, + so it must run in the worker pool rather than on the event loop.""" + loop = asyncio.get_running_loop() + started, release = asyncio.Event(), threading.Event() + + def get_filtered_ids(ids, owners): + assert threading.current_thread().name.startswith("exists-test") + assert list(owners) == ["agent-1"] + loop.call_soon_threadsafe(started.set) + if not release.wait(timeout=5): + raise TimeoutError("event loop did not release the worker") + return ["file-2"] + + store = SimpleNamespace(get_filtered_ids=get_filtered_ids) + monkeypatch.setattr(document_routes, "vector_store", store) + + with ThreadPoolExecutor(max_workers=1, thread_name_prefix="exists-test") as pool: + req = SimpleNamespace( + app=SimpleNamespace(state=SimpleNamespace(thread_pool=pool)), + state=SimpleNamespace(user={"id": "owner-1"}), + ) + task = asyncio.create_task( + document_routes.get_existing_file_ids( + req, + ExistsRequestBody(file_ids=["file-1", "file-2"], entity_id="agent-1"), + ) + ) + try: + await asyncio.wait_for(started.wait(), timeout=2) + assert not task.done() + finally: + release.set() + result = await asyncio.wait_for(task, timeout=2) + + assert result == {"file_ids": ["file-2"]}