fix: prompt-injection delimiter escape, legacy config coercion, gather resilience, true cross-thread concurrency bound

Addresses items 4-8 from the max-effort review of PR #272 (VectifyAI/PageIndex#272).

- agent.py: wrap_with_doc_context() strips '<'/'>' from doc_name/doc_description
  (untrusted: unsanitized filename / LLM-generated from document content) before
  inserting them into the <docs>...</docs> block, so embedded content can never
  form a literal </docs> that closes the delimiter early and escapes the
  untrusted-data boundary SCOPED_SYSTEM_PROMPT relies on. Deterministic
  per-field transform, doesn't touch the (cacheable) static system prompt.

- ConfigLoader.load() (legacy 0.2.x compat) now routes merged overrides through
  IndexConfig before returning, so a legacy 'no' string gets pydantic's bool
  coercion instead of surviving as a truthy non-empty string — page_index_main's
  bare `if opt.if_add_node_summary:` checks (changed from `== 'yes'` elsewhere
  in this PR) were silently inverting caller intent and firing unwanted billed
  LLM calls.

- verify_toc, process_large_node_recursively, tree_parser,
  generate_summaries_for_structure, generate_summaries_for_structure_md: added
  return_exceptions=True to their asyncio.gather calls (llm_completion/
  llm_acompletion raise RuntimeError on retry exhaustion, added earlier in this
  PR), each with a degrade path matching the pattern already used by sibling
  hardened gathers in the same files. One transient LLM failure no longer
  aborts the whole document's indexing.

- _llm_semaphore is now a true process-wide ceiling (threading.Semaphore,
  shared across every thread/event loop) instead of one asyncio.Semaphore per
  event loop -- concurrently indexing N documents on N threads no longer
  multiplies the effective cap by N. A max_concurrency_scope() override is
  layered as a second, nested, per-loop restriction that can only tighten the
  effective cap within the ceiling, never widen past it.

- set_llm_params() mutated a bare process-wide dict with no per-call isolation,
  unlike max_concurrency which already had ContextVar scoping. Added
  llm_params_scope() (mirrors max_concurrency_scope) + IndexConfig.llm_params,
  wired into build_index() the same way max_concurrency already was, so
  concurrent indexing jobs with different llm kwargs don't leak into each
  other.

Adds regression tests for all five. Full suite: 221 passed, 2 skipped (one
pre-existing, unrelated flaky cloud-streaming test intermittently fails on
rerun; confirmed independent of this change).

Claude-Session: https://claude.ai/code/session_01Kx5DgKbhK1N8autqXH8SmS
This commit is contained in:
mountain 2026-07-09 11:15:12 +08:00
parent 623ce927f1
commit b9d021916f
11 changed files with 408 additions and 55 deletions

View file

@ -929,13 +929,23 @@ async def verify_toc(page_list, list_result, start_index=1, N=None, model=None):
item_with_index['list_index'] = idx # Add the original index in list_result
indexed_sample_list.append(item_with_index)
# Run checks concurrently
# Run checks concurrently. return_exceptions=True: a transient LLM failure
# on one sampled item must degrade that item to 'no' (same as an
# unavailable physical_index above), not abort verification for the
# whole document.
tasks = [
check_title_appearance(item, page_list, start_index, model)
for item in indexed_sample_list
]
results = await asyncio.gather(*tasks)
raw_results = await asyncio.gather(*tasks, return_exceptions=True)
results = []
for item, result in zip(indexed_sample_list, raw_results):
if isinstance(result, Exception):
results.append({'list_index': item.get('list_index'), 'answer': 'no',
'title': item.get('title'), 'page_number': item.get('physical_index')})
else:
results.append(result)
# Process results
correct_count = 0
incorrect_results = []
@ -1022,8 +1032,14 @@ 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 asyncio.gather(*tasks)
# return_exceptions=True: one child subtree failing to expand further
# must not abort the whole document — it's left as a leaf at its
# current boundaries instead.
results = await asyncio.gather(*tasks, return_exceptions=True)
for child_node, result in zip(node['nodes'], results):
if isinstance(result, Exception) and logger:
logger.error(f"Failed to expand node '{child_node.get('title')}': {result}")
return node
async def tree_parser(page_list, opt, doc=None, logger=None):
@ -1058,8 +1074,13 @@ 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 asyncio.gather(*tasks)
# return_exceptions=True: one top-level node failing to expand further
# must not abort indexing the whole document.
results = await asyncio.gather(*tasks, return_exceptions=True)
for node, result in zip(toc_tree, results):
if isinstance(result, Exception) and logger:
logger.error(f"Failed to expand node '{node.get('title')}': {result}")
return toc_tree

