fix: enforce collection membership on cloud query doc_ids

Cloud query()/query_stream() now run the same _require_document membership
guard the four doc-scoped ops received in 929b3df, so a doc_id from another
collection raises DocumentNotFoundError (matching the local backend's
_scoped_docs) instead of silently answering from another collection's
document. The guard runs after doc_ids normalization and before the
completion request; query_stream runs it via asyncio.to_thread. Folders
unavailable on the plan -> guard skips, as with the sibling ops.

Also rename the streaming QueryEvent types answer_delta/answer_done ->
text_delta/text_done: text_done fires once per assistant text message, which
is honest for the local agent loop's multi-message stream (the last one
before the stream ends carries the final answer). Corrects the cloud
backend's now-inaccurate "single terminal contract" comment. reasoning
stays reserved for future model-reasoning tokens. QueryEvent is unreleased
API (0.3.0.dev1), so the rename is not a breaking change.

README.md

@@ -179,7 +179,7 @@ import asyncio
 
 async def main():
     async for ev in col.query("Explain multi-head attention", doc_ids=doc_id, stream=True):
-        if ev.type == "answer_delta":
+        if ev.type == "text_delta":
             print(ev.data, end="", flush=True)
         elif ev.type == "tool_call":
             print(f"\n[tool] {ev.data['name']}")
