fix(index): bound LLM concurrency at the leaf, not per gather call

The previous approach (bounded_gather building a fresh semaphore per call)
did NOT compose: the indexing call graph nests gathers
(tree_parser -> process_large_node_recursively -> recurse, plus each node's
check_title gather), and each level got its own independent cap. Peak
in-flight LLM calls grew ~N^depth, so a deep/wide document still exhausted
file descriptors (Errno 24) — the exact failure the cap was meant to
prevent — while the flat summary phase was over-serialized. Naively sharing
one semaphore across gather levels would instead deadlock (a parent holds a
slot while awaiting children that need slots).

Move the throttle to the single chokepoint every LLM call funnels through,
llm_acompletion: one shared semaphore per event loop, acquired only around
the litellm.acompletion network call. This gives a true global cap that
composes across any nesting and can't deadlock (a parent awaiting children
holds no slot). bounded_gather is gone; the call sites revert to plain
asyncio.gather.

Also:
- Reject bool in max_concurrency validation (bool is an int subclass, so
  set_max_concurrency(True) / IndexConfig(max_concurrency=True) previously
  became Semaphore(1) and silently serialized). Shared _validate_max_concurrency
  + a pydantic field_validator.
- Guard check_title_appearance_in_start_concurrent against an out-of-range
  or 0 physical_index (LLM can emit one): it was dereferenced during task
  construction, outside the gather's return_exceptions protection, aborting
  the whole build; 0 silently wrapped to the last page. Now marked 'no'.
- Propagate contextvars into agent.py's worker-thread run (mirrors
  pipeline._run_async) so ContextVar settings stay consistent.
- Tests rewritten to cover the nested case the old flat tests missed, the
  leaf-level throttle in llm_acompletion, bool rejection, and the
  out-of-range physical_index guard.

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

pageindex/agent.py

@@ -159,6 +159,10 @@ class AgentRunner:
             result = Runner.run_sync(agent, question)
         else:
             import concurrent.futures
+            import contextvars
+            # Copy the current context into the worker thread so ContextVar-based
+            # settings propagate (mirrors pipeline._run_async).
+            ctx = contextvars.copy_context()
             with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
-                result = pool.submit(asyncio.run, Runner.run(agent, question)).result()
+                result = pool.submit(ctx.run, asyncio.run, Runner.run(agent, question)).result()
         return result.final_output

pageindex/config.py

@@ -5,7 +5,7 @@ import os
 from contextlib import contextmanager
 from contextvars import ContextVar
 
-from pydantic import BaseModel
+from pydantic import BaseModel, field_validator
 
 
 class IndexConfig(BaseModel):
@@ -30,6 +30,16 @@ class IndexConfig(BaseModel):
     # An explicit value here wins for this client.
     max_concurrency: int | None = None
 