View file

@ -16,8 +16,14 @@ 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 asyncio.gather(*tasks)
# return_exceptions=True: one node's summary failing must not abort
# summarization for the whole document — fall back to its raw text.
raw_summaries = await asyncio.gather(*tasks, return_exceptions=True)
summaries = [
node.get('text', '') if isinstance(s, Exception) else s
for node, s in zip(nodes, raw_summaries)
]
for node, summary in zip(nodes, summaries):
if not node.get('nodes'):
node['summary'] = summary

View file

@ -68,15 +68,16 @@ def build_index(parsed: ParsedDocument, model: str = None, opt=None) -> dict:
from .utils import (write_node_id, add_node_text, remove_structure_text,
generate_summaries_for_structure, generate_doc_description,
create_clean_structure_for_description)
from ..config import IndexConfig, max_concurrency_scope
from ..config import IndexConfig, max_concurrency_scope, llm_params_scope
if opt is None:
opt = IndexConfig(model=model) if model else IndexConfig()
# Scope the per-index concurrency cap to THIS call only (per thread/async
# context), so concurrent indexing of other documents isn't affected and a
# one-off value never sticks as the process default.
with max_concurrency_scope(getattr(opt, "max_concurrency", None)):
# Scope the per-index concurrency cap AND llm kwargs to THIS call only (per
# thread/async context), so concurrent indexing of other documents isn't
# affected and a one-off value never sticks as the process default.
with max_concurrency_scope(getattr(opt, "max_concurrency", None)), \
llm_params_scope(getattr(opt, "llm_params", None)):
nodes = parsed.nodes
strategy = detect_strategy(nodes)

View file