@@ -187,7 +187,7 @@ async def main():
 asyncio.run(main())
 ```
 
-`ev.type` is one of: `tool_call`, `tool_result`, `answer_delta`, `answer_done`.
+`ev.type` is one of: `tool_call`, `tool_result`, `text_delta`, `text_done`. A `text_done` fires each time a text message completes — a local agentic query may emit several as the agent narrates between tool calls; the last `text_done` before the stream ends carries the final answer.
 
 ### Multi-document collections (experimental)
 

examples/cloud_demo.py

@@ -46,7 +46,7 @@ stream = col.query("What is the main contribution of this paper?", stream=True)
 async def main():
     streamed_text = False
     async for event in stream:
-        if event.type == "answer_delta":
+        if event.type == "text_delta":
             print(event.data, end="", flush=True)
             streamed_text = True
         elif event.type == "tool_call":
@@ -55,7 +55,7 @@ async def main():
                 streamed_text = False
             args = event.data.get("args", "")
             print(f"[tool call] {event.data['name']}({args})")
-        elif event.type == "answer_done":
+        elif event.type == "text_done":
             print()
             streamed_text = False
 

examples/demo_query_modes.py

@@ -68,7 +68,7 @@ async def stream_and_collect(coro_or_stream) -> list[str]:
         if ev.type == "tool_call":
             calls.append(ev.data["name"])
             print(f"  [tool] {ev.data['name']}({ev.data.get('args','')})")
-        elif ev.type == "answer_done":
+        elif ev.type == "text_done":
             text = str(ev.data)
             print(f"  [answer] {text[:160]}{'...' if len(text) > 160 else ''}")
     return calls

examples/local_demo.py

@@ -51,7 +51,7 @@ stream = col.query(
 async def main():
     streamed_text = False
     async for event in stream:
-        if event.type == "answer_delta":
+        if event.type == "text_delta":
             print(event.data, end="", flush=True)
             streamed_text = True
         elif event.type == "tool_call":
@@ -62,7 +62,7 @@ async def main():
         elif event.type == "tool_result":
             preview = str(event.data)[:200] + "..." if len(str(event.data)) > 200 else event.data
             print(f"[tool output] {preview}")
-        elif event.type == "answer_done":
+        elif event.type == "text_done":
             print()
             streamed_text = False
 

pageindex/agent.py

@@ -89,7 +89,7 @@ class QueryStream:
     Usage:
         stream = col.query("question", stream=True)
         async for event in stream:
-            if event.type == "answer_delta":
+            if event.type == "text_delta":
                 print(event.data, end="", flush=True)
     """
 
@@ -117,7 +117,7 @@ class QueryStream:
         async for event in streamed_run.stream_events():
             if isinstance(event, RawResponsesStreamEvent):
                 if isinstance(event.data, ResponseTextDeltaEvent):
-                    yield QueryEvent(type="answer_delta", data=event.data.delta)
+                    yield QueryEvent(type="text_delta", data=event.data.delta)
             elif isinstance(event, RunItemStreamEvent):
                 item = event.item
                 if item.type == "tool_call_item":
@@ -130,7 +130,7 @@ class QueryStream:
                 elif item.type == "message_output_item":
                     text = ItemHelpers.text_message_output(item)
                     if text:
-                        yield QueryEvent(type="answer_done", data=text)
+                        yield QueryEvent(type="text_done", data=text)
 
     def __aiter__(self):
         return self.stream_events()

pageindex/backend/cloud.py

@@ -256,6 +256,10 @@ class CloudBackend:
             )
         return meta
 
+    def _require_documents(self, collection: str, doc_ids: list[str]) -> None:
+        for doc_id in doc_ids:
+            self._require_document(collection, doc_id)
+
     def get_document(self, collection: str, doc_id: str, include_text: bool = False) -> dict:
         if include_text:
             import warnings
@@ -379,6 +383,8 @@ class CloudBackend:
             raise ValueError(
                 "doc_ids cannot be empty; pass None to query the whole collection"
             )
+        if doc_ids:
+            self._require_documents(collection, doc_ids)
         doc_id = doc_ids if doc_ids else self._get_all_doc_ids(collection)
         if not doc_id:
             raise ValueError("collection has no documents to query")
@@ -414,6 +420,8 @@ class CloudBackend:
             raise ValueError(
                 "doc_ids cannot be empty; pass None to query the whole collection"
             )
+        if doc_ids:
+            await asyncio.to_thread(self._require_documents, collection, doc_ids)
         doc_id = doc_ids if doc_ids else await asyncio.to_thread(
             self._get_all_doc_ids, collection
         )
@@ -516,11 +524,11 @@ class CloudBackend:
 
                     elif block_type == "text" and content:
                         answer_parts.append(content)
-                        _put(QueryEvent(type="answer_delta", data=content))
+                        _put(QueryEvent(type="text_delta", data=content))
 
-                # Same terminal contract as the local backend: a final
-                # answer_done event carrying the full answer text.
-                _put(QueryEvent(type="answer_done", data="".join(answer_parts)))
+                # The whole cloud answer is one text message, so its text_done
+                # carries the full answer text.
+                _put(QueryEvent(type="text_done", data="".join(answer_parts)))
 
             except requests.RequestException as e:
                 _put(CloudAPIError(f"Cloud streaming request failed: {e}"))

pageindex/events.py

@@ -5,5 +5,5 @@ from typing import Literal, Any
 @dataclass
 class QueryEvent:
     """Event emitted during streaming query."""
-    type: Literal["reasoning", "tool_call", "tool_result", "answer_delta", "answer_done"]
+    type: Literal["reasoning", "tool_call", "tool_result", "text_delta", "text_done"]
     data: Any

tests/test_cloud_backend.py

@@ -331,9 +331,37 @@ def test_doc_ops_skip_guard_when_folders_unavailable(monkeypatch):
     assert log == [("DELETE", "/doc/d1/")]
 
 
+def test_query_rejects_doc_from_another_collection(monkeypatch):
+    backend, log = _membership_backend(monkeypatch, doc_folder="f-other")
+    with pytest.raises(DocumentNotFoundError, match="collection 'col'"):
+        backend.query("col", "q", doc_ids=["d1"])
+    # guard must fail before the completion request goes out
+    assert all(method == "GET" and path.endswith("/metadata/") for method, path in log)
+
+
+def test_query_stream_rejects_doc_from_another_collection(monkeypatch):
+    backend, log = _membership_backend(monkeypatch, doc_folder="f-other")
+
+    async def _run():
+        async for _ in backend.query_stream("col", "q", doc_ids=["d1"]):
+            pass
+
+    with pytest.raises(DocumentNotFoundError, match="collection 'col'"):
+        asyncio.run(_run())
+    assert all(method == "GET" and path.endswith("/metadata/") for method, path in log)
+
+
+def test_query_checks_membership_then_completes(monkeypatch):
+    backend, log = _membership_backend(monkeypatch)
+    backend.query("col", "q", doc_ids=["d1"])
+    assert ("GET", "/doc/d1/metadata/") in log
+    assert ("POST", "/chat/completions/") in log
+
+
 # ── query_stream: terminal contract and error propagation ───────────────────
 
 def _collect_events(backend, **kwargs):
+    backend._folder_id_cache.setdefault("col", None)  # skip folder lookup
     async def _run():
         events = []
         async for ev in backend.query_stream("col", "q", doc_ids=["d1"], **kwargs):
@@ -349,7 +377,7 @@ def _sse(block_type, content):
     })
 
 
-def test_query_stream_emits_answer_done(monkeypatch):
+def test_query_stream_emits_text_done(monkeypatch):
     backend = CloudBackend(api_key="pi-test")
     lines = [_sse("text", "Hello "), _sse("text", "world"), "data: [DONE]"]
     monkeypatch.setattr(
@@ -357,13 +385,14 @@ def test_query_stream_emits_answer_done(monkeypatch):
         lambda *a, **k: FakeResponse(status_code=200, lines=lines),
     )
     events = _collect_events(backend)
-    assert [e.type for e in events] == ["answer_delta", "answer_delta", "answer_done"]
-    # Same contract as the local backend: answer_done carries the full text.
+    assert [e.type for e in events] == ["text_delta", "text_delta", "text_done"]
+    # The whole cloud answer is one text message, so text_done carries the full text.
     assert events[-1].data == "Hello world"
 
 
 def test_query_stream_http_error_raises(monkeypatch):
     backend = CloudBackend(api_key="pi-test")
+    backend._folder_id_cache["col"] = None  # skip folder lookup
     monkeypatch.setattr(
         cloud_mod.requests, "post",
         lambda *a, **k: FakeResponse(status_code=401, text="unauthorized"),
@@ -379,6 +408,7 @@ def test_query_stream_http_error_raises(monkeypatch):
 
 def test_query_stream_connect_failure_raises_instead_of_hanging(monkeypatch):
     backend = CloudBackend(api_key="pi-test")
+    backend._folder_id_cache["col"] = None  # skip folder lookup
 
     def fake_post(*a, **k):
         raise cloud_mod.requests.ConnectionError("dns failure")
@@ -397,6 +427,7 @@ def test_query_uses_long_timeout_and_single_attempt(monkeypatch):
     """Non-streaming chat completion is non-idempotent and slow: it must get
     a long timeout and must NOT be retried (each retry re-bills the query)."""
     backend = CloudBackend(api_key="pi-test")
+    backend._folder_id_cache["col"] = None  # skip folder lookup
     calls = []
 
     def fake_request(method, url, headers=None, **kwargs):
@@ -415,6 +446,7 @@ def test_query_stream_early_break_stops_background_thread(monkeypatch):
     drain the whole stream in the background."""
     import threading
     backend = CloudBackend(api_key="pi-test")
+    backend._folder_id_cache["col"] = None  # skip folder lookup
     drained_all = threading.Event()
 
     class SlowResponse:

tests/test_events.py

@@ -3,13 +3,13 @@ from pageindex.backend.protocol import AgentTools
 
 
 def test_query_event():
-    event = QueryEvent(type="answer_delta", data="hello")
-    assert event.type == "answer_delta"
+    event = QueryEvent(type="text_delta", data="hello")
+    assert event.type == "text_delta"
     assert event.data == "hello"
 
 
 def test_query_event_types():
-    for t in ["reasoning", "tool_call", "tool_result", "answer_delta", "answer_done"]:
+    for t in ["reasoning", "tool_call", "tool_result", "text_delta", "text_done"]:
         event = QueryEvent(type=t, data="test")
         assert event.type == t