fix: six P2 correctness/robustness cleanups from the SDK review

- legacy call_llm: open AsyncOpenAI via `async with` so the client (and
  its HTTP connection pool) is closed instead of leaked.
- LocalBackend.add_document: fail fast with CollectionNotFoundError when
  the collection doesn't exist, before the expensive parse + LLM index
  (previously the missing FK only tripped at save time, after paying for
  the LLM work). Also raise builtin FileNotFoundError for a missing path
  instead of FileTypeError (which now means only "unsupported extension").
- Collection.query(doc_ids=None): the empty-collection guard now always
  runs — previously it was skipped once PAGEINDEX_EXPERIMENTAL_MULTIDOC
  was set. A single list_documents call serves both the guard and the
  multi-doc warning (no separate call just to decide whether to warn).
- CloudBackend.query_stream: on early consumer break / GeneratorExit,
  signal the background SSE thread to stop and force-close the response,
  so it no longer drains the whole stream in the background.

Adds regression tests for each (client closed, fail-fast on unknown
collection, FileNotFoundError, empty-check under the multidoc env flag,
single list call, early-break thread stop).

Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS

pageindex/backend/cloud.py

@@ -344,6 +344,11 @@ class CloudBackend:
         # Queue carries QueryEvent, an Exception to re-raise, or None (end).
         queue: asyncio.Queue[QueryEvent | Exception | None] = asyncio.Queue()
         loop = asyncio.get_running_loop()
+        # Set when the consumer stops early (break / GeneratorExit) so the
+        # background thread stops draining the SSE stream instead of pulling
+        # it to completion in the background.
+        stop = threading.Event()
+        resp_holder: dict[str, requests.Response] = {}
 
         def _put(item: QueryEvent | Exception | None) -> None:
             try:
@@ -374,6 +379,7 @@ class CloudBackend:
                     stream=True,
                     timeout=120,
                 )