@ -21,47 +21,103 @@ from pprint import pprint
# `pageindex.config` submodule for those modules.
from types import SimpleNamespace as _config
from ..config import get_llm_params, get_max_concurrency
from contextlib import asynccontextmanager
from ..config import get_llm_params, get_max_concurrency, _process_wide_max_concurrency
from ..tokens import count_tokens # re-exported for backward compat
logger = logging.getLogger(__name__)
# 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()
# TRUE process-wide ceiling on concurrent in-flight LLM calls, shared across
# EVERY thread and event loop (a plain threading.Semaphore, not an
# asyncio.Semaphore — those are bound to the loop that created them, so one per
# loop would let N concurrently-indexing threads each get their own full-size
# cap and multiply the effective bound by N). Resized lazily when the
# process-wide default changes; resizing isn't perfectly atomic against
# in-flight acquires, which is fine since it only happens on an explicit
# set_max_concurrency() config change, not on the hot path.
_PROCESS_LLM_SEMAPHORE: threading.Semaphore | None = None
_PROCESS_LLM_SEMAPHORE_SIZE: int | None = None
_PROCESS_LLM_SEMAPHORE_LOCK = threading.Lock()
# Per-loop, per-size semaphores for a max_concurrency_scope() override that's
# narrower than the process ceiling — isolates one call's own subtree to a
# tighter self-imposed limit without needing to be cross-thread itself (it can
# never let MORE calls through than the process ceiling above already allows,
# since both are held simultaneously; see _llm_semaphore). Keyed by (loop, size)
# rather than just loop so a later scope with a different size in the same loop
# isn't silently ignored.
_SCOPED_LLM_SEMAPHORES: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
_SCOPED_LLM_SEMAPHORES_LOCK = threading.Lock()
def _llm_semaphore() -> asyncio.Semaphore:
"""Shared per-loop cap on concurrent in-flight LLM calls.
def _process_ceiling_semaphore() -> threading.Semaphore:
global _PROCESS_LLM_SEMAPHORE, _PROCESS_LLM_SEMAPHORE_SIZE
size = _process_wide_max_concurrency()
with _PROCESS_LLM_SEMAPHORE_LOCK:
if _PROCESS_LLM_SEMAPHORE is None or _PROCESS_LLM_SEMAPHORE_SIZE != size:
_PROCESS_LLM_SEMAPHORE = threading.Semaphore(size)
_PROCESS_LLM_SEMAPHORE_SIZE = size
return _PROCESS_LLM_SEMAPHORE
def _scoped_llm_semaphore(size: int) -> asyncio.Semaphore:
loop = asyncio.get_running_loop()
per_loop = _SCOPED_LLM_SEMAPHORES.get(loop)
if per_loop is None:
with _SCOPED_LLM_SEMAPHORES_LOCK:
per_loop = _SCOPED_LLM_SEMAPHORES.get(loop)
if per_loop is None:
per_loop = {}
_SCOPED_LLM_SEMAPHORES[loop] = per_loop
sem = per_loop.get(size)
if sem is None:
with _SCOPED_LLM_SEMAPHORES_LOCK:
sem = per_loop.get(size)
if sem is None:
sem = asyncio.Semaphore(size)
per_loop[size] = sem
return sem
@asynccontextmanager
async def _llm_semaphore():
"""Bound concurrent in-flight LLM calls to a TRUE process-wide ceiling,
optionally narrowed further by an active max_concurrency_scope() override.
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.
so the cap holds no matter how deeply the indexing gathers nest
(``tree_parser`` ``process_large_node_recursively`` ) AND no matter how
many threads are each running their own indexing job concurrently. 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).
The process ceiling (threading.Semaphore, shared cross-thread) is sized from
the process-wide default only; a narrower max_concurrency_scope() override
is enforced as a second, nested, per-loop restriction it can only
*tighten* the effective cap for its own call tree, never widen it past the
ceiling. Without the outer bound a many-node document opens one socket per
node at once and exhausts the process file-descriptor limit (Errno 24).
"""
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
ceiling_sem = _process_ceiling_semaphore()
# threading.Semaphore.acquire() blocks the calling thread, so run it off
# the event loop thread — otherwise it would freeze every other coroutine
# on this loop while waiting for a slot. release() is non-blocking and
# safe to call directly from any thread.
await asyncio.to_thread(ceiling_sem.acquire)
try:
effective = get_max_concurrency()
ceiling = _process_wide_max_concurrency()
if effective < ceiling:
async with _scoped_llm_semaphore(effective):
yield
else:
yield
finally:
ceiling_sem.release()
def llm_completion(model, prompt, chat_history=None, return_finish_reason=False):
@ -258,7 +314,14 @@ 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 asyncio.gather(*tasks)
# return_exceptions=True: one node's summary failing (e.g. a transient LLM
# error) must not abort summarization for the whole document — fall back
# to the node's own raw text so retrieval still has something usable.
raw_summaries = await asyncio.gather(*tasks, return_exceptions=True)
summaries = [
node.get('text', '') if isinstance(s, Exception) else s
for node, s in zip(nodes, raw_summaries)
]
for node, summary in zip(nodes, summaries):
node['summary'] = summary
@ -829,7 +892,13 @@ class ConfigLoader:
self._validate_keys(user_dict)
merged = {**self._default_dict, **user_dict}
return _config(**merged)
# Route through IndexConfig so legacy 'yes'/'no' string overrides get
# pydantic's bool coercion (a bare 'no' is otherwise a truthy string —
# page_index_main's `if opt.if_add_node_summary:` checks would silently
# invert the caller's intent).
from ..config import IndexConfig
validated = IndexConfig(**merged)
return _config(**validated.model_dump())
def create_node_mapping(tree, include_page_ranges=False, max_page=None):