+    @field_validator("max_concurrency", mode="before")
+    @classmethod
+    def _validate_max_concurrency_field(cls, v):
+        # Reject bool before pydantic coerces True->1 / False->0, and reject
+        # non-positive ints, so a bad value fails loudly instead of silently
+        # serializing (Semaphore(1)) or crashing (Semaphore(0)).
+        if v is not None:
+            _validate_max_concurrency(v)
+        return v
+
 
 def _env_drop_params_default() -> bool:
     return os.getenv("PAGEINDEX_DROP_PARAMS", "true").strip().lower() not in (
@@ -90,6 +100,17 @@ _MAX_CONCURRENCY_OVERRIDE: ContextVar[int | None] = ContextVar(
 )
 
 
+def _validate_max_concurrency(value) -> None:
+    """Raise ValueError unless ``value`` is a positive int.
+
+    ``bool`` is an ``int`` subclass, so it's rejected explicitly — otherwise
+    ``set_max_concurrency(True)`` would pass and become ``Semaphore(1)``,
+    silently serializing all indexing instead of failing loudly.
+    """
+    if isinstance(value, bool) or not isinstance(value, int) or value <= 0:
+        raise ValueError("max_concurrency must be a positive integer")
+
+
 def get_max_concurrency() -> int:
     """Return the effective cap on concurrent in-flight LLM calls during indexing.
 
@@ -103,8 +124,7 @@ def get_max_concurrency() -> int:
 def set_max_concurrency(value: int) -> None:
     """Set the process-wide default cap on concurrent in-flight LLM calls."""
     global _MAX_CONCURRENCY
-    if not isinstance(value, int) or value <= 0:
-        raise ValueError("max_concurrency must be a positive integer")
+    _validate_max_concurrency(value)
     _MAX_CONCURRENCY = value
 
 
@@ -117,8 +137,8 @@ def max_concurrency_scope(value: int | None):
     indexing doesn't leak across documents and a one-off value never becomes
     the sticky new default.
     """
-    if value is not None and (not isinstance(value, int) or value <= 0):
-        raise ValueError("max_concurrency must be a positive integer")
+    if value is not None:
+        _validate_max_concurrency(value)
     token = _MAX_CONCURRENCY_OVERRIDE.set(value)
     try:
         yield

pageindex/index/page_index.py

@@ -75,21 +75,28 @@ async def check_title_appearance_in_start_concurrent(structure, page_list, model
     if logger:
         logger.info("Checking title appearance in start concurrently")
     
-    # skip items without physical_index
+    # Mark items we can't check as 'no' up front: missing physical_index, or one
+    # out of range for page_list. An out-of-range index (the LLM can emit one)
+    # would otherwise raise IndexError below — during task-list construction,
+    # outside the gather's return_exceptions protection — and abort the build.
+    def _valid_physical_index(item):
+        idx = item.get('physical_index')
+        return idx is not None and 1 <= idx <= len(page_list)
+
     for item in structure:
-        if item.get('physical_index') is None:
+        if not _valid_physical_index(item):
             item['appear_start'] = 'no'
 
-    # only for items with valid physical_index
+    # only for items with a valid, in-range physical_index
     tasks = []
     valid_items = []
     for item in structure:
-        if item.get('physical_index') is not None:
+        if _valid_physical_index(item):
             page_text = page_list[item['physical_index'] - 1][0]
             tasks.append(check_title_appearance_in_start(item['title'], page_text, model=model, logger=logger))
             valid_items.append(item)
 
-    results = await bounded_gather(tasks, return_exceptions=True)
+    results = await asyncio.gather(*tasks, return_exceptions=True)
     for item, result in zip(valid_items, results):
         if isinstance(result, Exception):
             if logger:
@@ -832,7 +839,7 @@ async def fix_incorrect_toc(toc_with_page_number, page_list, incorrect_results,
         process_and_check_item(item)
         for item in incorrect_results
     ]
-    results = await bounded_gather(tasks, return_exceptions=True)
+    results = await asyncio.gather(*tasks, return_exceptions=True)
     for item, result in zip(incorrect_results, results):
         if isinstance(result, Exception):
             print(f"Processing item {item} generated an exception: {result}")
@@ -927,7 +934,7 @@ async def verify_toc(page_list, list_result, start_index=1, N=None, model=None):
         check_title_appearance(item, page_list, start_index, model)
         for item in indexed_sample_list
     ]
-    results = await bounded_gather(tasks)
+    results = await asyncio.gather(*tasks)
     
     # Process results
     correct_count = 0
@@ -1015,7 +1022,7 @@ async def process_large_node_recursively(node, page_list, opt=None, logger=None)
             process_large_node_recursively(child_node, page_list, opt, logger=logger)
             for child_node in node['nodes']
         ]
-        await bounded_gather(tasks)
+        await asyncio.gather(*tasks)
     
     return node
 
@@ -1051,7 +1058,7 @@ async def tree_parser(page_list, opt, doc=None, logger=None):
         process_large_node_recursively(node, page_list, opt, logger=logger)
         for node in toc_tree
     ]
-    await bounded_gather(tasks)
+    await asyncio.gather(*tasks)
     
     return toc_tree
 

pageindex/index/page_index_md.py

@@ -16,7 +16,7 @@ async def get_node_summary(node, summary_token_threshold=200, model=None):
 async def generate_summaries_for_structure_md(structure, summary_token_threshold, model=None):
     nodes = structure_to_list(structure)
     tasks = [get_node_summary(node, summary_token_threshold=summary_token_threshold, model=model) for node in nodes]
-    summaries = await bounded_gather(tasks)
+    summaries = await asyncio.gather(*tasks)
     
     for node, summary in zip(nodes, summaries):
         if not node.get('nodes'):

pageindex/index/utils.py

@@ -7,6 +7,8 @@ import json
 import copy
 import re
 import asyncio