+                resp_holder["resp"] = resp
                 if resp.status_code != 200:
                     body = resp.text[:500] if resp.text else ""
                     raise CloudAPIError(
@@ -385,6 +391,8 @@ class CloudBackend:
                 current_tool_args: list[str] = []
 
                 for line in resp.iter_lines(decode_unicode=True):
+                    if stop.is_set():
+                        return  # consumer abandoned the stream
                     if not line or not line.startswith("data: "):
                         continue
                     data_str = line[6:]
@@ -439,15 +447,26 @@ class CloudBackend:
         thread = threading.Thread(target=_stream, daemon=True)
         thread.start()
 
-        while True:
-            item = await queue.get()
-            if item is None:
-                break
-            if isinstance(item, Exception):
-                raise item
-            yield item
-
-        thread.join(timeout=5)
+        try:
+            while True:
+                item = await queue.get()
+                if item is None:
+                    break
+                if isinstance(item, Exception):
+                    raise item
+                yield item
+        finally:
+            # On early break / GeneratorExit / raised error: tell the thread to
+            # stop and force-close the response so a read blocked mid-stream
+            # unblocks instead of draining the rest in the background.
+            stop.set()
+            resp = resp_holder.get("resp")
+            if resp is not None:
+                try:
+                    resp.close()
+                except Exception:
+                    pass
+            thread.join(timeout=5)
 
     def _get_all_doc_ids(self, collection: str) -> list[str]:
         """Get all document IDs in a collection."""

pageindex/backend/local.py

@@ -13,7 +13,8 @@ from ..storage.protocol import StorageEngine
 from ..index.pipeline import build_index
 from ..index.utils import parse_pages, get_pdf_page_content, get_md_page_content, remove_fields
 from ..backend.protocol import AgentTools
-from ..errors import FileTypeError, DocumentNotFoundError, IndexingError, PageIndexError
+from ..errors import (FileTypeError, DocumentNotFoundError, CollectionNotFoundError,
+                      IndexingError, PageIndexError)
 
 _COLLECTION_NAME_RE = re.compile(r'^[a-zA-Z0-9_-]{1,128}$')
 
@@ -79,7 +80,16 @@ class LocalBackend:
     def add_document(self, collection: str, file_path: str) -> str:
         file_path = os.path.realpath(file_path)
         if not os.path.isfile(file_path):
-            raise FileTypeError(f"Not a regular file: {file_path}")
+            # Missing path is a file-not-found error, not an unsupported-type one.
+            raise FileNotFoundError(f"No such file: {file_path}")
+        # Fail fast before the expensive parse + LLM indexing if the collection
+        # doesn't exist — otherwise the FK constraint only trips at save time,
+        # after the LLM work (and its cost) is already spent.
+        if collection not in self._storage.list_collections():
+            raise CollectionNotFoundError(
+                f"Collection '{collection}' does not exist; "
+                f"create it first (e.g. client.collection('{collection}'))."
+            )
         parser = self._resolve_parser(file_path)
 
         # Dedup is content-only — same file is reused regardless of IndexConfig

pageindex/collection.py

@@ -94,14 +94,16 @@ class Collection:
             raise ValueError(
                 "doc_ids cannot be empty; pass None to query the whole collection"
             )
-        if doc_ids is None and not _multidoc_acked():
+        if doc_ids is None:
+            # One list_documents call serves both the empty-collection guard
+            # (always) and the multi-doc warning (only when not acknowledged).
             docs = self._backend.list_documents(self._name)
             if not docs:
                 raise ValueError(
                     f"Cannot query collection '{self._name}': it is empty. "
                     "Add documents with col.add(...) first."
                 )
-            if len(docs) > 1:
+            if len(docs) > 1 and not _multidoc_acked():
                 warnings.warn(_MULTIDOC_WARNING, UserWarning, stacklevel=2)
         if stream:
             return QueryStream(self._backend, self._name, question, doc_ids)

pageindex/index/utils.py

@@ -460,12 +460,12 @@ async def call_llm(prompt, api_key, model="gpt-4.1", temperature=0):
     """
     import openai
 
-    client = openai.AsyncOpenAI(api_key=api_key)
-    response = await client.chat.completions.create(
-        model=model,
-        messages=[{"role": "user", "content": prompt}],
-        temperature=temperature,
-    )
+    async with openai.AsyncOpenAI(api_key=api_key) as client:
+        response = await client.chat.completions.create(
+            model=model,
+            messages=[{"role": "user", "content": prompt}],
+            temperature=temperature,
+        )
     return response.choices[0].message.content.strip()
 
 

tests/test_cloud_backend.py

@@ -219,3 +219,32 @@ def test_query_uses_long_timeout_and_single_attempt(monkeypatch):
         backend.query("col", "q", doc_ids=["d1"])
     assert len(calls) == 1
     assert calls[0]["timeout"] == 300
+
+
+def test_query_stream_early_break_stops_background_thread(monkeypatch):
+    """Consumer breaking early must signal the SSE thread to stop, not let it
+    drain the whole stream in the background."""
+    import threading
+    import time as _real_time  # autouse fixture stubs cloud_mod.time.sleep, not this
+    backend = CloudBackend(api_key="pi-test")
+    drained_all = threading.Event()
+
+    class SlowResponse:
+        status_code = 200
+        text = ""
+        def iter_lines(self, decode_unicode=True):
+            for i in range(1000):
+                yield _sse("text", f"chunk{i} ")
+                _real_time.sleep(0.002)  # pace so the consumer reliably breaks first
+            drained_all.set()  # only reached if the thread was NOT stopped
+        def close(self):
+            pass
+
+    monkeypatch.setattr(cloud_mod.requests, "post", lambda *a, **k: SlowResponse())
+
+    async def _run():
+        async for _ in backend.query_stream("col", "q", doc_ids=["d1"]):
+            break  # consume one event, then abandon
+
+    asyncio.run(_run())
+    assert not drained_all.is_set()

tests/test_collection.py

@@ -94,3 +94,26 @@ def test_query_accepts_str_doc_id(col):
 def test_query_rejects_empty_list(col):
     with pytest.raises(ValueError, match="cannot be empty"):
         col.query("what?", doc_ids=[])
+
+
+def test_empty_collection_check_runs_even_when_multidoc_acked(monkeypatch):
+    from unittest.mock import MagicMock
+    from pageindex.collection import Collection
+    monkeypatch.setenv("PAGEINDEX_EXPERIMENTAL_MULTIDOC", "1")
+    backend = MagicMock()
+    backend.list_documents.return_value = []
+    col = Collection(name="papers", backend=backend)
+    with pytest.raises(ValueError, match="empty"):
+        col.query("q")  # doc_ids=None, collection empty -> must still raise
+
+
+def test_whole_collection_query_lists_documents_once(monkeypatch):
+    from unittest.mock import MagicMock
+    from pageindex.collection import Collection
+    monkeypatch.setenv("PAGEINDEX_EXPERIMENTAL_MULTIDOC", "1")  # silence warning path
+    backend = MagicMock()
+    backend.list_documents.return_value = [{"doc_id": "d1"}, {"doc_id": "d2"}]
+    backend.query.return_value = "ans"
+    col = Collection(name="papers", backend=backend)
+    col.query("q")
+    assert backend.list_documents.call_count == 1  # single call at the collection layer

tests/test_legacy_utils_contract.py

@@ -75,6 +75,7 @@ def test_print_tree_keeps_legacy_exclude_fields(capsys):
 
 def test_call_llm_keeps_legacy_async_openai_contract(monkeypatch):
     calls = []
+    closed = []
 
     class FakeCompletions:
         async def create(self, **kwargs):
@@ -88,6 +89,15 @@ def test_call_llm_keeps_legacy_async_openai_contract(monkeypatch):
             self.api_key = api_key
             self.chat = SimpleNamespace(completions=FakeCompletions())
 
+        # call_llm must open the client as an async context manager so it is
+        # closed (no leaked HTTP connection pool).
+        async def __aenter__(self):
+            return self
+
+        async def __aexit__(self, *exc):
+            closed.append(True)
+            return False
+
     fake_openai = SimpleNamespace(AsyncOpenAI=FakeAsyncOpenAI)
     monkeypatch.setitem(sys.modules, "openai", fake_openai)
 
@@ -99,6 +109,7 @@ def test_call_llm_keeps_legacy_async_openai_contract(monkeypatch):
     ))
 
     assert result == "answer"
+    assert closed == [True]  # client was closed
     assert calls == [{
         "model": "gpt-test",
         "messages": [{"role": "user", "content": "hello"}],

tests/test_local_backend.py

@@ -165,3 +165,18 @@ def test_delete_collection_rejects_path_traversal(backend, tmp_path):
     with pytest.raises(PageIndexError, match="Invalid collection name"):
         backend.delete_collection("../..")
     assert canary.exists()
+
+
+def test_add_document_missing_file_raises_file_not_found(backend, tmp_path):
+    backend.get_or_create_collection("papers")
+    with pytest.raises(FileNotFoundError):
+        backend.add_document("papers", str(tmp_path / "nope.pdf"))
+
+
+def test_add_document_unknown_collection_fails_fast(backend, tmp_path):
+    from pageindex.errors import CollectionNotFoundError
+    pdf = tmp_path / "doc.pdf"
+    pdf.write_bytes(b"%PDF-1.4")
+    # Collection never created -> must raise before any parse/LLM work.
+    with pytest.raises(CollectionNotFoundError, match="does not exist"):
+        backend.add_document("ghost-collection", str(pdf))