fix: resolve concurrency and DoS pitfalls from PR #272 review

- index/utils.py: fix an asyncio deadlock in _sync_llm_semaphore. Sync LLM
  calls (check_toc / process_no_toc / toc_transformer) run on the event-loop
  thread nested inside the async meta_processor, and its blocking ceiling
  acquire could wait forever for a permit held by async _llm_semaphore holders
  that can only release it once the (now-frozen) loop runs. Take the slot
  non-blocking when on a running loop; keep the blocking acquire off-loop.
- index/utils.py: bound parse_pages ranges before materializing range() into
  the list, so a huge span like '1-2000000000' is rejected up front instead of
  exhausting memory before the 1000-page cap is ever checked (DoS).
- storage/sqlite.py: bump a generation counter on close() so a thread that
  cached a connection in thread-local storage reconnects on its next call
  instead of reusing a closed handle (ProgrammingError).
- backend/cloud.py: after connect, bail out of the SSE background thread if the
  consumer already abandoned the stream, instead of draining it in the background.

Adds regression tests for each fix.

pageindex/backend/cloud.py

@@ -396,6 +396,13 @@ class CloudBackend:
                     timeout=120,
                 )
                 resp_holder["resp"] = resp
+                # The consumer may have abandoned the stream while we were still
+                # blocked in requests.post() (its connect phase, before resp
+                # existed to close). Now that resp exists, bail immediately
+                # rather than reading/draining a stream nobody is listening to;
+                # the finally block closes resp and pushes the sentinel.
+                if stop.is_set():
+                    return
                 if resp.status_code != 200:
                     body = resp.text[:500] if resp.text else ""
                     raise CloudAPIError(

pageindex/index/utils.py

@@ -116,11 +116,34 @@ def _sync_llm_semaphore():
 
     It uses the same process-wide ceiling so sync and async LLM calls share one
     real cap. A scoped override can only narrow that cap for the active context.
+
+    A *blocking* ceiling acquire is only safe OFF the event-loop thread. Several
+    sync LLM helpers (``check_toc`` → ``toc_detector_single_page``,
+    ``process_no_toc`` → ``generate_toc_init``, ``toc_transformer``, …) are
+    called synchronously from inside async coroutines (``meta_processor`` →
+    ``process_large_node_recursively``), i.e. ON the running loop. There, the
+    async ``_llm_semaphore`` holders own the ceiling permits and can only
+    release them by resuming on that same loop — so a blocking acquire here
+    would freeze the loop and *deadlock*: the permit it waits for can never be
+    freed. When we detect a running loop we therefore take a slot only if one is
+    immediately free (non-blocking) and otherwise proceed without it. That's
+    safe: a sync call monopolizes the loop thread while it runs, so it's already
+    serialized on this loop and can't multiply the in-flight count beyond one
+    extra per loop.
     """
+    try:
+        asyncio.get_running_loop()
+        on_event_loop = True
+    except RuntimeError:
+        on_event_loop = False
+
     ceiling_sem = _process_ceiling_semaphore()
-    ceiling_sem.acquire()
-    # Only set once the permit is actually held (mirrors _llm_semaphore): guards
-    # against releasing a permit we never acquired if acquire() is interrupted.
+    # Blocking acquire() (off-loop) always returns True; acquire(False) (on-loop)
+    # may return False, meaning "no free permit — proceed without one" rather
+    # than block the loop into a deadlock.
+    held_ceiling = ceiling_sem.acquire(False) if on_event_loop else ceiling_sem.acquire()
+    # Only track a permit we actually hold (mirrors _llm_semaphore): guards
+    # against releasing one we never acquired.
     scoped_sem = None
     try:
         effective = get_max_concurrency()
@@ -128,13 +151,19 @@ def _sync_llm_semaphore():
         if effective < ceiling:
             candidate = _max_concurrency_scope_semaphore()
             if candidate is not None:
-                candidate.acquire()
-                scoped_sem = candidate
+                # Same rule for the scoped cap: never block the loop for it.
+                if on_event_loop:
+                    if candidate.acquire(False):
+                        scoped_sem = candidate
+                else:
+                    candidate.acquire()
+                    scoped_sem = candidate
         yield
     finally:
         if scoped_sem is not None:
             scoped_sem.release()
-        ceiling_sem.release()
+        if held_ceiling:
+            ceiling_sem.release()
 
 
 def llm_completion(model, prompt, chat_history=None, return_finish_reason=False):
@@ -530,6 +559,9 @@ def remove_structure_text(data):
 
 # ── Functions migrated from retrieve.py ──────────────────────────────────────
 
+_MAX_PAGES = 1000
+
+
 def parse_pages(pages: str) -> list[int]:
     """Parse a pages string like '5-7', '3,8', or '12' into a sorted list of ints."""
     result = []
@@ -539,13 +571,23 @@ def parse_pages(pages: str) -> list[int]:
             start, end = int(part.split('-', 1)[0].strip()), int(part.split('-', 1)[1].strip())
             if start > end:
                 raise ValueError(f"Invalid range '{part}': start must be <= end")
+            # Bound the span BEFORE materializing range() into the list. Checking
+            # len(result) only after `result.extend(range(...))` is too late: a
+            # single huge span like '1-2000000000' allocates billions of ints
+            # and exhausts memory before the cap is ever reached (DoS). page_nums
+            # is attacker/LLM-reachable via get_page_content.
+            span = end - start + 1
+            if span > _MAX_PAGES or len(result) + span > _MAX_PAGES:
+                raise ValueError(f"Page range too large: max {_MAX_PAGES} pages")
             result.extend(range(start, end + 1))
         else:
+            if len(result) + 1 > _MAX_PAGES:
+                raise ValueError(f"Page range too large: max {_MAX_PAGES} pages")
             result.append(int(part))
     result = [p for p in result if p >= 1]
     result = sorted(set(result))
-    if len(result) > 1000:
-        raise ValueError(f"Page range too large: {len(result)} pages (max 1000)")
+    if len(result) > _MAX_PAGES:
+        raise ValueError(f"Page range too large: {len(result)} pages (max {_MAX_PAGES})")
     return result
 
 

pageindex/storage/sqlite.py

@@ -27,6 +27,14 @@ class SQLiteStorage:
         self._local = threading.local()
         self._connections: list[sqlite3.Connection] = []
         self._conn_lock = threading.Lock()
+        # Bumped by close(). A thread caches its connection in thread-local
+        # storage, so after close() every OTHER thread's thread-local still
+        # points at a now-closed connection. Comparing the cached generation
+        # against this counter lets _get_conn detect that and reconnect, instead
+        # of handing back a closed connection (sqlite3.ProgrammingError). close()
+        # can only touch its OWN thread-local, so this is the only way to
+        # invalidate the others consistently.
+        self._generation = 0
         # Serializes the (fast) write operations within this process so
         # concurrent indexing threads don't collide on WAL's single writer
         # ("database is locked"). Reads stay concurrent; the expensive LLM
@@ -36,8 +44,13 @@ class SQLiteStorage:
         self._init_schema()
 
     def _get_conn(self) -> sqlite3.Connection:
-        """Return a thread-local SQLite connection."""
-        if not hasattr(self._local, "conn"):
+        """Return a thread-local SQLite connection.
+
+        Reconnects if this thread has no connection yet OR its cached connection
+        was invalidated by a close() on another thread (generation mismatch).
+        """
+        if (not hasattr(self._local, "conn")
+                or getattr(self._local, "generation", None) != self._generation):
             # Each thread gets its own connection (threading.local), so
             # statements never race. check_same_thread=False exists solely so
             # close() can close every tracked connection from whichever thread
@@ -55,6 +68,7 @@ class SQLiteStorage:
             conn.execute("PRAGMA foreign_keys=ON")
             conn.execute("PRAGMA busy_timeout=10000")
             self._local.conn = conn
+            self._local.generation = self._generation
             with self._conn_lock:
                 self._connections.append(conn)
         return self._local.conn
@@ -213,6 +227,10 @@ class SQLiteStorage:
                 except Exception:
                     pass
             self._connections.clear()
+            # Invalidate every thread's cached connection. close() can only
+            # del its OWN thread-local, so the bump is what makes _get_conn on
+            # any other thread reconnect instead of reusing a closed handle.
+            self._generation += 1
         if hasattr(self._local, "conn"):
             del self._local.conn
 

tests/test_concurrency.py

@@ -290,6 +290,56 @@ def test_llm_completion_holds_the_shared_semaphore(monkeypatch):
     assert state["peak"] == 1
 
 
+def test_sync_llm_completion_on_event_loop_does_not_deadlock(monkeypatch):
+    # Regression: _sync_llm_semaphore used a BLOCKING ceiling acquire. Sync LLM
+    # helpers (check_toc, process_no_toc, toc_transformer, …) run synchronously
+    # ON the event loop (nested inside the async meta_processor). If async
+    # llm_acompletion holders occupy every ceiling permit across their awaits,
+    # a blocking acquire froze the loop -> the holders could never resume to
+    # release their permits -> permanent deadlock. The sync path must never
+    # block the running loop.
+    set_max_concurrency(2)
+
+    async def fake_acompletion(**kwargs):
+        await asyncio.sleep(0.3)  # hold a ceiling permit across the await
+        return SimpleNamespace(
+            choices=[SimpleNamespace(message=SimpleNamespace(content="ok"))]
+        )
+
+    def fake_completion(**kwargs):
+        return SimpleNamespace(
+            choices=[SimpleNamespace(
+                message=SimpleNamespace(content="sync-ok"), finish_reason="stop")]
+        )
+
+    monkeypatch.setattr("litellm.acompletion", fake_acompletion)
+    monkeypatch.setattr("litellm.completion", fake_completion)
+
+    async def run():
+        # Both ceiling permits taken by async holders, held across their await.
+        holders = [asyncio.create_task(llm_acompletion("m", f"p{i}")) for i in range(2)]
+        await asyncio.sleep(0.05)  # let them acquire the permits
+        # Sync call on the loop thread: pre-fix this blocks forever waiting for a
+        # permit the holders own and can't release (loop is frozen).
+        result = llm_completion("m", "sync")
+        await asyncio.gather(*holders)
+        return result
+
+    # Run in a thread with a join timeout so a regression FAILS instead of
+    # hanging CI: a real deadlock freezes the loop, so asyncio.wait_for can't
+    # cancel it (its timeout callback never runs on the frozen loop).
+    box = {}
+
+    def target():
+        box["result"] = asyncio.run(run())
+
+    t = threading.Thread(target=target, daemon=True)
+    t.start()
+    t.join(timeout=8)
+    assert not t.is_alive(), "deadlock: sync llm_completion blocked the event loop"
+    assert box["result"] == "sync-ok"
+
+
 def test_run_async_propagates_scope_into_worker_thread():
     # When build_index runs inside an already-running loop, _run_async hops to a
     # worker thread. The max_concurrency_scope override must ride along (copied

tests/test_page_content.py

@@ -50,6 +50,31 @@ def test_retrieve_parse_pages_delegates_to_canonical_and_enforces_dos_cap():
         _parse_pages("1-99999999")
 
 
+def test_parse_pages_caps_a_huge_range_without_materializing_it():
+    """Regression: the 1000-page cap was checked only AFTER
+    `result.extend(range(start, end + 1))`, so a single huge span like
+    '1-2000000000' allocated billions of ints and OOM'd before the check ran.
+    The span must be rejected up front, quickly, without building the list."""
+    import time
+    import pytest
+    from pageindex.index.utils import parse_pages
+
+    start = time.monotonic()
+    with pytest.raises(ValueError, match="too large"):
+        parse_pages("1-2000000000")
+    # Must be near-instant (no billion-element allocation). Generous bound to
+    # avoid flakiness while still failing loudly on a re-materializing regression.
+    assert time.monotonic() - start < 1.0
+
+    # Boundary: exactly 1000 pages is allowed; 1001 is rejected.
+    assert parse_pages("1-1000") == list(range(1, 1001))
+    with pytest.raises(ValueError, match="too large"):
+        parse_pages("1-1001")
+    # A range that fits but whose accumulation across parts crosses the cap.
+    with pytest.raises(ValueError, match="too large"):
+        parse_pages("1-600,700-1400")
+
+
 def test_retrieve_get_pdf_page_content_falls_back_to_canonical(tmp_path, monkeypatch):
     """When no cached 'pages' are present, the file-read fallback must
     delegate to the canonical get_pdf_page_content instead of re-implementing

tests/test_sqlite_storage.py

@@ -139,6 +139,46 @@ def test_close_closes_connections_created_in_other_threads(storage):
         conns["worker"].execute("SELECT 1")
 
 
+def test_worker_reconnects_via_get_conn_after_close(storage):
+    """Regression: after close(), a thread that had already cached a connection
+    in thread-local storage would get that now-CLOSED handle back from
+    _get_conn (close() can only del its own thread-local), raising
+    ProgrammingError instead of transparently reconnecting. A generation bump
+    on close() must make the SAME thread's next _get_conn hand back a fresh,
+    working connection."""
+    import threading
+
+    storage.create_collection("papers")
+
+    cached = threading.Event()
+    closed = threading.Event()
+    result = {}
+
+    def worker():
+        # 1. cache a connection in this thread's thread-local
+        storage._get_conn().execute("SELECT 1")
+        cached.set()
+        # 2. wait until the main thread closed the storage (invalidating it)
+        closed.wait(timeout=5)
+        # 3. reuse from the SAME thread -> must reconnect, not reuse closed conn
+        try:
+            result["val"] = storage._get_conn().execute("SELECT 1").fetchone()[0]
+            result["list"] = storage.list_collections()
+        except Exception as e:  # noqa: BLE001 - record for assertion
+            result["err"] = f"{type(e).__name__}: {e}"
+
+    t = threading.Thread(target=worker)
+    t.start()
+    cached.wait(timeout=5)
+    storage.close()  # closes + invalidates the worker's cached connection
+    closed.set()
+    t.join(timeout=5)
+
+    assert "err" not in result, f"reconnect after close failed: {result.get('err')}"
+    assert result["val"] == 1
+    assert result["list"] == ["papers"]
+
+
 def test_duplicate_file_hash_in_collection_raises(storage):
     """UNIQUE(collection_name, file_hash) guards the add-same-file race."""
     import sqlite3