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

@ -29,6 +29,11 @@ class IndexConfig(BaseModel):
# default (get_max_concurrency(), overridable via PAGEINDEX_MAX_CONCURRENCY).
# An explicit value here wins for this client.
max_concurrency: int | None = None
# Per-call litellm completion kwargs for this client's indexing calls only
# (e.g. {"temperature": 1}). None = use the process-wide defaults
# (get_llm_params(), overridable via set_llm_params()). Scoped via
# llm_params_scope so it doesn't leak into other concurrent indexing calls.
llm_params: dict | None = None
@field_validator("max_concurrency", mode="before")
@classmethod
@ -57,6 +62,16 @@ def _env_drop_params_default() -> bool:
# PAGEINDEX_DROP_PARAMS env shortcut.
_LLM_PARAMS: dict = {"temperature": 0, "drop_params": _env_drop_params_default()}
# Per-call override, isolated per thread / async context — mirrors
# _MAX_CONCURRENCY_OVERRIDE below. Without this, set_llm_params() is the only
# way to change llm params and it mutates the process-wide dict directly, so
# concurrently indexing two documents with different llm_params_scope() would
# otherwise leak one caller's settings (e.g. temperature) into the other's
# in-flight calls. None = no override -> fall back to the process-wide _LLM_PARAMS.
_LLM_PARAMS_OVERRIDE: ContextVar[dict | None] = ContextVar(
"pageindex_llm_params_override", default=None
)
# Structural kwargs PageIndex always supplies itself — not overridable here.
_RESERVED_LLM_PARAMS = ("model", "messages")
@ -121,6 +136,16 @@ def get_max_concurrency() -> int:
return override if override is not None else _MAX_CONCURRENCY
def _process_wide_max_concurrency() -> int:
"""The process-wide default cap, ignoring any active max_concurrency_scope
override. This is the TRUE ceiling shared across every thread/event loop in
the process (see index/utils.py's _llm_semaphore) — a per-call override may
only narrow the effective cap within that ceiling, never widen it, so the
ceiling itself must not vary with a context-local override.
"""
return _MAX_CONCURRENCY
def set_max_concurrency(value: int) -> None:
"""Set the process-wide default cap on concurrent in-flight LLM calls."""
global _MAX_CONCURRENCY
@ -147,19 +172,53 @@ def max_concurrency_scope(value: int | None):
def get_llm_params() -> dict:
"""Return a copy of the per-call kwargs PageIndex passes to litellm."""
return dict(_LLM_PARAMS)
"""Return a copy of the effective per-call kwargs PageIndex passes to litellm.
A per-index override (llm_params_scope) is merged over the process-wide
defaults for the current context; otherwise just the process-wide defaults
apply.
"""
params = dict(_LLM_PARAMS)
override = _LLM_PARAMS_OVERRIDE.get()
if override:
params.update(override)
return params
def set_llm_params(**kwargs) -> None:
"""Override or extend the litellm completion kwargs PageIndex sends per call.
"""Override or extend the process-wide default litellm completion kwargs.
e.g. ``set_llm_params(drop_params=False, temperature=1, num_retries=5)``.
Applied per call; never writes litellm's global state, so it can't leak into
other litellm users in the same process. ``model`` / ``messages`` are
reserved (PageIndex supplies them) and rejected.
Never writes litellm's global state, so it can't leak into other litellm
users in the same process but it DOES mutate PageIndex's own process-wide
default, so it affects every concurrent caller in this process. For a
one-off override scoped to a single indexing call, use ``llm_params_scope``
instead. ``model`` / ``messages`` are reserved (PageIndex supplies them) and
rejected.
"""
reserved = [k for k in kwargs if k in _RESERVED_LLM_PARAMS]
if reserved:
raise ValueError(f"cannot override reserved litellm kwargs: {reserved}")
_LLM_PARAMS.update(kwargs)
@contextmanager
def llm_params_scope(overrides: dict | None):
"""Scope a per-index override of the litellm completion kwargs to the
current context.
``overrides=None`` (or ``{}``) means "no override" (fall back to the
process-wide defaults). Isolated per thread / async context and reset on
exit, so concurrent indexing doesn't leak one call's kwargs into another's
and a one-off override never becomes the sticky new process default
mirrors ``max_concurrency_scope``.
"""
if overrides:
reserved = [k for k in overrides if k in _RESERVED_LLM_PARAMS]
if reserved:
raise ValueError(f"cannot override reserved litellm kwargs: {reserved}")
token = _LLM_PARAMS_OVERRIDE.set(overrides or None)
try:
yield
finally:
_LLM_PARAMS_OVERRIDE.reset(token)