+import threading
+import weakref
 import PyPDF2
 import pymupdf
 import yaml
@@ -25,28 +27,41 @@ from ..tokens import count_tokens  # re-exported for backward compat
 logger = logging.getLogger(__name__)
 
 
-async def bounded_gather(coros, *, return_exceptions=False):
-    """``asyncio.gather`` with a cap on how many coroutines run concurrently.
+# One shared semaphore per event loop, bounding concurrent in-flight LLM calls.
+# Keyed by the loop object (WeakKeyDictionary drops the entry once the loop is
+# closed and garbage-collected) so each asyncio.run() gets its own, correctly
+# loop-bound semaphore. The lock only guards the tiny get-or-create against two
+# threads (each driving its own loop) racing to insert; within a single loop
+# everything is single-threaded, so no lock is needed on the hot path.
+_LLM_SEMAPHORES: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
+_LLM_SEMAPHORES_LOCK = threading.Lock()
 
-    Each coroutine acquires a shared semaphore before running, so no more than
-    ``get_max_concurrency()`` LLM calls are ever in flight at once. Without this
-    a many-node document schedules every node's LLM call simultaneously, opening
-    one socket per node and exhausting the process file-descriptor limit
-    (Errno 24, "Too many open files").
 
-    The semaphore is created inside the running loop, so this stays correct when
-    the caller drives each document in its own ``asyncio.run()`` loop. Order of
-    results matches input order, mirroring ``asyncio.gather``.
+def _llm_semaphore() -> asyncio.Semaphore:
+    """Shared per-loop cap on concurrent in-flight LLM calls.
+
+    Acquired only around the leaf ``litellm.acompletion`` call in
+    ``llm_acompletion`` — the single point every LLM request funnels through —
+    so the cap is a TRUE global bound no matter how deeply the indexing gathers
+    nest (``tree_parser`` → ``process_large_node_recursively`` → …). Bounding at
+    the leaf rather than at each gather call site is also deadlock-free: a parent
+    coroutine awaiting its children holds no slot, so children can always
+    acquire one.
+
+    Sized from ``get_max_concurrency()`` the first time it's needed in a loop, so
+    a per-index ``max_concurrency_scope`` override in effect at that moment is
+    honored. Without this bound a many-node document opens one socket per node at
+    once and exhausts the process file-descriptor limit (Errno 24).
     """
-    semaphore = asyncio.Semaphore(get_max_concurrency())
-
-    async def _run(coro):
-        async with semaphore:
-            return await coro
-
-    return await asyncio.gather(
-        *(_run(c) for c in coros), return_exceptions=return_exceptions
-    )
+    loop = asyncio.get_running_loop()
+    sem = _LLM_SEMAPHORES.get(loop)
+    if sem is None:
+        with _LLM_SEMAPHORES_LOCK:
+            sem = _LLM_SEMAPHORES.get(loop)
+            if sem is None:
+                sem = asyncio.Semaphore(get_max_concurrency())
+                _LLM_SEMAPHORES[loop] = sem
+    return sem
 
 
 def llm_completion(model, prompt, chat_history=None, return_finish_reason=False):
@@ -86,11 +101,14 @@ async def llm_acompletion(model, prompt):
     messages = [{"role": "user", "content": prompt}]
     for i in range(max_retries):
         try:
-            response = await litellm.acompletion(
-                model=model,
-                messages=messages,
-                **get_llm_params(),  # per-call kwargs; never the litellm global
-            )
+            # Hold a concurrency slot only around the actual network call — not
+            # across retry backoff — so the cap counts real in-flight requests.
+            async with _llm_semaphore():
+                response = await litellm.acompletion(
+                    model=model,
+                    messages=messages,
+                    **get_llm_params(),  # per-call kwargs; never the litellm global
+                )
             return response.choices[0].message.content
         except Exception as e:
             logger.warning("Retrying async LLM completion (%d/%d)", i + 1, max_retries)
@@ -240,7 +258,7 @@ async def generate_node_summary(node, model=None):
 async def generate_summaries_for_structure(structure, model=None):
     nodes = structure_to_list(structure)
     tasks = [generate_node_summary(node, model=model) for node in nodes]
-    summaries = await bounded_gather(tasks)
+    summaries = await asyncio.gather(*tasks)
 
     for node, summary in zip(nodes, summaries):
         node['summary'] = summary