fix: Aug 18-19 review hardening — five max-effort rounds across tools, chat, and flash (#413)

Squash-port of feat/local-chat's post-dev5 wave (03ffab3..49a24e1): five
review rounds of reproduced-then-fixed findings. Highlights: standalone
agent_instructions back on the strict shadow check (the *_agent_config
bundles keep the relaxed in-set check they can prove); caller-owned
http_client survives the per-call backend closes; get_tree keeps
key_items under the flash merge default; OpenAI-protocol classification
follows litellm's own routing (azure/openrouter/deepseek/groq/xai);
thinking-aware max_tokens default shared by messages() and
anthropic_runner_config(thinking=); 401/403 re-raise instead of
retry-coaching envelopes, bridge errors carry status_code; cloud
discovery + instructions ride the ?tools=read endpoint matching the
tool gate (live-verified); the chat lane's litellm model resolution
deduplicated into utils._litellm_model; explicit optimize= wins over
the deprecated optimize_expand (DeprecationWarning added); mcp 2.0
compatibility; spawn-worker and doc-scope guard fixes; the suite is
.env-independent and pins the LitellmModel._fetch_response seam.

Per-finding rationale in the ported commit messages on feat/local-chat.

.github/workflows/tests.yml

@@ -30,6 +30,9 @@ jobs:
           python-version: ${{ matrix.python-version }}
           cache: pip
       - run: pip install -r requirements.txt pytest
+      - if: matrix.agent-frameworks == 'without'
+        # requirements.txt carries it; this leg tests the no-framework paths
+        run: pip uninstall -y openai-agents
       - if: matrix.agent-frameworks == 'with'
         run: pip install openai-agents claude-agent-sdk anthropic
       - run: python -m pytest -q

.gitignore

@@ -6,4 +6,3 @@ __pycache__
 logs/
 .pageindex/
 dist/
-*.doc_id

README.md

@@ -173,7 +173,7 @@ python3 run_pageindex.py --pdf_path /path/to/your/document.pdf
 <details>
 <summary>Optional parameters</summary>
 <br>
-You can customize the processing with additional optional arguments (the structure-tuning flags below require <code>--mode standard</code>):
+You can customize the processing with additional optional arguments (the structure-tuning flags from <code>--toc-check-pages</code> down require <code>--mode standard</code>):
 
 ```
 --mode                  Processing mode: flash (default) or standard

examples/agentic_vectorless_rag_demo.py

@@ -32,14 +32,13 @@ from agents import Agent, Runner, set_tracing_disabled
 from agents.stream_events import RawResponsesStreamEvent, RunItemStreamEvent
 from openai.types.responses import ResponseTextDeltaEvent, ResponseReasoningSummaryTextDeltaEvent
 
-from pageindex import PageIndexAPIError, PageIndexLocalClient
+from pageindex import PageIndexLocalClient
 import pageindex.utils as utils
 
 PDF_URL = "https://arxiv.org/pdf/2603.15031"
 
 _EXAMPLES_DIR = Path(__file__).parent
 PDF_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.pdf"
-DOC_ID_PATH = _EXAMPLES_DIR / "documents" / "attention-residuals.doc_id"
 STORAGE_PATH = _EXAMPLES_DIR / ".pageindex"
 
 
@@ -127,27 +126,13 @@ if __name__ == "__main__":
     print("=" * 60)
     print("Step 1: Index PDF and view tree structure")
     print("=" * 60)
-    doc_id = None
-    if DOC_ID_PATH.exists():
-        cached = DOC_ID_PATH.read_text().strip()
-        try:
-            client.get_document(cached)
-            doc_id = cached
-        except PageIndexAPIError:
-            DOC_ID_PATH.unlink()
-    if doc_id is None:
-        # The .doc_id cache is gitignored — on a fresh clone with an
-        # existing store, find the already-indexed copy by name instead of
-        # re-indexing it.
-        doc_id = next(
-            (doc["id"] for doc in client.list_documents(limit=100)["documents"]
-             if doc["name"] == PDF_PATH.name), None)
+    doc_id = next(
+        (doc["id"] for doc in client.list_documents(limit=100)["documents"]
+         if doc["name"] == PDF_PATH.name), None)
     if doc_id:
-        DOC_ID_PATH.write_text(doc_id)
         print(f"\nLoaded cached doc_id: {doc_id}")
     else:
         doc_id = client.submit_document(str(PDF_PATH), wait=True)["doc_id"]
-        DOC_ID_PATH.write_text(doc_id)
         print(f"\nIndexed. doc_id: {doc_id}")
     print("\nTree Structure (top-level sections):")
     structure = client.get_tree(doc_id, node_summary=True)["result"]

pageindex/__init__.py

@@ -1,11 +1,6 @@
 """PageIndex SDK."""
-import os as _os
 from typing import TYPE_CHECKING as _TYPE_CHECKING
 
-# LiteLLM's import otherwise fetches its model map over the network — seconds
-# of blocking (or a hang offline). setdefault, so an explicit user choice wins.
-_os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True")
-
 from .client import PageIndexClient, PageIndexCloudClient, PageIndexLocalClient
 from .errors import PageIndexAPIError
 
@@ -23,8 +18,6 @@ __all__ = [
 ]
 
 _LAZY = {
-    "page_index": ".page_index_classic",
-    "page_index_main": ".page_index_classic",
     "page_index_flash": ".flash",
     "optimize_tree": ".tree_optimize",
     "md_to_tree": ".page_index_md",
@@ -34,25 +27,18 @@ _SUBMODULES = {"agent_tools", "client", "cloud_api", "errors", "flash",
                "mcp_bridge", "page_index_classic", "page_index_md",
                "tree_optimize", "utils"}
 
+
 def __getattr__(name):
     if name.startswith("_"):
-        # Dunder probes (copy, pickle, inspect) are the frequent unknown
-        # names — they must not trigger the classic import below.
         raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
     import importlib
     if name in _SUBMODULES:
         return importlib.import_module(f".{name}", __name__)
-    # Pre-0.2.10 compat: unknown names fall through to the classic module,
-    # whose public surface (ConfigLoader, count_tokens, ...) resolved as
-    # package attributes. A non-underscore typo pays one classic import
-    # before its AttributeError — not worth an allowlist.
-    module = importlib.import_module(_LAZY.get(name, ".page_index_classic"),
-                                     __name__)
+    module = importlib.import_module(_LAZY.get(name, ".page_index_classic"), __name__)
     try:
         value = getattr(module, name)
     except AttributeError:
-        raise AttributeError(
-            f"module {__name__!r} has no attribute {name!r}") from None
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
     globals()[name] = value
     return value
 

pageindex/agent_tools.py

@@ -9,11 +9,10 @@ guidance strings (tool descriptions) adapt to the local surface the same
 way the agent instructions do — they never teach capabilities that only
 exist on the cloud.
 
-Tools never raise for any invocation their signatures accept: every
-outcome, including errors, is returned as the same JSON envelope the cloud
-emits ({"success": true, ...} / {"error": ...}). Arguments outside a pruned
-local signature fail at the Python call boundary; the call_tool path
-answers them with the guided error envelope instead.
+Tools never raise: every outcome, including errors, is returned as the
+same JSON envelope the cloud emits ({"success": true, ...} /
+{"error": ...}) — arguments outside a pruned local signature come back as
+that envelope too, on the direct and the call_tool path alike.
 """
 from __future__ import annotations
 
@@ -33,7 +32,6 @@ TOOL_RESPONSE_CHAR_LIMIT = 100_000
 STRUCTURE_FIRST_PAGE_THRESHOLD = 20
 
 _CHAR_BUDGET = int(TOOL_RESPONSE_CHAR_LIMIT * 0.95)
-_PAGES_SPEC_RE = re.compile(r"^(\d+(-\d+)?)(,\s*\d+(-\d+)?)*$")
 _MAX_REQUESTED_PAGES = 10_000
 _SIMILAR_NAMES_LIMIT = 3
 _TOOL_WAIT_TIMEOUT = 180.0  # "up to 3 minutes", per the wait_for_completion schema
@@ -328,15 +326,23 @@ def _dumps(payload: dict[str, Any]) -> str:
 
 # ── document listing / name resolution ──
 
-def _all_documents(client) -> list[dict[str, Any]]:
+def _all_documents(client, stop_ids=None) -> list[dict[str, Any]]:
     """Every document the client can list, newest first (both modes list
-    newest-first; paging preserves that order)."""
+    newest-first; paging preserves that order). With ``stop_ids``, paging
+    stops early once every one of those ids has been seen — for callers
+    that only need those entries; an id absent from the listing still
+    costs a full sweep."""
     documents: list[dict[str, Any]] = []
     offset = 0
+    remaining = {str(one_id) for one_id in stop_ids} if stop_ids else None
     while True:
         page = client.list_documents(limit=100, offset=offset)
         batch = page.get("documents") or []
         documents.extend(batch)
+        if remaining is not None:
+            remaining.difference_update(str(doc.get("id")) for doc in batch)
+            if not remaining:
+                return documents
         # Advance by what actually arrived — stepping by the requested
         # limit skips documents whenever a server caps its page size.
         offset += len(batch)
@@ -387,14 +393,14 @@ def _resolve_document(
     """Resolve doc_name to a list entry. Same-name duplicates resolve to the
     newest match. Returns (entry, None) or (None, error_payload_pair)."""
     if documents is None:
-        documents = _all_documents(client)
+        documents = _all_documents(client, stop_ids=allowed_ids)
     documents = _scope_documents(documents, allowed_ids)
     matches = [doc for doc in documents if doc.get("name") == doc_name]
     if matches:
         return max(matches, key=lambda d: d.get("createdAt") or ""), None
     names = [str(doc.get("name")) for doc in documents if doc.get("name")]
-    similar = difflib.get_close_matches(doc_name, names, n=_SIMILAR_NAMES_LIMIT,
-                                        cutoff=0.5)
+    similar = difflib.get_close_matches(str(doc_name), names,
+                                        n=_SIMILAR_NAMES_LIMIT, cutoff=0.5)
     message = (
         "Document not found. Did you mean: "
         + ", ".join(f'"{name}"' for name in similar) + "?"
@@ -433,7 +439,7 @@ def _await_completion(client, entry: dict[str, Any], wait: bool) -> dict[str, An
         time.sleep(_TOOL_WAIT_INTERVAL)
         refreshed = _refetch_entry(client, doc_id)
         if refreshed is None:
-            return current
+            continue  # transient refetch failure: poll on to the deadline
         if refreshed.get("metadata") is None:
             # Status refetches omit (or null out) custom metadata; keep the
             # listing's copy.
@@ -503,69 +509,102 @@ def _folder_unsupported(param: str) -> tuple[dict, bool]:
 
 # ── page spec handling ──
 
-def _parse_page_spec(
-    pages: str, doc_name: str,
-) -> "tuple[Optional[list[int]], Optional[_ToolResult]]":
-    """Expand '1-3,7' into a sorted, deduplicated page list, or an error."""
-    invalid = _failure(
-        "Invalid page specification format",
-        {"doc_name": doc_name},
-        {
-            "summary": "Failed to parse the pages parameter",
-            "options": [
-                'Use valid formats: "5", "3,7,10", "5-10", or "1-3,7,9-12"',
-                "Ensure page numbers are positive integers",
-            ],
-        },
-        "INVALID_INPUT",
-    )
-    if not isinstance(pages, str) or not _PAGES_SPEC_RE.match(pages.strip()):
-        return None, invalid
-    too_many = _failure(
-        f"Too many pages requested (over {_MAX_REQUESTED_PAGES})",
-        {"doc_name": doc_name},
-        {
-            "summary": "The page specification spans too many pages",
-            "options": [
-                "Request a narrower page range",
-                "The response holds only a few pages per call - page through with several smaller requests",
-            ],
-        },
-        "INVALID_INPUT",
-    )
+class _PageSpecError(ValueError):
+    """Shared page-spec rejection; ``code`` picks the caller's rendering."""
+
+    def __init__(self, code: str, message: str):
+        super().__init__(message)
+        self.code = code
+
+
+def _expand_pages(pages) -> list[int]:
+    """Expand '1-3,7' into sorted distinct pages — the one parser for the
+    SDK surface and the tool layer. Raises _PageSpecError with code
+    'invalid', 'too_many', or 'nonpositive'."""
+    if not isinstance(pages, str):
+        raise _PageSpecError("invalid",
+                             f"Invalid page specification: {pages!r}")
+    too_many = (f"Page specification '{pages}' spans more than "
+                f"{_MAX_REQUESTED_PAGES} pages; request a narrower range")
     expanded: set[int] = set()
     for part in pages.split(","):
         part = part.strip()
-        if "-" in part:
-            start, end = (int(x) for x in part.split("-", 1))
-            if start > end:
-                return None, invalid
-        else:
-            start = end = int(part)
+        try:
+            if "-" in part:
+                start, end = (int(x) for x in part.split("-", 1))
+                if start > end:
+                    raise _PageSpecError(
+                        "invalid",
+                        f"Invalid range '{part}': start must be <= end")
+            else:
+                start = end = int(part)
+        except _PageSpecError:
+            raise
+        except ValueError as exc:
+            raise _PageSpecError(
+                "invalid", f"Invalid page specification '{pages}'") from exc
         # Bound each part arithmetically before materializing it: a spec like
         # "1-1000000000" would otherwise expand to billions of integers
         # inside the caller's process. The cap is on distinct pages, so
         # overlapping parts (a parent section plus its children) don't
         # double-count.
         if end - start + 1 > _MAX_REQUESTED_PAGES:
-            return None, too_many
+            raise _PageSpecError("too_many", too_many)
         expanded.update(range(start, end + 1))
         if len(expanded) > _MAX_REQUESTED_PAGES:
-            return None, too_many
-    if any(page < 1 for page in expanded):
+            raise _PageSpecError("too_many", too_many)
+    if min(expanded) < 1:
+        raise _PageSpecError(
+            "nonpositive",
+            "Invalid page numbers. Page numbers must be positive integers")
+    return sorted(expanded)
+
+
+def _parse_page_spec(
+    pages: str, doc_name: str,
+) -> "tuple[Optional[list[int]], Optional[_ToolResult]]":
+    """Expand '1-3,7' into a sorted, deduplicated page list, or an error."""
+    try:
+        return _expand_pages(pages), None
+    except _PageSpecError as exc:
+        if exc.code == "too_many":
+            return None, _failure(
+                f"Too many pages requested (over {_MAX_REQUESTED_PAGES})",
+                {"doc_name": doc_name},
+                {
+                    "summary": "The page specification spans too many pages",
+                    "options": [
+                        "Request a narrower page range",
+                        "The response holds only a few pages per call - page through with several smaller requests",
+                    ],
+                },
+                "INVALID_INPUT",
+            )
+        if exc.code == "nonpositive":
+            return None, _failure(
+                "Invalid page numbers. Page numbers must be positive integers",
+                {"doc_name": doc_name},
+                {
+                    "summary": "Invalid page numbers provided",
+                    "options": [
+                        "Page numbers must be positive integers (>= 1)",
+                        "Check the page specification format",
+                    ],
+                },
+                "INVALID_INPUT",
+            )
         return None, _failure(
-            "Invalid page numbers. Page numbers must be positive integers",
+            "Invalid page specification format",
             {"doc_name": doc_name},
             {
-                "summary": "Invalid page numbers provided",
+                "summary": "Failed to parse the pages parameter",
                 "options": [
-                    "Page numbers must be positive integers (>= 1)",
-                    "Check the page specification format",
+                    'Use valid formats: "5", "3,7,10", "5-10", or "1-3,7,9-12"',
+                    "Ensure page numbers are positive integers",
                 ],
             },
             "INVALID_INPUT",
         )
-    return sorted(expanded), None
 
 
 def _format_page_spec(pages: list[int]) -> str:
@@ -702,7 +741,8 @@ def _browse_documents(client, folder_id: str = "root", recursive: bool = False,
         window = listing.get("documents") or []
         total = listing.get("total")
     else:
-        scoped = _scope_documents(_all_documents(client), _allowed_ids)
+        scoped = _scope_documents(_all_documents(client, stop_ids=_allowed_ids),
+                                  _allowed_ids)
         window, total = scoped[offset:offset + limit], len(scoped)
     window_end = offset + len(window)
     has_more = bool(window) and (window_end < total if isinstance(total, int)
@@ -863,7 +903,9 @@ def _get_document_structure(client, doc_name: str,
         raw_tree = getattr(getattr(client, "_api", None), "raw_tree", None)
         tree = raw_tree(entry["id"]) if raw_tree is not None else None
         if tree is None:
-            tree = client.get_tree(entry["id"], node_summary=True).get("result")
+            # _format_structure strips text anyway — don't download it.
+            tree = client.get_tree(entry["id"], node_summary=True,
+                                   include_text=False).get("result")
     except PageIndexAPIError as exc:
         return _failure(
             f"Failed to retrieve document structure: {exc}",
@@ -1020,10 +1062,12 @@ def _get_page_content(client, doc_name: str, pages: str,
         markdown = item.get("markdown") if item else None
         text = (markdown if isinstance(markdown, str)
                 else f"Page {page} content not available")
-        if not included or budget - len(text) >= 0:
-            content.append({"page": page, "text": text})
+        entry = {"page": page, "text": text}
+        size = _serialized_size(entry) + 2  # +2: json ", " item separator
+        if not included or budget - size >= 0:
+            content.append(entry)
             included.append(page)
-            budget -= len(text)
+            budget -= size
         else:
             remaining.append(page)
 
@@ -1087,7 +1131,7 @@ def _remove_document(client, doc_names: list[str],
                         {"summary": "Too many documents in one call",
                          "options": ["Delete at most 10 documents per call"]},
                         "INVALID_INPUT")
-    documents = _all_documents(client)
+    documents = _all_documents(client, stop_ids=_allowed_ids)
     results = []
     for doc_name in doc_names:
         entry, error = _resolve_document(client, doc_name, documents=documents,
@@ -1126,11 +1170,11 @@ def tool_names(include_management: bool = False) -> tuple[str, ...]:
     return _READ_TOOLS + (_MANAGEMENT_TOOLS if include_management else ())
 
 
-def _coerce_bool_args(name: str, kwargs: dict[str, Any]) -> None:
+def _coerce_bool_args(schema: dict, kwargs: dict[str, Any]) -> None:
     """Models routinely send booleans as JSON strings ("false"); the bare
-    truthiness tests downstream would read those as True."""
-    properties = TOOL_CONTRACT.get(name, {}).get("schema", {}).get(
-        "properties", {})
+    truthiness tests downstream would read those as True. Runs on both
+    dispatch paths — call_tool and the cloud bridge invoker."""
+    properties = (schema or {}).get("properties", {})
     for key, spec in properties.items():
         value = kwargs.get(key)
         if spec.get("type") == "boolean" and isinstance(value, str):
@@ -1168,7 +1212,7 @@ def call_tool(client, name: str, arguments: dict[str, Any],
     # "omit if ..." semantics, same as the cloud bridge invoker).
     kwargs = {key: value for key, value in (arguments or {}).items()
               if not key.startswith("_") and value is not None}
-    _coerce_bool_args(name, kwargs)
+    _coerce_bool_args(TOOL_CONTRACT.get(name, {}).get("schema", {}), kwargs)
     try:
         if doc_ids is not None:
             ids = [doc_ids] if isinstance(doc_ids, str) else doc_ids
@@ -1229,9 +1273,11 @@ _LOCAL_DESCRIPTIONS: dict[str, str] = {
         'Folder browsing and semantic ranking (sort="relevance") are not '
         "supported in local mode yet — they work on PageIndex cloud."
     ),
-    "get_page_content": TOOL_CONTRACT["get_page_content"]["description"]
-    .replace(" Embedded image paths in the response feed into "
-             "`get_document_image()`.", ""),
+    # Drop the sentence naming the cloud-only image tool, whatever its
+    # wording; the contract-refresh test pins that something was removed.
+    "get_page_content": re.sub(
+        r"\s*[^.]*`get_document_image\(\)`[^.]*\.", "",
+        TOOL_CONTRACT["get_page_content"]["description"]),
 }
 
 _LOCAL_PARAM_DESCRIPTIONS: dict[tuple[str, str], str] = {
@@ -1261,9 +1307,6 @@ def _local_schema(name: str) -> dict[str, Any]:
     return schema
 
 
-def _docstring(name: str) -> str:
-    return _tool_docstring(_local_description(name),
-                           _local_schema(name)["properties"])
 
 
 _SCHEMA_TYPE_MAP = {"string": str, "integer": int, "number": float,
@@ -1298,17 +1341,23 @@ def _annotation_for(spec: dict) -> Any:
     return Optional[base] if nullable else base
 
 
-def _bridge_invoker(bridge, name: str) -> "Callable[[dict], tuple[str, bool]]":
-    """One cloud tool call proxied over MCP: None-valued arguments are
-    dropped (None ≡ omitted, matching the contract's "omit if ..."
-    semantics) and failures are contained in the error envelope. Returns
-    (envelope_text, is_error), like call_tool."""
+def _bridge_invoker(bridge, name: str, schema: dict,
+                    ) -> "Callable[[dict], tuple[str, bool]]":
+    """One cloud tool call proxied over MCP: string booleans are coerced
+    (same as call_tool), None-valued arguments are dropped (None ≡ omitted,
+    matching the contract's "omit if ..." semantics) and failures are
+    contained in the error envelope — except 401/403, which re-raise.
+    Returns (envelope_text, is_error), like call_tool."""
     def _invoke(arguments: dict[str, Any]) -> tuple[str, bool]:
         try:
             arguments = {key: value for key, value in arguments.items()
                          if value is not None}
+            _coerce_bool_args(schema, arguments)
             return bridge.call_tool(name, arguments)
         except Exception as exc:
+            if (isinstance(exc, PageIndexAPIError)
+                    and exc.status_code in (401, 403)):
+                raise
             payload, _ = _failure(
                 f"{name} failed: {exc}", None,
                 {"summary": "Unexpected error while running the tool",
@@ -1321,22 +1370,23 @@ def _bridge_invoker(bridge, name: str) -> "Callable[[dict], tuple[str, bool]]":
     return _invoke
 
 
-def _make_bridge_function(bridge, meta: dict) -> Callable[..., str]:
-    """One plain function for a cloud tool: real signature and docstring from
-    the server's schema, invocation proxied over MCP, errors contained."""
+def _make_tool_function(name: str, description: str, schema: dict,
+                        invoke: "Callable[[dict], tuple[str, bool]]",
+                        ) -> Callable[..., str]:
+    """One plain function for a tool: real signature and docstring from the
+    schema, errors contained by the invoker; arguments the signature
+    rejects come back as the guided envelope instead of raising."""
     import keyword
 
-    name = str(meta.get("name") or "")
-    schema = meta.get("inputSchema") or {}
     properties: dict[str, Any] = schema.get("properties") or {}
     required = set(schema.get("required") or [])
-    _invoke = _bridge_invoker(bridge, name)
+    _invoke = invoke
 
     params_usable = all(param.isidentifier() and not keyword.iskeyword(param)
                         and param != "_invoke"
                         for param in properties)
     if not params_usable:
-        def proxy(**kwargs: Any) -> str:
+        def inner(**kwargs: Any) -> str:
             return _invoke(kwargs)[0]
     else:
         ordered = ([p for p in properties if p in required]
@@ -1349,7 +1399,9 @@ def _make_bridge_function(bridge, meta: dict) -> Callable[..., str]:
         namespace: dict[str, Any] = {"_invoke": _invoke}
         exec(f"def _synthesized({rendered}):\n"
              f"    return _invoke({args_literal})[0]", namespace)
-        proxy = namespace["_synthesized"]
+        inner = namespace["_synthesized"]
+        # binding TypeErrors quote __qualname__, not __name__
+        inner.__name__ = inner.__qualname__ = name or "tool"
         annotations: dict[str, Any] = {}
         for p in ordered:
             annotation = _annotation_for(properties[p])
@@ -1359,9 +1411,26 @@ def _make_bridge_function(bridge, meta: dict) -> Callable[..., str]:
                 annotation = Optional[annotation]
             annotations[p] = annotation
         annotations["return"] = str
-        proxy.__annotations__ = annotations
+        inner.__annotations__ = annotations
+
+    def proxy(*args: Any, **kwargs: Any) -> str:
+        # The invoker lets only 401/403 auth failures through, so a
+        # TypeError here is the binding rejecting the arguments.
+        try:
+            return inner(*args, **kwargs)
+        except TypeError as exc:
+            payload, _ = _failure(
+                f"Invalid arguments for {name}: {exc}", None,
+                {"summary": "Invalid tool arguments",
+                 "options": [f"Check the {name}() parameter names "
+                             "and types"]},
+                "INVALID_INPUT",
+            )
+            return _dumps(payload)
+    proxy.__signature__ = inspect.signature(inner)  # type: ignore[attr-defined]
+    proxy.__annotations__ = dict(inner.__annotations__)
     proxy.__name__ = proxy.__qualname__ = name or "tool"
-    proxy.__doc__ = _tool_docstring(meta.get("description") or "", properties)
+    proxy.__doc__ = _tool_docstring(description or "", properties)
     return proxy
 
 
@@ -1369,19 +1438,26 @@ _BRIDGES: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
 _BRIDGES_LOCK = threading.Lock()
 
 
-def _cloud_bridge(client):
-    """One bridge per client: tool discovery and instructions share a single
-    MCP session. Weak-keyed off the instance so clients stay picklable; the
+def _cloud_bridge(client, gated: bool = False):
+    """One bridge per client and endpoint gate (``gated`` = the read-only
+    ?tools=read endpoint); tool discovery and instructions share a session
+    per gate. Weak-keyed off the instance so clients stay picklable; the
     lock closes the check-then-set race under concurrent first calls."""
     with _BRIDGES_LOCK:
-        bridge = _BRIDGES.get(client)
+        # A rotated api_key or moved BASE_URL rebuilds the bridges.
+        auth = (client.BASE_URL, client.api_key)
+        bridges, seen = _BRIDGES.get(client) or ({}, None)
+        if seen != auth:
+            bridges = {}
+        bridge = bridges.get(gated)
         if bridge is None:
             from .mcp_bridge import McpBridge
             bridge = McpBridge(
-                f"{client.BASE_URL}/mcp",
-                {"Authorization": f"Bearer {client.api_key}"},
+                f"{auth[0]}/mcp" + ("?tools=read" if gated else ""),
+                {"Authorization": f"Bearer {auth[1]}"},
             )
-            _BRIDGES[client] = bridge
+            bridges[gated] = bridge
+            _BRIDGES[client] = (bridges, auth)
         return bridge
 
 
@@ -1401,17 +1477,15 @@ def _read_only_tools(tools_meta: list[dict]) -> list[dict]:
     return filtered
 
 
-def _build_cloud_agent_tools(client, include_management: bool) -> list[Callable[..., str]]:
-    bridge = _cloud_bridge(client)
-    tools_meta = bridge.list_tools()
-    if not include_management:
-        tools_meta = _read_only_tools(tools_meta)
-    return [_make_bridge_function(bridge, meta) for meta in tools_meta]
-
-
 def _require_local_scope(client, doc_ids) -> None:
     """The allowlist is enforced in-process; cloud lookups run server-side,
-    so accepting doc_ids there would be advisory-only — refuse loudly."""
+    so accepting doc_ids there would be advisory-only — refuse loudly.
+    An empty allowlist is refused too: it would scope the agent to
+    nothing, with no signal to the caller."""
+    if doc_ids is not None and not doc_ids:
+        raise PageIndexAPIError(
+            "doc_id is empty. Pass one or more document IDs, or omit "
+            "doc_id to give the agent the whole library.")
     if doc_ids is not None and getattr(client, "api_key", None):
         raise PageIndexAPIError(
             "doc_ids scoping applies to local tools only — cloud calls "
@@ -1427,7 +1501,7 @@ def _tool_specs(client, include_management: bool = False, doc_ids=None,
     is the local chat scope; cloud scoping is server-side."""
     _require_local_scope(client, doc_ids)
     if getattr(client, "api_key", None):
-        bridge = _cloud_bridge(client)
+        bridge = _cloud_bridge(client, gated=not include_management)
         tools_meta = bridge.list_tools()
         if not include_management:
             tools_meta = _read_only_tools(tools_meta)
@@ -1435,7 +1509,8 @@ def _tool_specs(client, include_management: bool = False, doc_ids=None,
                  meta.get("description") or "",
                  copy.deepcopy(meta.get("inputSchema"))
                  or {"type": "object", "properties": {}},
-                 _bridge_invoker(bridge, str(meta.get("name") or "tool")))
+                 _bridge_invoker(bridge, str(meta.get("name") or "tool"),
+                                 meta.get("inputSchema") or {}))
                 for meta in tools_meta]
 
     def local_invoke(name: str) -> "Callable[[dict], tuple[str, bool]]":
@@ -1448,7 +1523,8 @@ def _tool_specs(client, include_management: bool = False, doc_ids=None,
             for name in tool_names(include_management)]
 
 
-def build_agent_tools(client, include_management: bool = False) -> list[Callable[..., str]]:
+def build_agent_tools(client, include_management: bool = False,
+                      doc_ids=None) -> list[Callable[..., str]]:
     """Plain synchronous functions bound to `client`.
 
     Cloud: one function per tool of the live cloud MCP tool set, signatures
@@ -1457,53 +1533,11 @@ def build_agent_tools(client, include_management: bool = False) -> list[Callable
     the JSON envelope as a string and never raises for arguments its
     signature accepts (cloud-only parameters are absent from the local
     signatures; the call_tool path answers them with the guided envelope).
+    ``doc_ids`` is the local allowlist, as in ``_tool_specs``.
     """
-    if getattr(client, "api_key", None):
-        return _build_cloud_agent_tools(client, include_management)
-
-    def browse_documents(offset: int = 0, limit: int = 10) -> str:
-        return call_tool(client, "browse_documents", {
-            "offset": offset, "limit": limit,
-        })[0]
-
-    def get_document(doc_name: str, wait_for_completion: bool = False) -> str:
-        return call_tool(client, "get_document", {
-            "doc_name": doc_name,
-            "wait_for_completion": wait_for_completion,
-        })[0]
-
-    def get_document_structure(doc_name: str, part: int = 1,
-                               wait_for_completion: bool = False) -> str:
-        return call_tool(client, "get_document_structure", {
-            "doc_name": doc_name, "part": part,
-            "wait_for_completion": wait_for_completion,
-        })[0]
-
-    def get_page_content(doc_name: str, pages: str,
-                         wait_for_completion: bool = False) -> str:
-        return call_tool(client, "get_page_content", {
-            "doc_name": doc_name, "pages": pages,
-            "wait_for_completion": wait_for_completion,
-        })[0]
-
-    def remove_document(doc_names: list[str]) -> str:
-        return call_tool(client, "remove_document", {
-            "doc_names": doc_names,
-        })[0]
-
-    functions = {
-        "browse_documents": browse_documents,
-        "get_document": get_document,
-        "get_document_structure": get_document_structure,
-        "get_page_content": get_page_content,
-        "remove_document": remove_document,
-    }
-    tools = []
-    for name in tool_names(include_management):
-        function = functions[name]
-        function.__doc__ = _docstring(name)
-        tools.append(function)
-    return tools
+    return [_make_tool_function(name, description, schema, invoke)
+            for name, description, schema, invoke
+            in _tool_specs(client, include_management, doc_ids)]
 
 
 # ── agent instructions ──
@@ -1557,12 +1591,13 @@ AGENT_INSTRUCTIONS = "\n\n".join([
 ])
 
 
-def _base_instructions(client) -> str:
-    """Cloud: the live instructions the MCP server serves for this key's
-    tool set. Local: the built-in subset instructions."""
+def _base_instructions(client, include_management: bool = False) -> str:
+    """Cloud: the live instructions the MCP server serves for the tool set
+    actually shipped. Local: the built-in subset instructions."""
     if not getattr(client, "api_key", None):
         return AGENT_INSTRUCTIONS
-    instructions = _cloud_bridge(client).instructions()
+    instructions = _cloud_bridge(
+        client, gated=not include_management).instructions()
     if not isinstance(instructions, str) or not instructions.strip():
         raise PageIndexAPIError(
             "The MCP server returned no agent instructions — refusing to "
@@ -1585,9 +1620,29 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]:
         return None
     doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id)
     if not doc_ids:
-        return None
-    details = [client.get_document(one_id) for one_id in doc_ids]
-    listing = _all_documents(client)
+        # An empty selection must fail loud: washing it to None would mean
+        # "everything", and the tool-layer allowlist would mean "nothing".
+        raise PageIndexAPIError(
+            "doc_id is empty. Pass one or more document IDs, or omit "
+            "doc_id to give the agent the whole library.")
+    details = []
+    missing = []
+    for one_id in doc_ids:
+        try:
+            details.append(client.get_document(one_id))
+        except PageIndexAPIError as exc:
+            # Batch only a definite not-found/denied (local raises carry no
+            # status); a cloud transport failure (429/5xx) propagates raw.
+            if exc.status_code not in (None, 403, 404):
+                raise
+            missing.append(str(one_id))
+    if missing:
+        raise PageIndexAPIError(
+            "Documents not found or access denied: " + ", ".join(missing))
+    # Scoped: the listing only backfills the target docs' metadata (list
+    # entries carry it, get_document does not — cloud parity), so paging
+    # can stop at those ids. Unscoped needs it all for the shadow check.
+    listing = _all_documents(client, stop_ids=doc_ids if scoped else None)
     documents = ([{**detail, "id": one_id}
                   for one_id, detail in zip(doc_ids, details)]
                  if scoped else listing)
@@ -1625,9 +1680,10 @@ def doc_targeting_block(client, doc_id, scoped: bool = False) -> Optional[str]:
     )
 
 
-def build_agent_instructions(client, doc_id=None, scoped: bool = False) -> str:
+def build_agent_instructions(client, doc_id=None, scoped: bool = False,
+                             include_management: bool = False) -> str:
     """Orchestration guidance for document QA agents; with doc_id, appends
     the target documents and directs the agent to work within them."""
-    base = _base_instructions(client)
+    base = _base_instructions(client, include_management)
     block = doc_targeting_block(client, doc_id, scoped=scoped)
     return base if block is None else base + "\n\n" + block

pageindex/client.py

@@ -10,35 +10,37 @@ from typing import Any, Callable, Iterator, Optional, Union, cast
 from .errors import PageIndexAPIError
 
 
+_litellm_preload_started = False
+
+
 def _preload_litellm() -> None:
-    try:
-        import litellm  # noqa: F401
-    except Exception:
-        pass
+    """Start litellm's multi-second import in the background, once per
+    process — a per-client thread would churn under per-request clients."""
+    # LiteLLM's import otherwise fetches its model map over the network —
+    # seconds of blocking (or a hang offline). Stamped here, not at package
+    # import, so merely importing pageindex leaves the host process's own
+    # litellm untouched; setdefault, so an explicit user choice wins.
+    os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True")
+    global _litellm_preload_started
+    if _litellm_preload_started:
+        return
+    _litellm_preload_started = True
+
+    def _import() -> None:
+        try:
+            import litellm  # noqa: F401
+        except Exception:
+            pass
+
+    threading.Thread(target=_import, daemon=True).start()
 
 
 def _parse_pages(pages: str) -> list[int]:
-    result: set[int] = set()
-    too_many = (f"Page specification '{pages}' spans more than "
-                "10000 pages; request a narrower range")
-    for part in pages.split(","):
-        part = part.strip()
-        if "-" in part:
-            start, end = (int(x) for x in part.split("-", 1))
-            if start > end:
-                raise ValueError(f"Invalid range '{part}': start must be <= end")
-        else:
-            start = end = int(part)
-        # Bound each part arithmetically before materializing it — a spec
-        # like "1-999999999" would otherwise expand to a billion integers.
-        # The cap is on distinct pages, so overlapping parts (a parent
-        # section plus its children) don't double-count.
-        if end - start + 1 > 10_000:
-            raise ValueError(too_many)
-        result.update(range(start, end + 1))
-        if len(result) > 10_000:
-            raise ValueError(too_many)
-    return sorted(result)
+    from .agent_tools import _PageSpecError, _expand_pages
+    try:
+        return _expand_pages(pages)
+    except _PageSpecError as exc:
+        raise PageIndexAPIError(str(exc)) from exc
 
 
 def _agents_sdk_model_name(model: str) -> str:
@@ -155,7 +157,7 @@ class PageIndexClient:
             self.model = opt.model
             self.index_model = opt.index_model
             self.summary_model = opt.summary_model
-            self.chat_model = _agents_sdk_model_name(opt.chat_model)
+            self.chat_model = opt.chat_model
             self.chat_backend = chat_backend
             self.storage_path = storage_path or ".pageindex"
             from .local_api import LocalAPI
@@ -163,18 +165,22 @@ class PageIndexClient:
                 storage_path=self.storage_path,
                 model=self.model,
                 summary_model=self.summary_model,
-                retrieve_model=self.chat_model,
                 index_backend=index_backend,
             )
             # LiteLLM's multi-second import would otherwise land on the
             # first chat call; failures resurface there with real context.
-            threading.Thread(target=_preload_litellm, daemon=True).start()
+            _preload_litellm()
 
     @property
     def retrieve_model(self):
         """Legacy name for ``chat_model``."""
         return self.chat_model
 
+    @retrieve_model.setter
+    def retrieve_model(self, value):
+        # 0.2.9 allowed assignment; keep the write path working too.
+        self.chat_model = value
+
     # ---------- DOCUMENT SUBMISSION ----------
 
     def submit_document(
@@ -432,7 +438,12 @@ class PageIndexClient:
         if stream:
             return cast(Iterator[str], result)
         envelope = cast(dict[str, Any], result)
-        return envelope["choices"][0]["message"]["content"] or ""
+        try:
+            return envelope["choices"][0]["message"]["content"] or ""
+        except (KeyError, IndexError, TypeError) as exc:
+            raise PageIndexAPIError(
+                "The chat response carries no answer: "
+                f"{str(envelope)[:200]}") from exc
 
     def chat_completions(
         self,
@@ -598,11 +609,12 @@ class PageIndexClient:
             model: Backend model name (defaults to ``chat_model``).
             stream: Yield Responses stream events as dicts — one logical
                 response per call: per-turn backend lifecycle events are
-                collapsed, sequence numbers are reassigned monotonically,
-                and ``output_index`` is re-based onto the single logical
-                ``output``. The single final event is the terminal
-                ``response.*`` for the run's status; its ``response``
-                carries the tool outputs in ``items``.
+                collapsed to one opening ``response.created`` and one
+                final terminal event, sequence numbers are reassigned
+                monotonically, and ``output_index`` is re-based onto the
+                single logical ``output``. The final event is the
+                terminal ``response.*`` for the run's status; its
+                ``response`` carries the tool outputs in ``items``.
             doc_id: Document ID or list of IDs to scope the conversation.
                 Keep it identical across a conversation's calls — the
                 targeting block it adds is re-set each call and is part
@@ -673,7 +685,10 @@ class PageIndexClient:
         final message envelope with cross-turn aggregated ``usage`` plus a
         ``messages`` field — the full new turn sequence, valid for verbatim
         append to your history. The managed system prompt carries a
-        ``cache_control`` breakpoint.
+        ``cache_control`` breakpoint, and the request sets the top-level
+        ``cache_control`` so each turn re-reads the growing conversation
+        from cache — skipped when your own blocks already use the three
+        remaining breakpoints (the managed prompt holds the fourth).
 
         Args:
             messages: Native Messages-format history (including prior
@@ -683,7 +698,9 @@ class PageIndexClient:
             max_tokens: Per-turn output budget the Messages API requires on
                 the wire; the default is resolved per model (8192, or 4096
                 for the claude-3 generation whose ceiling is lower) so the
-                simple call needs only a question. Passed through.
+                simple call needs only a question, and rises to
+                budget_tokens + 8192 when ``thinking`` is enabled (the wire
+                requires max_tokens above the budget). Passed through.
             stream: Yield the Anthropic SDK's event stream across turns
                 (its native event objects, including SDK-synthesized
                 convenience events), one message sequence per turn.
@@ -776,13 +793,16 @@ class PageIndexClient:
 
     # ---------- AGENT INTEGRATION ----------
 
-    def agent_tools(self, include_management: bool = False) -> list[Callable[..., str]]:
+    def agent_tools(
+        self, include_management: bool = False,
+        doc_id: Optional[Union[str, list[str]]] = None,
+    ) -> list[Callable[..., str]]:
         """
         Plain functions for any agent framework (LangChain, PydanticAI, ...).
         For the OpenAI / Claude Agent SDKs, prefer ``as_openai_tools()`` /
         ``as_claude_mcp()``.
 
-        Cloud: the full cloud tool set, discovered live from the PageIndex
+        Cloud: the full live read tool set, discovered from the PageIndex
         MCP server when this method is called — one function per tool,
         signature and docstring synthesized from the server's schemas, calls
         executed from your process over MCP. Raises PageIndexAPIError if the
@@ -798,9 +818,13 @@ class PageIndexClient:
                 library. Local: adds ``remove_document``. Cloud: by default
                 only tools the server marks read-only are exposed; True
                 exposes the server's complete list (upload, delete, ...).
+            doc_id: Local only — restrict the tools to this document ID
+                (or list of IDs), enforced at the tool layer: out-of-scope
+                lookups return NOT_FOUND. Raises on cloud, where scoping
+                is server-side.
         """
         from .agent_tools import build_agent_tools
-        return build_agent_tools(self, include_management)
+        return build_agent_tools(self, include_management, doc_ids=doc_id)
 
     def as_openai_tools(self, include_management: bool = False,
                         hosted: bool = False,
@@ -851,14 +875,14 @@ class PageIndexClient:
         """doc_id for the tool layer: passed through locally (structural
         allowlist), dropped on cloud where scoping is server-side and the
         config helpers keep prompt-level targeting."""
-        if not getattr(self, "api_key", None):
-            return doc_id
         if doc_id is not None and not doc_id:
-            # Cloud has no tool-layer allowlist to make an empty scope mean
-            # "nothing"; dropping it would silently mean "everything".
+            # An empty scope means "nothing" locally (empty allowlist) and
+            # cannot be represented on cloud; both refuse it loudly.
             raise PageIndexAPIError(
                 "doc_id is empty. Pass one or more document IDs, or omit "
                 "doc_id to give the agent the whole library.")
+        if not getattr(self, "api_key", None):
+            return doc_id
         return None
 
     def openai_agent_config(
@@ -882,6 +906,16 @@ class PageIndexClient:
         environment, so its model auth comes from there —
         ``chat_backend`` does not travel with it.
 
+        Prompt caching configures itself for most destinations (OpenAI
+        server-side; Anthropic- and Bedrock-hosted Claude via LiteLLM's
+        defaults). Vertex-hosted Claude is the exception — pass the
+        injection points yourself::
+
+            Agent(**config, model_settings=ModelSettings(extra_args={
+                "cache_control_injection_points": [
+                    {"location": "message", "role": "system"},
+                    {"location": "message", "index": -1}]}))
+
         Args:
             doc_id: Document ID or list of IDs to target, as in
                 ``agent_instructions``. Local: also enforced at the tool
@@ -889,19 +923,27 @@ class PageIndexClient:
                 (tool scoping is server-side).
             include_management (bool): Also expose tools that modify the
                 library.
-            model: Backend model name; overrides the local default.
+            model: Backend model name; overrides the local default. Same
+                grammar as ``chat_model`` (LiteLLM names; bare names are
+                OpenAI-compatible shorthand).
         """
         from .agent_tools import build_agent_instructions
         scope = self._local_doc_scope(doc_id)
         config: dict[str, Any] = {
             "name": "PageIndex",
-            "instructions": build_agent_instructions(self, doc_id,
-                                                     scoped=scope is not None),
+            "instructions": build_agent_instructions(
+                self, doc_id, scoped=scope is not None,
+                include_management=include_management),
             "tools": self.as_openai_tools(include_management, doc_id=scope),
         }
         model = model or getattr(self, "chat_model", None)
         if model:
-            config["model"] = model
+            config["model"] = _agents_sdk_model_name(model)
+            if config["model"].startswith("litellm/"):
+                # The runner resolves this model through LiteLLM in the
+                # caller's process, outside our completion helpers.
+                from .utils import _repair_litellm_types
+                _repair_litellm_types()
         return config
 
     def as_anthropic_tools(self, include_management: bool = False,
@@ -961,6 +1003,7 @@ class PageIndexClient:
         asynchronous: bool = False,
         max_tokens: Optional[int] = None,
         max_turns: Optional[int] = None,
+        thinking: Optional[dict] = None,
     ) -> dict[str, Any]:
         """
         Document QA ``tool_runner`` kwargs for the Anthropic SDK in one
@@ -973,10 +1016,14 @@ class PageIndexClient:
 
         Sugar over the explicit form — ``agent_instructions`` (with
         ``doc_id`` targeting) as the system prompt and
-        ``as_anthropic_tools`` as the tools — plus the same defaults
-        ``messages()`` applies: a per-model ``max_tokens`` and a
-        ``max_iterations`` bound of 10. To customize further, switch to
-        those methods directly.
+        ``as_anthropic_tools`` as the tools — plus the ``max_tokens``
+        default and 10-turn ``max_iterations`` bound ``messages()`` uses,
+        and a top-level ``cache_control`` so each loop turn re-reads the
+        growing prompt from cache (pop the key if you place your own
+        breakpoints — the API allows four). Unlike ``messages()``,
+        ``system`` here is the bare instructions string, without the chat
+        header or its block-level breakpoint. To customize further,
+        switch to those methods directly.
 
         Args:
             model: Backend model name (also resolves the ``max_tokens``
@@ -992,19 +1039,27 @@ class PageIndexClient:
             max_tokens: Per-turn output budget; default resolved per
                 model.
             max_turns: Agent-loop bound; default 10.
+            thinking: Anthropic ``thinking`` config, included in the
+                kwargs; an enabled budget also lifts the ``max_tokens``
+                default above it. Pass it here, not alongside the
+                unpacked config, so the default stays valid.
         """
         from .agent_tools import build_agent_instructions
-        from .local_chat import _default_max_tokens
+        from .local_chat import _default_max_tokens, _validate_max_turns
+        _validate_max_turns(max_turns)
         scope = self._local_doc_scope(doc_id)
         return {
             "model": model,
             "max_tokens": (max_tokens if max_tokens is not None
-                           else _default_max_tokens(model)),
-            "system": build_agent_instructions(self, doc_id,
-                                               scoped=scope is not None),
+                           else _default_max_tokens(model, thinking)),
+            "system": build_agent_instructions(
+                self, doc_id, scoped=scope is not None,
+                include_management=include_management),
             "tools": self.as_anthropic_tools(include_management, asynchronous,
                                              doc_id=scope),
             "max_iterations": max_turns if max_turns is not None else 10,
+            **({"thinking": thinking} if thinking is not None else {}),
+            "cache_control": {"type": "ephemeral"},
         }
 
     def as_claude_mcp(self, include_management: bool = False,
@@ -1072,8 +1127,9 @@ class PageIndexClient:
         from .agent_tools import build_agent_instructions
         scope = self._local_doc_scope(doc_id)
         return {
-            "system_prompt": build_agent_instructions(self, doc_id,
-                                                      scoped=scope is not None),
+            "system_prompt": build_agent_instructions(
+                self, doc_id, scoped=scope is not None,
+                include_management=include_management),
             "mcp_servers": {server_name: self.as_claude_mcp(
                 include_management, doc_id=scope)},
             # Pre-approval only — the server itself is already gated (the
@@ -1081,7 +1137,10 @@ class PageIndexClient:
             "allowed_tools": [f"mcp__{server_name}"],
         }
 
-    def agent_instructions(self, doc_id: Optional[Union[str, list[str]]] = None) -> str:
+    def agent_instructions(
+        self, doc_id: Optional[Union[str, list[str]]] = None,
+        include_management: bool = False,
+    ) -> str:
         """
         Orchestration guidance for document QA agents — pass as the agent's
         system prompt (or append to your own).
@@ -1094,12 +1153,19 @@ class PageIndexClient:
 
         With ``doc_id`` (str or list, same shape as ``chat_completions``),
         appends the target documents' names and metadata and directs the
-        agent to work within them. Raises PageIndexAPIError if a doc_id does
-        not exist, or if its name is shadowed by a newer same-name document
-        (the name-addressed tools could not reach it).
+        agent to work within them. Raises PageIndexAPIError if a doc_id
+        does not exist, or if its name is shadowed by a newer same-name
+        document — the name-addressed tools could not reach it (the
+        ``*_agent_config`` bundles, whose tools carry the doc_id scope,
+        relax this to duplicates within the targeted set).
+
+        ``include_management``: fetch the guidance for the full tool set,
+        matching tools built with ``include_management=True`` (cloud;
+        local guidance is a single set).
         """
         from .agent_tools import build_agent_instructions
-        return build_agent_instructions(self, doc_id)
+        return build_agent_instructions(
+            self, doc_id, include_management=include_management)
 
     # ---------- FOLDER MANAGEMENT ----------
 

pageindex/cloud_api.py

@@ -59,7 +59,7 @@ class CloudAPI:
                 (a taken name gains a numeric suffix), when the server
                 returns it.
         """
-        data = {'if_retrieval': True}
+        data: Dict[str, Any] = {'if_retrieval': True}
         if mode is not None:
             data['mode'] = mode
         if beta_headers is not None:
@@ -225,7 +225,7 @@ class CloudAPI:
             headers=self._headers(),
             json=payload,
             stream=stream,
-            timeout=120 if stream else 300
+            timeout=120 if stream else 600
         )
 
         if response.status_code != 200:
@@ -312,7 +312,9 @@ class CloudAPI:
             timeout=30
         )
         if response.status_code != 200:
-            raise PageIndexAPIError(f"Failed to get document metadata: {response.text}")
+            raise PageIndexAPIError(
+                f"Failed to get document metadata: {response.text}",
+                status_code=response.status_code)
         return response.json()
 
     def delete_document(self, doc_id: str) -> Dict[str, Any]:
@@ -356,7 +358,7 @@ class CloudAPI:
         if offset < 0:
             raise ValueError("offset must be non-negative")
 
-        params = {"limit": limit, "offset": offset}
+        params: Dict[str, Any] = {"limit": limit, "offset": offset}
         if folder_id is not None:
             params["folder_id"] = folder_id
 

pageindex/config.yaml

@@ -2,8 +2,8 @@
 # chat_model answers questions on the chat surfaces; set model to use one
 # for both. Unset keys use the SDK defaults shown below. Legacy keys
 # (model, summary_model, retrieve_model) keep working.
-# Indexing names without a provider prefix use the OpenAI SDK directly;
-# for other providers, use "provider/model" (e.g. "anthropic/claude-sonnet-4-6").
+# Model names are LiteLLM's: bare names are OpenAI models; for other
+# providers, use "provider/model" (e.g. "anthropic/claude-sonnet-4-6").
 # index_model: "gpt-5.6-luna"
 # chat_model: "gpt-5.6-sol"
 toc_check_page_num: 20

pageindex/errors.py

@@ -1,2 +1,7 @@
 class PageIndexAPIError(Exception):
-    pass
+    """status_code carries the HTTP status when the failure came from a
+    non-200 cloud response; None otherwise (local mode, client-side)."""
+
+    def __init__(self, *args: object, status_code: int | None = None) -> None:
+        super().__init__(*args)
+        self.status_code = status_code

pageindex/flash/api.py

@@ -96,19 +96,33 @@ def _optimize(structure, page_texts, do_expand, model):
 
 
 def page_index_flash(pdf, summary=True, summary_model=None,
-                     optimize: str | bool = "full", optimize_expand=None,
+                     optimize: str | bool | None = None, optimize_expand=None,
                      optimize_model=None, summary_concurrency=None,
                      use_embedded_toc=True) -> dict:
-    """Build a PageIndex tree structure from a PDF using layout statistics, without an LLM. Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand, ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
-    if optimize is True:
-        optimize = "full"
+    """Build a PageIndex tree structure from a PDF using layout statistics. The tree extraction itself uses no LLM; by default an LLM writes node summaries and expands the tree (``summary=False, optimize=False`` runs fully LLM-free). Args: pdf: path to a PDF file (``str`` or ``pathlib.Path``) or an in-memory binary stream (``io.BytesIO``). summary: if True, generate LLM summaries for each node (requires ``summary_model``). summary_model: the LLM model identifier to use for summary generation. optimize: ``"full"`` for merge + LLM expand (fails fast with ``PageIndexAPIError`` when no LLM key is configured), ``"merge"`` for deterministic merge only, ``False`` to disable. ``True`` is accepted as ``"full"`` for backward compatibility; defaults to ``"full"``. optimize_expand: deprecated — use ``optimize``. Honored only when ``optimize`` is not passed (or is the legacy ``True``): ``False`` maps to ``"merge"``, ``True`` to ``"full"``. optimize_model: the LLM model for expand (defaults to the summary model). summary_concurrency: maximum simultaneous summary model calls; None uses the library default. use_embedded_toc: if True, consume the PDF's embedded bookmarks when trustworthy: deep bookmarks become the frame and the detected sections they lack are grafted back in after noise filtering, coarse ones become the chapter frame with detected nodes re-hung under them (deeper sparse entries are filled in when the page text confirms them, and garbled extracted titles are repaired from the bookmark strings), garbage ones are ignored; adds a ``toc_source`` key to the result. On by default; pass False for the pure detected structure. Returns: dict with keys ``doc_name``, ``doc_title``, ``structure`` (a list of nested ``{"title", "start_index", "end_index", "nodes"}`` dicts; page indexes are 1-based) and ``has_abstract_or_references_section`` (True when a top-level entry is an abstract or references heading). With ``optimize`` an ``optimize`` key reports merge/expand counts and before/after search-cost metrics. """
+    if optimize_expand is not None:
+        import warnings
+        warnings.warn(
+            "optimize_expand is deprecated: pass optimize='full', 'merge', "
+            "or False.", DeprecationWarning, stacklevel=2)
+    if optimize is None or optimize is True:
+        # legacy spellings only — an explicit 'full'/'merge' wins
+        optimize = "merge" if optimize_expand is False else "full"
     if not optimize:
         optimize = False
     elif optimize not in ("full", "merge"):
         raise ValueError(
             f"optimize must be 'full', 'merge', or False, got {optimize!r}")
-    if optimize_expand is not None and optimize:
-        optimize = "full" if optimize_expand else "merge"
+    if optimize == "full":
+        from ..errors import PageIndexAPIError
+        from ..utils import ConfigLoader, _llm_backend, _openai_missing_keys
+        model = (optimize_model or summary_model
+                 or ConfigLoader().load().summary_model)
+        if not _llm_backend.get() and _openai_missing_keys(model):
+            raise PageIndexAPIError(
+                "optimize='full' runs LLM expand and no LLM key is "
+                "configured — set OPENAI_API_KEY, or pass optimize='merge' "
+                "or optimize=False for the LLM-free tree.")
     result = extract_toc(_validate_pdf(pdf), use_embedded_toc=use_embedded_toc)
     structure = result.get("structure", [])
     if optimize and structure:

pageindex/flash/embedded_toc.py

@@ -123,13 +123,19 @@ def read_bookmarks(doc_handle: Union[str, Path, BytesIO]) -> list[dict]:
         doc = pdfium.PdfDocument(handle)
         entries = []
         for item in doc.get_toc():
-            if item.page_index is None:
+            if hasattr(item, "get_dest"):
+                dest = item.get_dest()
+                page_idx = dest.get_index() if dest else None
+                title = item.get_title()
+            else:
+                page_idx, title = item.page_index, item.title
+            if page_idx is None:
                 continue
-            title = (item.title or "").strip()
+            title = (title or "").strip()
             if not title:
                 continue
             entries.append(
-                {"title": title, "level": item.level + 1, "page": item.page_index + 1}
+                {"title": title, "level": item.level + 1, "page": page_idx + 1}
             )
         return entries
     except (pdfium.PdfiumError, OSError, ValueError, TypeError):

pageindex/flash/parser_pdfium_charlevel/char_extract.py

@@ -56,10 +56,21 @@ def _extract_raw_chars(page, text_page) -> tuple[list[dict], list[dict]]:
     name_cache: dict[bytes, str] = {}
     raw_chars: list[dict] = []
     last_obj: dict | None = None
+    skip_next = False
     for index_value in range(count_item):
+        if skip_next:
+            skip_next = False
+            continue
         codepoint = get_unicode(text_page, index_value)
         if codepoint < 0:
             continue
+        # PDFium returns astral characters (U+10000+) as two UTF-16 surrogate
+        # code units in consecutive textpage slots. Reassemble before chr().
+        if 0xD800 <= codepoint <= 0xDBFF and index_value + 1 < count_item:
+            low = get_unicode(text_page, index_value + 1)
+            if 0xDC00 <= low <= 0xDFFF:
+                codepoint = ((codepoint & 0x3FF) << 10) + (low & 0x3FF) + 0x10000
+                skip_next = True
         # u == 0 (PDFium found no unicode for the glyph) is KEPT as '\x00':
         # text extraction emits the raw charcode for unmapped codes, so its items
         # really contain chr(0) for extension-font pieces at code 0, and the

pageindex/flash/parser_pdfium_charlevel/geometry.py

@@ -6,6 +6,8 @@ import ctypes
 import math
 import pypdfium2.raw as pdfium_c
 
+_get_font_name = getattr(pdfium_c, "FPDFFont_GetBaseFontName", None) or pdfium_c.FPDFFont_GetFontName
+
 
 def _obj_rotation(value: float, other_item: float, candidate_item: float, reference_item: float) -> int:
     """Classify a text-object matrix as upright, cardinal rotation, or oblique. Near-cardinal matrices snap to the cardinal bucket; genuinely oblique matrices use the baseline remerge path."""
@@ -128,7 +130,7 @@ def _collect_text_objs(page, text_page) -> list[dict]:
         # Snap back to the shortest decimal so knife-edge font-size comparisons
         # match the content-stream value.
         fs_eff = float(f"{fs_eff:.6g}")
-        name = pdfium_c.FPDFFont_GetFontName(font, font_name_buffer, 256)
+        name = _get_font_name(font, font_name_buffer, 256)
         font_name = (
             bytes(font_name_buffer[:name]).decode("latin-1", errors="replace").rstrip("\x00")
             if name > 1 else ""

pageindex/flash/parser_pdfium_charlevel/pipeline.py

@@ -95,7 +95,7 @@ def _page_pass1(pdf, pdf_doc, page_idx: int, type3_ext: dict, font_map_cache: di
                 # PDFium's output).
                 if raw_chars:
                     _apply_font_unicode(
-                        text_page.raw, raw_chars, objects, show_codes, pdf_doc,
+                        raw_chars, objects, show_codes, pdf_doc,
                         font_map_cache)
         except Exception:
             pass

pageindex/flash/parser_pdfium_charlevel/unicode_apply.py

@@ -5,7 +5,6 @@ from __future__ import annotations
 import bisect
 import difflib
 from collections import Counter
-import pypdfium2.raw as pdfium_c
 
 from .text_normalize import _is_whitespace
 from .font_unicode import _font_unicode_map
@@ -16,7 +15,6 @@ from .code_walk import (
 
 
 def _apply_font_unicode(
-    text_page,
     raw_chars: list[dict],
     objects: list[dict],
     show_codes: list[tuple[int | None, tuple[int, ...], float]],
@@ -287,14 +285,12 @@ def _apply_font_unicode(
                 _synthesize_dropped_glyphs(kept, raw_chars, chars_by_index)
         return
 
-    # Page mode.
-    seq: list[tuple[int, str]] = []
-    char_count = pdfium_c.FPDFText_CountChars(text_page)
-    for char_index in range(char_count):
-        if pdfium_c.FPDFText_IsGenerated(text_page, char_index) == 1:
-            continue
-        codepoint = pdfium_c.FPDFText_GetUnicode(text_page, char_index)
-        seq.append((char_index, chr(codepoint) if codepoint > 0 else "\x00"))
+    # Page mode. Walk the char census char_extract built (surrogate pairs
+    # already merged there): re-reading the textpage would split astral
+    # chars back into two lone-surrogate slots, desync the walk against
+    # their one-char cmap targets, and drop the whole page's patch.
+    seq = [(raw_char["i"], raw_char["ch"])
+           for raw_char in raw_chars if not raw_char["is_gen"]]
     targets: list[str] = []
     for font_index, encoded_text, _tz in show_codes:
         if not encoded_text:

pageindex/flash/parser_pdfium_parallel.py

@@ -23,7 +23,9 @@ from __future__ import annotations
 
 import multiprocessing
 import os
+import sys
 from concurrent.futures import ProcessPoolExecutor
+from contextlib import contextmanager
 from io import BytesIO
 from pathlib import Path
 from typing import Union
@@ -54,6 +56,29 @@ _worker_pdf_doc = None
 _worker_font_maps: dict = {}
 
 
+@contextmanager
+def _anonymous_main():
+    """Hide __main__'s import identity while workers spawn: spawn re-executes
+    the caller's script in every worker otherwise, which for an unguarded
+    script means one duplicate full run per worker. Our workers import
+    everything by module name and never need __main__.
+
+    ponytail: window covers the whole map; a concurrent pool spawned from
+    another thread whose tasks live in __main__ would break during it."""
+    main = sys.modules.get("__main__")
+    if main is None:
+        yield
+        return
+    d = main.__dict__
+    saved = {k: d.pop(k) for k in ("__file__", "__spec__") if k in d}
+    d["__spec__"] = None  # get_preparation_data reads it via attribute access
+    try:
+        yield
+    finally:
+        d.pop("__spec__", None)
+        d.update(saved)
+
+
 def _init_worker(kind: str, payload) -> None:
     global _worker_pdf, _worker_pdf_doc, _worker_font_maps
     # Open the document exactly as parse_charlevel_meta does, including
@@ -122,12 +147,17 @@ def parse_charlevel_meta_parallel(
         initargs=src,
     )
     try:
-        results = list(executor.map(_run_page, range(n_pages)))
+        with _anonymous_main():
+            results = list(executor.map(_run_page, range(n_pages)))
     except Exception:
         # _Type3Detected or any worker/pool failure. Cancel what is queued
         # and rerun sequentially; in-flight pages finish in their workers
         # and are discarded (separate processes, no shared PDFium state).
         executor.shutdown(wait=False, cancel_futures=True)
+        if getattr(multiprocessing.current_process(), "_inheriting", False):
+            # Spawn child re-importing an unguarded __main__; a sequential rerun
+            # here would silently duplicate the caller's whole run per worker.
+            raise
         return parse_charlevel_meta(doc_handle)
     executor.shutdown()
 

pageindex/integrations/claude_agent_sdk.py

@@ -34,14 +34,11 @@ def build_claude_mcp(client, include_management: bool = False, doc_ids=None):
             "as_claude_mcp in local mode requires the Claude Agent SDK — "
             "pip install claude-agent-sdk (or pip install 'pageindex[claude]')."
         ) from exc
-    from ..agent_tools import (TOOL_CONTRACT, _local_description,
-                               _local_schema, call_tool, tool_names)
+    from ..agent_tools import TOOL_CONTRACT, _tool_specs
 
-    def make_handler(name: str):
+    def make_handler(invoke):
         async def handler(arguments: dict[str, Any]) -> dict[str, Any]:
-            text, is_error = await asyncio.to_thread(
-                call_tool, client, name, arguments or {}, doc_ids
-            )
+            text, is_error = await asyncio.to_thread(invoke, arguments or {})
             result: dict[str, Any] = {"content": [{"type": "text", "text": text}]}
             if is_error:
                 result["is_error"] = True
@@ -59,9 +56,10 @@ def build_claude_mcp(client, include_management: bool = False, doc_ids=None):
         return {"annotations": ToolAnnotations(**annotations)}
 
     tools = [
-        tool(name, _local_description(name),
-             _local_schema(name), **tool_kwargs(name))(make_handler(name))
-        for name in tool_names(include_management)
+        tool(name, description, schema,
+             **tool_kwargs(name))(make_handler(invoke))
+        for name, description, schema, invoke
+        in _tool_specs(client, include_management, doc_ids)
     ]
     return create_sdk_mcp_server(name="pageindex", version=sdk_version(),
                                  tools=tools)

pageindex/integrations/openai_agents.py

@@ -65,6 +65,8 @@ def build_openai_tools(client, include_management: bool = False,
                 return _dumps(payload)
             arguments = {key: value for key, value in parsed.items()
                          if value is not None}
+            # is_error has no per-result channel on hand-built FunctionTools
+            # (raising aborts the run — see above); the text is the signal.
             text, _ = await asyncio.to_thread(invoke, arguments)
             return text
 

pageindex/local_api.py

@@ -1,17 +1,17 @@
 """Local implementation of the PageIndex SDK surface."""
 from __future__ import annotations
 
-import asyncio
 import json
 import logging
+import multiprocessing
 import os
 import uuid
-from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime, timezone
 from typing import Any
 
 from .errors import PageIndexAPIError
 from .local_store import DocStore
+from .utils import run_off_loop
 
 logger = logging.getLogger(__name__)
 
@@ -22,31 +22,23 @@ def _now_iso() -> str:
     return now.replace(microsecond=now.microsecond // 1000 * 1000).isoformat()
 
 
-def _run_indexer(func, *args, **kwargs):
-    try:
-        asyncio.get_running_loop()
-    except RuntimeError:
-        return func(*args, **kwargs)
-    with ThreadPoolExecutor(max_workers=1) as pool:
-        return pool.submit(func, *args, **kwargs).result()
 
 
 class LocalAPI:
     """Backs PageIndexClient's local mode. One instance per client."""
 
     def __init__(self, storage_path: str, model: str, summary_model: str,
-                 retrieve_model: str, index_backend: dict | None = None):
+                 index_backend: dict | None = None):
         self._store = DocStore(storage_path)
         self._model = model
         self._summary_model = summary_model
-        self._retrieve_model = retrieve_model
         self._index_backend = index_backend
         from .utils import ConfigLoader
         self._config_loader = ConfigLoader()
 
     def _with_backend(self, func, *args):
         """Scope the indexing lane's connection overrides around one
-        operation — runs inside whatever thread _run_indexer picked."""
+        operation — runs inside whatever thread run_off_loop picked."""
         from .utils import _llm_backend
         token = _llm_backend.set(self._index_backend)
         try:
@@ -64,6 +56,12 @@ class LocalAPI:
         folder_id: str | None = None,
         metadata: dict | None = None,
     ) -> dict[str, Any]:
+        if getattr(multiprocessing.current_process(), "_inheriting", False):
+            raise PageIndexAPIError(
+                "Failed to submit document: called again while a spawned worker "
+                "process was importing your script. Put your top-level code under "
+                "if __name__ == '__main__': so worker processes do not re-run it."
+            )
         if beta_headers is not None:
             raise PageIndexAPIError(
                 "Failed to submit document: beta_headers is not supported in local mode."
@@ -114,11 +112,11 @@ class LocalAPI:
 
         try:
             if mode == "flash":
-                structure, description = _run_indexer(
-                    self._with_backend, self._index_flash, file_path, page_texts
+                structure, description = run_off_loop(
+                    self._with_backend, self._index_flash, file_path
                 )
             else:
-                structure, description = _run_indexer(
+                structure, description = run_off_loop(
                     self._with_backend, self._index_standard, file_path,
                     page_texts
                 )
@@ -128,22 +126,25 @@ class LocalAPI:
             raise PageIndexAPIError(f"Failed to submit document: {e}") from e
 
         doc_id = "pi-" + uuid.uuid4().hex
-        meta = {
-            "id": doc_id,
-            "name": self._unique_doc_name(os.path.basename(file_path)),
-            "description": description,
-            "status": "completed",
-            "createdAt": _now_iso(),
-            "pageNum": len(page_texts),
-            "folderId": None,
-            "metadata": metadata,
-            "mode": mode,
-        }
         pages = [{"page_index": i + 1, "markdown": text}
                  for i, text in enumerate(page_texts)]
         from .utils import remove_fields
-        self._store.save_document(
-            doc_id, meta, remove_fields(structure, fields=["text"]), pages)
+        # Check-then-write under the store lock; the early pre-check above
+        # is advisory only.
+        with self._store.lock():
+            meta = {
+                "id": doc_id,
+                "name": self._unique_doc_name(os.path.basename(file_path)),
+                "description": description,
+                "status": "completed",
+                "createdAt": _now_iso(),
+                "pageNum": len(page_texts),
+                "folderId": None,
+                "metadata": metadata,
+                "mode": mode,
+            }
+            self._store.save_document(
+                doc_id, meta, remove_fields(structure, fields=["text"]), pages)
         return {"doc_id": doc_id, "name": meta["name"]}
 
     def _unique_doc_name(self, name: str) -> str:
@@ -190,17 +191,10 @@ class LocalAPI:
             )
         return structure, result.get("doc_description")
 
-    def _index_flash(self, file_path: str, page_texts: list[str]) -> tuple[list, str | None]:
+    def _index_flash(self, file_path: str) -> tuple[list, str | None]:
         from .flash import page_index_flash
-        from .utils import (add_node_text, create_clean_structure_for_description,
+        from .utils import (create_clean_structure_for_description,
                             generate_doc_description, write_node_id)
-        import litellm
-        if not self._index_backend:
-            env = litellm.validate_environment(self._summary_model)
-            if not env["keys_in_environment"]:
-                raise PageIndexAPIError(
-                    f"Failed to submit document: missing API key for "
-                    f"{self._summary_model}: {', '.join(env['missing_keys'])}")
         result = page_index_flash(file_path, summary=True,
                                   summary_model=self._summary_model,
                                   optimize="full",
@@ -212,7 +206,6 @@ class LocalAPI:
                 "a structure from this PDF."
             )
         write_node_id(structure)
-        add_node_text(structure, [(text, 0) for text in page_texts])
         description = generate_doc_description(
             create_clean_structure_for_description(structure),
             model=self._summary_model,
@@ -356,6 +349,8 @@ def _format_tree_node(node: dict, node_summary: bool) -> dict:
         "node_id": node.get("node_id"),
         "page_index": node.get("start_index"),
     }
+    if node.get("key_items"):
+        out["key_items"] = node["key_items"]
     if node_summary:
         summary = node.get("summary")
         if summary is not None:

pageindex/local_chat.py

@@ -1,27 +1,7 @@
-"""Managed local chat: document-QA agents over the local tools.
-
-Three methods, three backend protocols, routed 1:1: ``chat_completions``
-drives the backend's /chat/completions (any OpenAI-compatible backend,
-final answer only), ``responses`` drives /responses (official-shape
-envelope; the full process transcript rides in ``items`` — round-trip it
-for provider prompt-cache continuation and agent memory), ``messages``
-drives Anthropic's /v1/messages via the SDK's own tool runner
-(tool_use/tool_result round-trip is the format's native behavior).
-
-Content passes through untouched — the caller's messages, the model's
-answers, tool outputs. Native stop reasons pass through on ``messages``;
-the OpenAI engine's abstraction does not surface per-turn finish reasons,
-so ``chat_completions`` reports loop completion as ``"stop"``, while
-``responses`` reports the backend's terminal ``status`` where the wire
-surfaces one (recorded at the transport layer — the framework discards
-it). The SDK owns gatekeeping (structural validation), table-setting
-(managed instructions, tools, doc targeting), tool execution, and billing
-(usage aggregation, envelope ids).
-"""
+"""Managed local chat: document-QA agents over the local tools."""
 from __future__ import annotations
 
 import asyncio
-import concurrent.futures
 import hashlib
 import json
 import os
@@ -52,17 +32,6 @@ def _doc_block(client, doc_id) -> Optional[str]:
     if not isinstance(doc_id, (str, list)):
         raise PageIndexAPIError("doc_id must be a string or a list of "
                                 "strings.")
-    doc_ids = [doc_id] if isinstance(doc_id, str) else list(doc_id)
-    missing = []
-    for one_id in doc_ids:
-        try:
-            client.get_document(one_id)
-        except PageIndexAPIError:
-            missing.append(str(one_id))
-    if missing:
-        raise PageIndexAPIError(
-            "Documents not found or access denied: " + ", ".join(missing)
-        )
     # scoped: the chat surfaces also pass doc_id into the tool layer, so
     # name resolution happens inside the allowlist — only a duplicate name
     # within the targeted set shadows.
@@ -118,16 +87,8 @@ def _split_chat_messages(messages) -> "tuple[list[str], list[dict]]":
 
 
 def _run_sync(coro):
-    try:
-        asyncio.get_running_loop()
-    except RuntimeError:
-        has_loop = False
-    else:
-        has_loop = True
-    if not has_loop:
-        return asyncio.run(coro)
-    with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:
-        return pool.submit(asyncio.run, coro).result()
+    from .utils import run_off_loop
+    return run_off_loop(asyncio.run, coro)
 
 
 _SENTINEL = object()
@@ -212,25 +173,20 @@ def _require_openai_agents(method: str) -> None:
         ) from exc
 
 
+def _sdk_backend(backend) -> dict:
+    """chat_backend for an SDK constructor: LiteLLM takes either endpoint
+    spelling, the openai and anthropic SDKs only ``base_url``."""
+    return {("base_url" if key == "api_base" else key): value
+            for key, value in (backend or {}).items()}
+
+
 def _openai_model(protocol: str, model_name: str, backend=None):
-    """The backend protocol driver — the seam tests replace with a fake.
-
-    chat protocol: LiteLLM, full stop — model names mean what LiteLLM says
-    they mean. Bare names are OpenAI-compatible shorthand (wire form
-    ``openai/<name>``, so OPENAI_API_KEY / OPENAI_BASE_URL keep selecting
-    the backend), a ``litellm/`` prefix strips, and a first segment LiteLLM
-    does not know (a HuggingFace repo id like ``Qwen/...``) is refused with
-    the ``openai/`` form instead of failing inside LiteLLM at request time.
-
-    responses protocol: the Responses API is OpenAI-SDK native — LiteLLM's
-    completion surface speaks the chat.completions format, so
-    provider-prefixed models are refused instead of silently downgrading;
-    bare and ``openai/`` names drive the OpenAI SDK."""
+    """The backend protocol driver — the seam tests replace with a fake."""
     if protocol == "responses":
+        model_name = model_name.removeprefix("litellm/")
         if "/" in model_name and not model_name.startswith("openai/"):
             raise PageIndexAPIError(
-                f"responses() cannot drive "
-                f"'{model_name.removeprefix('litellm/')}': provider-prefixed "
+                f"responses() cannot drive '{model_name}': provider-prefixed "
                 "models route through LiteLLM, which speaks chat.completions, "
                 "not the Responses API. Use chat_completions() (or messages() "
                 "for Anthropic models), or point OPENAI_BASE_URL at a "
@@ -240,10 +196,12 @@ def _openai_model(protocol: str, model_name: str, backend=None):
         import openai
         model_name = model_name.removeprefix("openai/")
         try:
-            sdk_client = openai.AsyncOpenAI(**(backend or {}))
+            sdk_client = openai.AsyncOpenAI(**_sdk_backend(backend))
         except (openai.OpenAIError, TypeError) as exc:
             raise PageIndexAPIError(
                 f"The OpenAI backend is not configured: {exc}") from exc
+        # A caller-owned transport must survive the per-call close.
+        sdk_client._pageindex_caller_http = "http_client" in (backend or {})
         from agents.models.openai_responses import OpenAIResponsesModel
         return OpenAIResponsesModel(model_name, openai_client=sdk_client)
     try:
@@ -254,30 +212,20 @@ def _openai_model(protocol: str, model_name: str, backend=None):
             f"'{model_name}' routes through LiteLLM, but litellm is not "
             "installed. Run:  pip install 'litellm>=1.97'"
         )
-    from .utils import _repair_litellm_types
+    from .utils import _litellm_model, _repair_litellm_types
     _repair_litellm_types()
-    wire = model_name.removeprefix("litellm/")
-    if "/" not in wire or wire.startswith("openai/"):
-        if (not os.environ.get("OPENAI_API_KEY")
-                and not (backend or {}).get("api_key")):
-            raise PageIndexAPIError(
-                "The OpenAI backend is not configured: set the "
-                "OPENAI_API_KEY environment variable, pass an api_key "
-                "in chat_backend / backend (any value works for keyless "
-                "OPENAI_BASE_URL servers), or point chat_model at "
-                "another provider (e.g. 'anthropic/...')."
-            )
-    if "/" not in wire:
-        wire = f"openai/{wire}"
-    providers = getattr(litellm, "provider_list", None)
-    if providers and wire.split("/", 1)[0] not in providers:
+    try:
+        wire = _litellm_model(model_name, backend)
+    except litellm.AuthenticationError as exc:
         raise PageIndexAPIError(
-            f"'{wire}' routes through LiteLLM, but "
-            f"'{wire.split('/', 1)[0]}' is not a LiteLLM provider. For an "
-            "OpenAI-compatible server (vLLM, TGI, Ollama) serving this "
-            f"model id, use 'openai/{wire}' and point OPENAI_BASE_URL "
-            "at the server."
-        )
+            "The OpenAI backend is not configured: set the "
+            "OPENAI_API_KEY environment variable, pass an api_key "
+            "in chat_backend / backend (any value works for keyless "
+            "OPENAI_BASE_URL servers), or point chat_model at "
+            "another provider (e.g. 'anthropic/...')."
+        ) from exc
+    except litellm.NotFoundError as exc:
+        raise PageIndexAPIError(str(exc)) from exc
     return LitellmModel(wire, api_key=(backend or {}).get("api_key"),
                         base_url=(backend or {}).get("base_url"))
 
@@ -290,11 +238,11 @@ def _reported_model(model_name: str) -> str:
 def _cache_extra_args(model_name: str) -> Optional[dict]:
     """Claude's prompt caching is opt-in per request: on Claude models
     routed through LiteLLM (Anthropic direct, Bedrock, Vertex — each
-    channel live-verified), mark the managed system prefix via LiteLLM's
-    injection param so the loop's later turns and a conversation's next
-    calls read it instead of repaying full price. Provider resolution is
-    LiteLLM's own, so this predicate can never disagree with where the
-    request actually routes."""
+    channel live-verified), mark the managed system prefix and the newest
+    message via LiteLLM's injection param so the loop's later turns and a
+    conversation's next calls read them instead of repaying full price.
+    Provider resolution is LiteLLM's own, so this predicate can never
+    disagree with where the request actually routes."""
     if "/" not in model_name or model_name.startswith("openai/"):
         return None
     try:
@@ -305,11 +253,34 @@ def _cache_extra_args(model_name: str) -> Optional[dict]:
         return None
     if provider == "anthropic" or (provider in ("bedrock", "vertex_ai")
                                    and "claude" in model.lower()):
+        # The pair LiteLLM itself seeds for Anthropic and Bedrock: the
+        # stable prefix plus the newest message, so each turn re-reads
+        # the turns before it. Passing it explicitly extends it to Vertex.
         return {"cache_control_injection_points": [
-            {"location": "message", "role": "system"}]}
+            {"location": "message", "role": "system"},
+            {"location": "message", "index": -1}]}
     return None
 
 
+def _openai_protocol(model_name: str) -> bool:
+    """Destinations that speak the OpenAI protocol on the wire, where
+    prompt_cache_key means something and extra_body lands in the request
+    body. Resolution is LiteLLM's own (same as _cache_extra_args), so the
+    answer follows actual routing; azure and openrouter ride the OpenAI
+    protocol without appearing in openai_compatible_providers."""
+    wire = model_name.removeprefix("litellm/")
+    if "/" not in wire or wire.startswith("openai/"):
+        return True
+    try:
+        import litellm
+        _, provider, _, _ = litellm.get_llm_provider(model=wire)
+    except Exception:
+        return False
+    return (provider in ("openai", "azure", "openrouter")
+            or provider in getattr(litellm, "openai_compatible_providers",
+                                   ()))
+
+
 def _merged_backend(client, backend):
     """This call's connection overrides: the client's ``chat_backend``
     under the per-call dict, per-call keys winning."""
@@ -329,8 +300,7 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str,
     # and both OpenAI model classes pass extra_body through verbatim. OpenAI
     # destinations only — LiteLLM plants extra_body as a literal field in
     # other providers' request bodies, which Anthropic rejects as unknown.
-    wire = model_name.removeprefix("litellm/")
-    openai_backend = "/" not in wire or wire.startswith("openai/")
+    openai_backend = _openai_protocol(model_name)
     # Chat-lane effort rides extra_args: LiteLLM takes it as its own
     # top-level kwarg on every supported openai-agents version, and the
     # channel admits values outside the OpenAI enum ("none").
@@ -338,12 +308,12 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str,
     if reasoning_effort is not None:
         extra_args = {**(extra_args or {}),
                       "reasoning_effort": reasoning_effort}
-    conn = dict(backend) if backend else {}
+    conn = _sdk_backend(backend) if backend else {}
     if conn and protocol == "chat":
         # LiteLLM takes connection params per call, except the two names
         # LitellmModel pins as its own keywords — those ride its constructor.
         lifted = {"api_key": conn.pop("api_key", None),
-                  "base_url": conn.pop("base_url", conn.pop("api_base", None))}
+                  "base_url": conn.pop("base_url", None)}
         if conn:
             extra_args = {**(extra_args or {}), **conn}
         conn = {key: value for key, value in lifted.items()
@@ -365,6 +335,9 @@ def _openai_agent(client, protocol: str, model_name: str, instructions: str,
         model_settings=ModelSettings(
             temperature=temperature, top_p=top_p, max_tokens=max_tokens,
             reasoning=reasoning,
+            # Streamed runs otherwise carry no usage at all (agents forwards
+            # this as stream_options only on streaming calls).
+            include_usage=True,
             extra_body=body,
             extra_headers=extra_headers,
             extra_args=extra_args),
@@ -377,17 +350,21 @@ def _validate_max_turns(max_turns) -> None:
         raise PageIndexAPIError("max_turns must be a positive integer.")
 
 
-def _conversation_cache_key(model_name: str, instructions: str, items) -> str:
+def _conversation_cache_key(model_name: str, instructions: str, doc_id,
+                            items) -> str:
     """Stable per-conversation cache-routing key, sent as the OpenAI
     ``prompt_cache_key`` through ModelSettings.extra_body (openai-agents
     0.20 no longer derives it from RunConfig.group_id — verified against a
     captured wire). Keyed on the prefix identity — model, instructions,
-    first conversation item — so a conversation's continuations share one
-    route without pooling unrelated conversations. Callers pass the
-    conversation's own items, never the SDK-prepended doc-targeting block:
-    that block is byte-identical for every conversation about a document
-    and would pool them all under one key."""
-    seed = json.dumps([model_name, instructions,
+    doc targeting, first conversation item — so a conversation's
+    continuations share one route without pooling unrelated conversations.
+    Callers pass the conversation's own items, never the SDK-prepended
+    doc-targeting block: that block is byte-identical for every
+    conversation about a document and would pool them all under one key.
+    doc_id carries the targeting identity instead — the same opening
+    question against different documents is different conversations."""
+    scope = [doc_id] if isinstance(doc_id, str) else doc_id
+    seed = json.dumps([model_name, instructions, scope,
                        items[0] if items else None],
                       sort_keys=True, default=str)
     return "pageindex-" + hashlib.sha256(seed.encode()).hexdigest()[:16]
@@ -437,16 +414,72 @@ def _record_response_status(agent, recorded: dict) -> None:
                 value = getattr(response, field, None)
                 recorded[field] = (value.model_dump(mode="json")
                                    if hasattr(value, "model_dump") else value)
+        # The backend's echo of what it actually ran with.
+        for field in ("tool_choice", "parallel_tool_calls"):
+            value = getattr(response, field, None)
+            if value is not None:
+                recorded[field] = (value.model_dump(mode="json")
+                                   if hasattr(value, "model_dump") else value)
         return response
 
     responses.create = recording_create
 
 
+def _record_chat_finish(agent, recorded: dict) -> None:
+    """Capture each turn's native finish_reason from the raw LiteLLM
+    response: openai-agents' ModelResponse drops it, so a truncated or
+    content-filtered final turn would otherwise report as a clean "stop".
+    The chat protocol has no transport client to hook (cf.
+    _record_response_status), so this wraps the model's response fetch;
+    no-op if that private seam moves."""
+    model = getattr(agent, "model", None)
+    fetch = getattr(model, "_fetch_response", None)
+    if fetch is None:
+        return
+
+    def note(item) -> None:
+        choices = getattr(item, "choices", None)
+        finish = getattr(choices[0], "finish_reason", None) if choices else None
+        if finish:
+            recorded["finish_reason"] = finish
+
+    class _Tee:
+        """Iteration passthrough that notes each chunk; everything else
+        (aclose/close/...) delegates to the provider stream itself."""
+
+        def __init__(self, inner):
+            self._inner = inner
+
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            chunk = await self._inner.__anext__()
+            note(chunk)
+            return chunk
+
+        def __getattr__(self, name):
+            return getattr(self._inner, name)
+
+    async def recording_fetch(*args, **kwargs):
+        result = await fetch(*args, **kwargs)
+        if isinstance(result, tuple):
+            response, stream = result
+            return response, _Tee(stream)
+        note(result)
+        return result
+
+    model._fetch_response = recording_fetch
+
+
 async def _aclose_backend(agent) -> None:
     """Close the per-call AsyncOpenAI client before its event loop ends —
     otherwise httpx tears down pooled connections on a closed loop and
-    emits 'Task exception was never retrieved' noise."""
+    emits 'Task exception was never retrieved' noise. A client built on a
+    caller-owned http_client stays open."""
     backend = getattr(getattr(agent, "model", None), "_client", None)
+    if getattr(backend, "_pageindex_caller_http", False):
+        return
     close = getattr(backend, "close", None)
     if close is not None:
         try:
@@ -473,8 +506,10 @@ def _wrap_max_turns(max_turns) -> PageIndexAPIError:
 def _usage_sums(raw_responses) -> "tuple[int, int, int, int, int]":
     prompt = completion = cached = cache_write = reasoning = 0
     for r in raw_responses:
-        prompt += r.usage.input_tokens
-        completion += r.usage.output_tokens
+        if r.usage is None:
+            continue
+        prompt += r.usage.input_tokens or 0
+        completion += r.usage.output_tokens or 0
         details = getattr(r.usage, "input_tokens_details", None)
         cached += getattr(details, "cached_tokens", 0) or 0
         cache_write += getattr(details, "cache_write_tokens", 0) or 0
@@ -532,12 +567,14 @@ def run_chat_completions(client, messages, stream: bool = False,
     managed = _managed_instructions(system_texts)
     agent = _openai_agent(client, "chat", model_name, managed,
                           temperature, top_p, doc_ids=doc_id,
-                          cache_key=_conversation_cache_key(model_name,
-                                                            managed, history),
+                          cache_key=_conversation_cache_key(
+                              model_name, managed, doc_id, history),
                           reasoning_effort=reasoning_effort,
                           extra_body=extra_body, max_tokens=max_tokens,
                           backend=_merged_backend(client, backend),
                           extra_headers=extra_headers)
+    recorded: dict = {}
+    _record_chat_finish(agent, recorded)
     run_kwargs = _run_kwargs(max_turns)
     import openai
     from agents import Runner
@@ -562,7 +599,7 @@ def run_chat_completions(client, messages, stream: bool = False,
                 "index": 0,
                 "message": {"role": "assistant",
                             "content": result.final_output or ""},
-                "finish_reason": "stop",
+                "finish_reason": recorded.get("finish_reason") or "stop",
             }],
             "usage": _openai_usage(result.raw_responses),
         }
@@ -602,7 +639,7 @@ def run_chat_completions(client, messages, stream: bool = False,
             if not completed and hasattr(streamed, "cancel"):
                 streamed.cancel()  # abandoned/failed: stop the agent task
             await _aclose_backend(agent)
-        yield chunk({}, finish="stop")
+        yield chunk({}, finish=recorded.get("finish_reason") or "stop")
         yield {
             "id": chat_id, "object": "chat.completion.chunk",
             "created": created, "model": reported_model, "choices": [],
@@ -649,8 +686,8 @@ def run_responses(client, input, model: Optional[str] = None,
     managed = _managed_instructions(extra)
     agent = _openai_agent(client, "responses", model_name, managed,
                           temperature, top_p, doc_ids=doc_id,
-                          cache_key=_conversation_cache_key(model_name, managed,
-                                                            conversation),
+                          cache_key=_conversation_cache_key(
+                              model_name, managed, doc_id, conversation),
                           reasoning=reasoning, extra_body=extra_body,
                           max_tokens=max_output_tokens,
                           backend=_merged_backend(client, backend),
@@ -661,11 +698,14 @@ def run_responses(client, input, model: Optional[str] = None,
     from agents import Runner
     from agents.exceptions import AgentsException, MaxTurnsExceeded
 
+    response_id = f"resp_{uuid.uuid4().hex}"
+    created_at = int(time.time())
+
     def envelope(transcript: list, raw_responses) -> dict:
         return {
-            "id": f"resp_{uuid.uuid4().hex}",
+            "id": response_id,
             "object": "response",
-            "created_at": int(time.time()),
+            "created_at": created_at,
             "model": _reported_model(model_name),
             "status": recorded.get("status") or "completed",
             "output": [item for item in transcript
@@ -678,8 +718,9 @@ def run_responses(client, input, model: Optional[str] = None,
                        "parameters": tool.params_json_schema,
                        "strict": getattr(tool, "strict_json_schema", True)}
                       for tool in agent.tools],
-            "tool_choice": "auto",
-            "parallel_tool_calls": True,
+            # Backend echo when captured; the request sends neither param.
+            "tool_choice": recorded.get("tool_choice", "auto"),
+            "parallel_tool_calls": recorded.get("parallel_tool_calls", True),
             "temperature": temperature,
             "top_p": top_p,
             "reasoning": reasoning,
@@ -721,11 +762,24 @@ def run_responses(client, input, model: Optional[str] = None,
         # re-based by the count of items already committed by prior turns.
         output_offset = 0
         completed = False
+        opened = False
         try:
             async for event in streamed.stream_events():
                 if event.type == "raw_response_event":
                     data = event.data.model_dump(exclude_unset=True)
                     if data.get("type") in lifecycle:
+                        if data["type"] == "response.created" and not opened:
+                            # N per-turn openings collapse to one, carrying
+                            # the id and created_at the terminal event will
+                            # report.
+                            opened = True
+                            if data.get("response"):
+                                data["response"]["id"] = response_id
+                                data["response"]["created_at"] = created_at
+                            sequence += 1
+                            data["sequence_number"] = sequence
+                            yield data
+                            continue
                         if data["type"] in ("response.completed",
                                             "response.incomplete",
                                             "response.failed"):
@@ -735,6 +789,10 @@ def run_responses(client, input, model: Optional[str] = None,
                             for field in ("status", "incomplete_details",
                                           "error"):
                                 recorded[field] = state.get(field)
+                            for field in ("tool_choice",
+                                          "parallel_tool_calls"):
+                                if state.get(field) is not None:
+                                    recorded[field] = state[field]
                             output_offset += len(state.get("output") or [])
                         continue
                     if isinstance(data.get("output_index"), int):
@@ -793,7 +851,7 @@ def _anthropic_client(backend=None):
     """The backend client — the seam tests replace with a fake transport."""
     import anthropic
     try:
-        return anthropic.Anthropic(**(backend or {}))
+        return anthropic.Anthropic(**_sdk_backend(backend))
     except TypeError as exc:
         raise PageIndexAPIError(
             f"The Anthropic backend is not configured: {exc}") from exc
@@ -820,6 +878,18 @@ def _anthropic_system(extra_system, block: Optional[str]) -> list[dict]:
     raise PageIndexAPIError("system must be a string or a list of blocks.")
 
 
+def _cache_marks(system_blocks, messages) -> int:
+    """Breakpoints already on the request. The API allows 4 total; the
+    top-level moving breakpoint is only added when it fits."""
+    blocks = list(system_blocks)
+    for message in messages:
+        content = message.get("content")
+        if isinstance(content, list):
+            blocks += [b for b in content if isinstance(b, dict)]
+    return sum(1 for b in blocks
+               if isinstance(b, dict) and b.get("cache_control"))
+
+
 def _dump_block(block) -> Any:
     """A content block as a plain JSON dict, minus SDK-internal fields the
     API rejects (ParsedBetaTextBlock.__api_exclude__, e.g. parsed_output)."""
@@ -855,9 +925,15 @@ _CLAUDE_4096_MODELS = ("claude-3-opus", "claude-3-sonnet", "claude-3-haiku",
                        "claude-3-5-sonnet-20240620")
 
 
-def _default_max_tokens(model: str) -> int:
+def _default_max_tokens(model: str, thinking=None) -> int:
     """The wire-required per-turn budget when the caller sets none: 8192,
-    except the claude-3 generation whose output ceiling is 4096."""
+    except the claude-3 generation whose output ceiling is 4096. The wire
+    also requires max_tokens > thinking.budget_tokens, so an enabled
+    budget lifts the default above itself."""
+    budget = (thinking.get("budget_tokens")
+              if isinstance(thinking, dict) else None)
+    if isinstance(budget, int):
+        return budget + 8192
     return 4096 if model.startswith(_CLAUDE_4096_MODELS) else 8192
 
 
@@ -892,18 +968,38 @@ def run_messages(client, messages, model: str,
         "stop_sequences": stop_sequences, "thinking": thinking,
         "extra_body": extra_body, "extra_headers": extra_headers,
     }.items() if value is not None}
-    runner = _anthropic_client(_merged_backend(client, backend)) \
-        .beta.messages.tool_runner(
-        max_tokens=(max_tokens if max_tokens is not None
-                    else _default_max_tokens(model)),
+    system_blocks = _anthropic_system(system, block)
+    # Top-level cache_control: the server re-marks the newest block each
+    # turn, so the loop re-reads the growing conversation from cache.
+    # Counts toward the 4-breakpoint limit (live-verified 400 past it).
+    cached: dict[str, Any] = (
+        {"cache_control": {"type": "ephemeral"}}
+        if _cache_marks(system_blocks, prepared) < 4 else {})
+    merged = _merged_backend(client, backend)
+    # The SDK defers credential resolution to request time and raises a
+    # bare TypeError there — pre-check for the contract's PageIndexAPIError.
+    if not merged and not (os.environ.get("ANTHROPIC_API_KEY")
+                           or os.environ.get("ANTHROPIC_AUTH_TOKEN")):
+        raise PageIndexAPIError(
+            "The Anthropic backend is not configured: set the "
+            "ANTHROPIC_API_KEY environment variable, or pass an api_key "
+            "in chat_backend / backend.")
+    backend_client = _anthropic_client(merged)
+    # A caller-owned http_client must survive the per-call closes below.
+    owns_transport = "http_client" not in (merged or {})
+    if max_tokens is None:
+        max_tokens = _default_max_tokens(model, thinking)
+    runner = backend_client.beta.messages.tool_runner(
+        max_tokens=max_tokens,
         messages=prepared,
         model=model,
         tools=build_anthropic_tools(client, doc_ids=doc_id),
-        system=_anthropic_system(system, block),
+        system=system_blocks,
         stream=stream,
         # Bounded like the OpenAI surfaces (their framework default is 10).
         max_iterations=max_turns if max_turns is not None else 10,
         **passthrough,
+        **cached,
     )
 
     if stream:
@@ -915,6 +1011,10 @@ def run_messages(client, messages, model: str,
             except anthropic.AnthropicError as exc:
                 raise PageIndexAPIError(
                     f"The model backend failed: {exc}") from exc
+            finally:
+                # runs on exhaustion and abandonment (GeneratorExit) alike
+                if owns_transport:
+                    backend_client.close()
         return events()
 
     try:
@@ -922,6 +1022,10 @@ def run_messages(client, messages, model: str,
     except anthropic.AnthropicError as exc:
         raise PageIndexAPIError(
             f"The model backend failed: {exc}") from exc
+    finally:
+        # safe here: the params read-back below does no HTTP
+        if owns_transport:
+            backend_client.close()
     if not turns:
         raise PageIndexAPIError("The model returned no response.")
     captured: dict = {}
@@ -955,7 +1059,7 @@ def run_messages(client, messages, model: str,
     new_messages = [_dump_message(message)
                     for message in conversation[len(prepared):]]
     final_blocks = [_dump_block(item) for item in final.content]
-    final_ids = {block["id"] for block in final_blocks
+    final_ids = {block.get("id") for block in final_blocks
                  if block.get("type") == "tool_use"}
     history_ids = {block.get("id")
                    for message in new_messages

pageindex/local_store.py

@@ -6,6 +6,7 @@ import logging
 import os
 import shutil
 import uuid
+from contextlib import contextmanager
 from pathlib import Path
 
 logger = logging.getLogger(__name__)
@@ -89,6 +90,24 @@ class DocStore:
         except OSError:
             pass
 
+    @contextmanager
+    def lock(self):
+        """Cross-process mutex for check-then-write sequences (name
+        uniquing before save). fcntl is absent on Windows, where the
+        pre-existing best-effort behavior stays."""
+        try:
+            import fcntl
+        except ImportError:
+            yield
+            return
+        self._root.mkdir(parents=True, exist_ok=True)
+        with open(self._root / ".lock", "w") as handle:
+            fcntl.flock(handle, fcntl.LOCK_EX)
+            try:
+                yield
+            finally:
+                fcntl.flock(handle, fcntl.LOCK_UN)
+
     # ── documents ──
     def save_document(self, doc_id: str, meta: dict, tree: list, pages: list) -> None:
         doc_dir = self._doc_dir(doc_id)

pageindex/mcp_bridge.py

@@ -44,6 +44,7 @@ class McpBridge:
     def __init__(self, url: str, headers: dict[str, str]):
         self._url = url
         self._auth_headers = dict(headers)
+        self._session = requests.Session()  # agent tool calls come in bursts
         self._session_id: Optional[str] = None
         self._protocol_version: Optional[str] = None
         self._instructions: Optional[str] = None
@@ -65,8 +66,8 @@ class McpBridge:
         if protocol_version:
             headers["MCP-Protocol-Version"] = protocol_version
         try:
-            return requests.post(self._url, json=payload, headers=headers,
-                                 timeout=_TIMEOUT)
+            return self._session.post(self._url, json=payload,
+                                      headers=headers, timeout=_TIMEOUT)
         except requests.RequestException as exc:
             raise PageIndexAPIError(
                 f"Could not reach the PageIndex MCP server: {exc}"
@@ -85,7 +86,8 @@ class McpBridge:
             except ValueError as exc:
                 raise PageIndexAPIError(
                     f"MCP server returned a non-JSON response "
-                    f"(HTTP {response.status_code})."
+                    f"(HTTP {response.status_code}).",
+                    status_code=response.status_code,
                 ) from exc
         # Strict id correlation only — accepting any result-bearing message
         # would return a stale or mis-correlated reply as this call's.
@@ -129,7 +131,8 @@ class McpBridge:
         if response.status_code >= 400:
             raise PageIndexAPIError(
                 f"MCP request failed: HTTP {response.status_code} "
-                f"({response.text[:200]})"
+                f"({response.text[:200]})",
+                status_code=response.status_code,
             )
         return self._extract_result(response, request_id)
 
@@ -152,7 +155,8 @@ class McpBridge:
                 raise PageIndexAPIError(
                     f"Could not connect to the PageIndex MCP server: HTTP "
                     f"{response.status_code} ({response.text[:200]}). Check "
-                    "your API key."
+                    "your API key.",
+                    status_code=response.status_code,
                 )
             result = self._extract_result(response, request_id) or {}
             self._session_id = response.headers.get("Mcp-Session-Id")
@@ -179,13 +183,18 @@ class McpBridge:
     def list_tools(self) -> list[dict]:
         tools: list[dict] = []
         cursor: Optional[str] = None
-        while True:
+        # A server echoing its cursor (or cycling) must not hang the client:
+        # no-progress terminates, the page cap turns a cycle into an error.
+        for _ in range(50):
             params = {"cursor": cursor} if cursor else {}
             result = self._request("tools/list", params) or {}
             tools.extend(result.get("tools") or [])
-            cursor = result.get("nextCursor")
-            if not cursor:
+            next_cursor = result.get("nextCursor")
+            if not next_cursor or next_cursor == cursor:
                 return tools
+            cursor = next_cursor
+        raise PageIndexAPIError(
+            "MCP tools/list pagination did not terminate within 50 pages.")
 
     def call_tool(self, name: str, arguments: dict[str, Any]) -> "tuple[str, bool]":
         """Returns (text, is_error) — is_error is the server's MCP isError
@@ -205,6 +214,14 @@ class McpBridge:
                 kind = block.get("mimeType") or block.get("type") or "binary"
                 size_kb = max(1, len(block["data"]) * 3 // 4096)
                 texts.append(f"[{kind} content omitted: ~{size_kb} KB]")
+            elif (isinstance(block, dict)
+                  and isinstance(block.get("resource"), dict)
+                  and isinstance(block["resource"].get("blob"), str)):
+                # EmbeddedResource nests its base64 one level down.
+                resource = block["resource"]
+                kind = resource.get("mimeType") or "binary"
+                size_kb = max(1, len(resource["blob"]) * 3 // 4096)
+                texts.append(f"[{kind} content omitted: ~{size_kb} KB]")
             else:
                 texts.append(json.dumps(block, ensure_ascii=False))
         return "\n".join(texts), is_error

pageindex/tree_optimize.py

@@ -61,7 +61,7 @@ import re
 import sys
 from types import SimpleNamespace
 
-from .utils import (ConfigLoader, _is_openai_model, _is_unrecoverable,
+from .utils import (ConfigLoader, _is_unrecoverable, _openai_missing_keys,
                     llm_acompletion, strip_internal_keys)
 
 TRIGGER_PAGES = 5        # only look ahead on nodes larger than this
@@ -872,9 +872,11 @@ async def main():
     args = parser.parse_args()
 
     model = args.model or default_model()
-    if args.expand and not args.plan and _is_openai_model(model) \
-            and not os.getenv("OPENAI_API_KEY"):
-        sys.exit(f"OPENAI_API_KEY is not set (expand model: {model}).")
+    if args.expand and not args.plan:
+        missing = _openai_missing_keys(model)
+        if missing:
+            sys.exit(f"{', '.join(missing)} is not set "
+                     f"(expand model: {model}).")
 
     original = json.load(open(args.structure))
     structure = copy.deepcopy(original["structure"])

pageindex/utils.py

@@ -29,13 +29,6 @@ _llm_backend: contextvars.ContextVar = contextvars.ContextVar(
     "pageindex_llm_backend", default=None)
 
 
-def _openai_sdk_kwargs(backend: dict) -> dict:
-    """The same backend dict works on both gateway paths: LiteLLM accepts
-    either endpoint spelling, the openai SDK only ``base_url``."""
-    return {("base_url" if key == "api_base" else key): value
-            for key, value in backend.items()}
-
-
 def _repair_litellm_types() -> None:
     """litellm 1.97.0's Message/Delta annotations carry nested forward refs
     Python 3.10 cannot resolve (BerriAI/litellm#36384), so every completion
@@ -54,6 +47,9 @@ def _repair_litellm_types() -> None:
 
 # Backward compatibility: support CHATGPT_API_KEY as alias for OPENAI_API_KEY
 if not os.getenv("OPENAI_API_KEY") and os.getenv("CHATGPT_API_KEY"):
+    import warnings
+    warnings.warn("CHATGPT_API_KEY is deprecated — set OPENAI_API_KEY "
+                  "instead.", FutureWarning)
     os.environ["OPENAI_API_KEY"] = os.getenv("CHATGPT_API_KEY")
 
 def count_tokens(text, model=None):
@@ -69,16 +65,67 @@ def _strip_prefix(s, prefix):
     return s
 
 
-def _is_openai_model(model):
-    """Models without a provider prefix (no '/') use the openai SDK directly.
-    For other providers, use 'provider/model' format (e.g. 'anthropic/claude-sonnet-4-6')."""
-    if not model or model.startswith('litellm/'):
-        return False
-    return '/' not in model or model.startswith('openai/')
+def run_off_loop(func, *args):
+    """Run func now, or on a worker thread when this thread already runs an
+    asyncio loop (func may itself call asyncio.run)."""
+    try:
+        asyncio.get_running_loop()
+    except RuntimeError:
+        return func(*args)
+    from concurrent.futures import ThreadPoolExecutor
+    with ThreadPoolExecutor(max_workers=1) as pool:
+        return pool.submit(func, *args).result()
 
 
-_openai_sync_client = None
-_openai_async_client = None
+def _openai_missing_keys(model):
+    """Missing env keys for the pre-check, which covers only OpenAI-shaped
+    names (bare or ``openai/``): other providers resolve credentials their
+    own way at call time (IAM chains, ADC, Ollama's localhost default),
+    invisible to env inspection — the chat lane draws the same line.
+    ``litellm/``-prefixed names are exempt: the prefix is an explicit
+    routing choice, and litellm resolves credentials beyond the
+    environment (litellm.api_key, a keyless OPENAI_BASE_URL server).
+    Truthiness, not litellm's validate_environment, which reports a blank
+    exported key as present."""
+    if model.startswith("litellm/"):
+        return []
+    if "/" in model and not model.startswith("openai/"):
+        return []
+    return ([] if (os.getenv("OPENAI_API_KEY") or "").strip()
+            else ["OPENAI_API_KEY"])
+
+
+def _litellm_model(model, backend):
+    """Normalize to LiteLLM's grammar (``litellm/`` strips, bare names get
+    the ``openai/`` wire form — same as the chat lane) and fail fast on a
+    missing key or unknown provider, with status codes the retry loop and
+    the summary/optimize passes treat as unrecoverable."""
+    if not model:
+        return model
+    raw = model
+    model = _strip_prefix(model, "litellm/")
+    if "/" not in model:
+        model = f"openai/{model}"
+    import litellm
+    provider = model.split("/", 1)[0]
+    providers = getattr(litellm, "provider_list", None)
+    # custom_provider_map providers join provider_list only at call time.
+    custom = {entry.get("provider") for entry
+              in getattr(litellm, "custom_provider_map", None) or []}
+    if providers and provider not in providers and provider not in custom:
+        raise litellm.NotFoundError(
+            f"'{model}' routes through LiteLLM, but '{provider}' is not a "
+            f"LiteLLM provider. For an OpenAI-compatible server serving "
+            f"this model id, use 'openai/{model}' and point "
+            f"OPENAI_BASE_URL at the server.",
+            llm_provider=None, model=model)
+    if not backend:
+        missing = _openai_missing_keys(raw)
+        if missing:
+            raise litellm.AuthenticationError(
+                f"missing API key for {model}: {', '.join(missing)}",
+                llm_provider=None, model=model)
+    return model
 
 
 # Misconfiguration: no retry can fix a rejected key or a model that does not
@@ -92,42 +139,34 @@ def _is_unrecoverable(exc: Exception) -> bool:
     return getattr(exc, "status_code", None) in _UNRECOVERABLE_STATUS
 
 
+def _no_cache_seeding_kwargs(backend):
+    """litellm 1.97 auto-marks Claude requests for prompt caching (system +
+    last message); indexing prompts are single-shot and unique, so every call
+    would pay the cache-write premium with nothing ever read back. A
+    system-role-only injection point matches no indexing message, and its
+    presence stops litellm seeding its own defaults; backend keys still
+    win."""
+    return {"cache_control_injection_points":
+                [{"location": "message", "role": "system"}],
+            **(backend or {})}
+
+
 def llm_completion(model, prompt, chat_history=None, return_finish_reason=False):
-    use_openai_sdk = _is_openai_model(model)
-    if model:
-        model = _strip_prefix(model, "litellm/")
-        if use_openai_sdk:
-            model = _strip_prefix(model, "openai/")
+    import litellm
     max_retries = 10
     messages = list(chat_history) + [{"role": "user", "content": prompt}] if chat_history else [{"role": "user", "content": prompt}]
     backend = _llm_backend.get()
-    if use_openai_sdk:
-        import openai
-        if backend:
-            oai_client = openai.OpenAI(**{"max_retries": 0,
-                                          **_openai_sdk_kwargs(backend)})
-        else:
-            global _openai_sync_client
-            if _openai_sync_client is None:
-                _openai_sync_client = openai.OpenAI(max_retries=0)
-            oai_client = _openai_sync_client
+    model = _litellm_model(model, backend)
+    _repair_litellm_types()
     for i in range(max_retries):
         try:
-            if use_openai_sdk:
-                response = oai_client.chat.completions.create(
-                    model=model,
-                    messages=messages,
-                )
-            else:
-                import litellm
-                _repair_litellm_types()
-                response = litellm.completion(
-                    model=model,
-                    messages=messages,
-                    temperature=0,
-                    drop_params=True,
-                    **(backend or {}),
-                )
+            response = litellm.completion(
+                model=model,
+                messages=messages,
+                drop_params=True,
+                # the loop is the retry policy; the merge lets a backend override win
+                **{"max_retries": 0, **_no_cache_seeding_kwargs(backend)},
+            )
             content = response.choices[0].message.content
             if return_finish_reason:
                 finish_reason = "max_output_reached" if response.choices[0].finish_reason == "length" else "finished"
@@ -147,41 +186,20 @@ def llm_completion(model, prompt, chat_history=None, return_finish_reason=False)
 
 
 async def llm_acompletion(model, prompt):
-    use_openai_sdk = _is_openai_model(model)
-    if model:
-        model = _strip_prefix(model, "litellm/")
-        if use_openai_sdk:
-            model = _strip_prefix(model, "openai/")
+    import litellm
     max_retries = 10
     messages = [{"role": "user", "content": prompt}]
     backend = _llm_backend.get()
-    if use_openai_sdk:
-        import openai
-        if backend:
-            oai_client = openai.AsyncOpenAI(**{"max_retries": 0,
-                                               **_openai_sdk_kwargs(backend)})
-        else:
-            global _openai_async_client
-            if _openai_async_client is None:
-                _openai_async_client = openai.AsyncOpenAI(max_retries=0)
-            oai_client = _openai_async_client
+    model = _litellm_model(model, backend)
+    _repair_litellm_types()
     for i in range(max_retries):
         try:
-            if use_openai_sdk:
-                response = await oai_client.chat.completions.create(
-                    model=model,
-                    messages=messages,
-                )
-            else:
-                import litellm
-                _repair_litellm_types()
-                response = await litellm.acompletion(
-                    model=model,
-                    messages=messages,
-                    temperature=0,
-                    drop_params=True,
-                    **(backend or {}),
-                )
+            response = await litellm.acompletion(
+                model=model,
+                messages=messages,
+                drop_params=True,
+                **{"max_retries": 0, **_no_cache_seeding_kwargs(backend)},
+            )
             return response.choices[0].message.content
         except Exception as e:
             if _is_unrecoverable(e):
@@ -714,6 +732,8 @@ async def generate_summaries_for_structure(structure, model=None):
     summaries = await asyncio.gather(*tasks, return_exceptions=True)
 
     for node, summary in zip(nodes, summaries):
+        if isinstance(summary, Exception) and _is_unrecoverable(summary):
+            raise summary
         node['summary'] = "" if isinstance(summary, BaseException) else summary
     if nodes and not any(node['summary'] for node in nodes):
         raise RuntimeError(
@@ -887,17 +907,25 @@ async def summarize_tree(structure, pdf_pages, model=None,
     async def visit(node):
         children = node.get('nodes') or []
         if children:
-            await asyncio.gather(*(visit(child) for child in children),
-                                 return_exceptions=True)
+            done = await asyncio.gather(*(visit(child) for child in children),
+                                        return_exceptions=True)
+            for result in done:
+                if isinstance(result, Exception) and _is_unrecoverable(result):
+                    raise result
         if node.get('summary'):
             return
         try:
             node['summary'] = await (parent_summary(node) if children else leaf_summary(node))
-        except Exception:
+        except Exception as e:
             node['summary'] = ""
+            if _is_unrecoverable(e):
+                raise
 
-    await asyncio.gather(*(visit(root) for root in structure),
-                         return_exceptions=True)
+    results = await asyncio.gather(*(visit(root) for root in structure),
+                                    return_exceptions=True)
+    for r in results:
+        if isinstance(r, Exception) and _is_unrecoverable(r):
+            raise r
 
     def _any_summary(nodes):
         return any(n.get('summary') or _any_summary(n.get('nodes') or [])

pyproject.toml

@@ -32,7 +32,7 @@ openai = ">=1.70.0"
 openai-agents = ">=0.18.1"
 litellm = ">=1.97.0"
 PyPDF2 = ">=3.0.0"
-pypdfium2 = ">=4.30.0"
+pypdfium2 = ">=5"
 sortedcontainers = ">=2.4.0"
 regex = ">=2024.0.0"
 python-dotenv = ">=1.0.0"

requirements.txt

@@ -4,7 +4,7 @@ requests>=2.28.0
 openai-agents>=0.18.1
 # pymupdf  # optional
 PyPDF2==3.0.1
-pypdfium2==4.30.0
+pypdfium2==5.13.0
 python-dotenv==1.2.2
 pyyaml==6.0.2
 regex>=2024.0.0

run_pageindex.py

@@ -3,7 +3,11 @@ import os
 import json
 from pageindex import *
 from pageindex.page_index_md import md_to_tree
-from pageindex.utils import ConfigLoader
+from pageindex.utils import ConfigLoader, _openai_missing_keys
+
+# Keep LiteLLM's import off the network (frozen bundled model-cost map);
+# an explicit user setting wins.
+os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "True")
 
 if __name__ == "__main__":
     # Set up argument parser
@@ -62,7 +66,7 @@ if __name__ == "__main__":
         raise ValueError("Either --pdf_path or --md_path must be specified")
     if args.pdf_path and args.md_path:
         raise ValueError("Only one of --pdf_path or --md_path can be specified")
-    if args.optimize in ('full', 'merge') and not (args.pdf_path and args.mode == 'flash'):
+    if args.optimize is not None and not (args.pdf_path and args.mode == 'flash'):
         raise ValueError("--optimize requires Flash mode with --pdf_path")
     if args.optimize is None:
         args.optimize = 'full' if args.mode == 'flash' else 'off'
@@ -90,14 +94,15 @@ if __name__ == "__main__":
             
         if args.mode == 'flash':
             from pageindex.flash import page_index_flash
-            summary_model = args.summary_model or args.index_model or args.model
+            summary_model = (args.summary_model or args.index_model
+                             or args.model
+                             or ConfigLoader().load().summary_model)
             will_summarize = args.summary if args.summary is not None else True
-            if summary_model and (will_summarize or args.optimize == 'full'):
-                import litellm
-                env = litellm.validate_environment(summary_model)
-                if not env["keys_in_environment"]:
+            if will_summarize or args.optimize == 'full':
+                missing = _openai_missing_keys(summary_model)
+                if missing:
                     raise SystemExit(
-                        f"Missing API key for {summary_model}: {', '.join(env['missing_keys'])}")
+                        f"Missing API key for {summary_model}: {', '.join(missing)}")
             toc_with_page_number = page_index_flash(
                 args.pdf_path,
                 optimize=args.optimize if args.optimize != 'off' else False,
@@ -164,25 +169,24 @@ if __name__ == "__main__":
         user_opt = {
             'index_model': args.index_model,
             'model': args.model,
-            'if_add_node_summary': args.if_add_node_summary,
-            'if_add_doc_description': args.if_add_doc_description,
-            'if_add_node_text': args.if_add_node_text,
-            'if_add_node_id': args.if_add_node_id
         }
         
         # Load config with defaults from config.yaml
-        opt = config_loader.load(user_opt)
+        opt = config_loader.load({k: v for k, v in user_opt.items() if v is not None})
         
+        # if_add_* pass through as given (absent = off, as before this CLI
+        # used config.yaml): the PDF defaults there must not switch on LLM
+        # passes the markdown CLI never ran.
         toc_with_page_number = asyncio.run(md_to_tree(
             md_path=args.md_path,
             if_thinning=args.if_thinning.lower() == 'yes',
             min_token_threshold=args.thinning_threshold,
-            if_add_node_summary=opt.if_add_node_summary,
+            if_add_node_summary=args.if_add_node_summary,
             summary_token_threshold=args.summary_token_threshold,
             model=opt.model,
-            if_add_doc_description=opt.if_add_doc_description,
-            if_add_node_text=opt.if_add_node_text,
-            if_add_node_id=opt.if_add_node_id
+            if_add_doc_description=args.if_add_doc_description,
+            if_add_node_text=args.if_add_node_text,
+            if_add_node_id=args.if_add_node_id
         ))
         
         print('Parsing done, saving to file...')

tests/conftest.py

@@ -5,6 +5,7 @@ import pytest
 def _llm_key(monkeypatch):
     """Deterministic key presence for every test; missing-key tests delenv."""
     monkeypatch.setenv("OPENAI_API_KEY", "test-key")
+    monkeypatch.setenv("ANTHROPIC_API_KEY", "test-key")
 
 
 def build_pdf(page_texts):

tests/test_agent_tools.py

@@ -69,7 +69,10 @@ def run(client, name, **arguments):
 
 # ── contract parity ──
 
-def test_contract_matches_snapshot():
+def test_contract_edits_are_deliberate():
+    """The committed snapshot cannot detect drift from the live cloud
+    server — both copies live in this repo. It exists so a TOOL_CONTRACT
+    edit must touch two files in one change, never land by accident."""
     snapshot = json.loads(SNAPSHOT_PATH.read_text(encoding="utf-8"))
     assert snapshot["tools"] == TOOL_CONTRACT
 
@@ -405,9 +408,11 @@ def test_created_at_accepts_z_suffixed_input(client, store_path):
 
 
 def test_page_content_char_budget(client, store_path):
+    # Escape-dense pages: raw length fits the budget, JSON-serialized
+    # length does not — the budget must count emitted characters.
     pages = [
-        {"page_index": 1, "markdown": "x" * 96_000},
-        {"page_index": 2, "markdown": "short"},
+        {"page_index": 1, "markdown": '"' * 30_000},
+        {"page_index": 2, "markdown": '"' * 20_000},
     ]
     seed_doc(store_path, "pi-a", "huge.pdf", pages=pages)
     payload, is_error = run(client, "get_page_content", doc_name="huge.pdf",
@@ -417,6 +422,8 @@ def test_page_content_char_budget(client, store_path):
     assert "size limits" in payload["next_steps"]["summary"]
     assert any("For remaining pages, request: 2" in option
                for option in payload["next_steps"]["options"])
+    emitted = json.dumps(payload, ensure_ascii=False)
+    assert len(emitted) <= agent_tools_module.TOOL_RESPONSE_CHAR_LIMIT
 
 
 def test_page_content_reports_truncation_and_out_of_range_together(
@@ -763,6 +770,54 @@ def test_openai_agent_config_model_speaks_the_agents_sdk_grammar(tmp_path):
                                   retrieve_model="anthropic/claude-x")
     assert (client.openai_agent_config()["model"]
             == "litellm/anthropic/claude-x")
+    # The per-call override speaks the same grammar as chat_model.
+    assert (client.openai_agent_config(model="anthropic/claude-y")["model"]
+            == "litellm/anthropic/claude-y")
+    assert (client.openai_agent_config(model="litellm/groq/llama-x")["model"]
+            == "litellm/groq/llama-x")
+
+
+def test_plain_functions_answer_bad_arguments_with_the_envelope(client,
+                                                                store_path):
+    """agent_tools() functions must not raise into a framework loop:
+    cloud-only parameters pruned from the local signature come back as
+    the guided envelope, and the schema-bearing signature survives."""
+    import inspect
+    seed_doc(store_path, "pi-a", "report.pdf")
+    tools = {f.__name__: f for f in client.agent_tools()}
+    fn = tools["get_document"]
+    assert "doc_name" in inspect.signature(fn).parameters
+    payload = json.loads(fn(doc_name="report.pdf", folder_id="root"))
+    assert payload["errorCode"] == "INVALID_INPUT"
+    ok = json.loads(fn(doc_name="report.pdf"))
+    assert not ok.get("errorCode")
+
+
+def test_non_string_doc_name_stays_not_found(client, store_path):
+    """A type-loose model argument must not turn NOT_FOUND into the
+    retry-inviting INTERNAL_ERROR (strict_json_schema is off on the
+    OpenAI adapter, so nothing upstream validates the type)."""
+    seed_doc(store_path, "pi-a", "report.pdf")
+    payload, is_error = run(client, "get_document", doc_name=5)
+    assert is_error and payload["errorCode"] == "NOT_FOUND"
+
+
+def test_openai_agent_config_repairs_litellm_types(tmp_path, monkeypatch):
+    """The BYO path resolves its model through LiteLLM in the caller's
+    process, outside our completion helpers — the py3.10 type repair must
+    run at config time, and only for LiteLLM-routed models."""
+    pytest.importorskip("agents")
+    import pageindex.utils
+    calls = []
+    monkeypatch.setattr(pageindex.utils, "_repair_litellm_types",
+                        lambda: calls.append(True))
+    client = PageIndexLocalClient(storage_path=str(tmp_path / "s"),
+                                  chat_model="anthropic/claude-x")
+    client.openai_agent_config()
+    assert calls
+    calls.clear()
+    client.openai_agent_config(model="gpt-plain")
+    assert not calls
 
 
 def test_openai_agent_config_cloud_omits_model(cloud_with_fake_bridge):
@@ -784,6 +839,7 @@ def test_anthropic_runner_config_shapes(client, store_path):
                                             doc_id="pi-a")
     assert config["max_tokens"] == 4096
     assert config["max_iterations"] == 10
+    assert config["cache_control"] == {"type": "ephemeral"}
     assert "report.pdf" in config["system"]
     assert [tool.name for tool in config["tools"]] == list(tool_names())
     assert (client.anthropic_runner_config(model="claude-sonnet-4-5")
@@ -846,22 +902,35 @@ def test_anthropic_runner_config_doc_scope_enforced_in_tools(client,
     assert [doc["name"] for doc in browse["documents"]] == ["report.pdf"]
 
 
-def test_claude_agent_config_doc_scope_enforced_in_tools(client, store_path):
-    pytest.importorskip("claude_agent_sdk")
-    from mcp.types import CallToolRequest, CallToolRequestParams
+def test_claude_agent_config_doc_scope_enforced_in_tools(client, store_path,
+                                                         monkeypatch):
+    """claude_agent_config(doc_id=...) must wire scope all the way into the
+    handlers it registers — an out-of-scope document returns NOT_FOUND. The
+    assertion has to drive those handlers, or build_claude_mcp's doc_ids
+    pass-through goes unguarded."""
+    claude_agent_sdk = pytest.importorskip("claude_agent_sdk")
     seed_doc(store_path, "pi-a", "report.pdf")
     seed_doc(store_path, "pi-b", "payroll.pdf",
              created_at="2026-08-02T10:00:00.123000")
+
+    registered = {}
+    create_server = claude_agent_sdk.create_sdk_mcp_server
+
+    def capture(**kwargs):
+        registered.update(kwargs)
+        return create_server(**kwargs)
+
+    monkeypatch.setattr(claude_agent_sdk, "create_sdk_mcp_server", capture)
     config = client.claude_agent_config(doc_id="pi-a")
-    server = config["mcp_servers"]["pageindex"]
-    handler = server["instance"].request_handlers[CallToolRequest]
-    result = asyncio.run(handler(CallToolRequest(
-        method="tools/call",
-        params=CallToolRequestParams(
-            name="get_page_content",
-            arguments={"doc_name": "payroll.pdf", "pages": "1"}))))
-    payload = json.loads(result.root.content[0].text)
-    assert payload["errorCode"] == "NOT_FOUND"
+    assert "report.pdf" in config["system_prompt"]
+    handlers = {spec.name: spec.handler for spec in registered["tools"]}
+    result = asyncio.run(handlers["get_page_content"](
+        {"doc_name": "payroll.pdf", "pages": "1"}))
+    assert result.get("is_error")
+    assert json.loads(result["content"][0]["text"])["errorCode"] == "NOT_FOUND"
+    browse = asyncio.run(handlers["browse_documents"]({}))
+    listed = json.loads(browse["content"][0]["text"])["documents"]
+    assert [doc["name"] for doc in listed] == ["report.pdf"]
 
 
 def test_openai_agent_config_scoped_shadow_check(client, store_path):
@@ -897,6 +966,60 @@ def test_claude_agent_config_scoped_shadow_check(client, store_path):
     assert "report.pdf" in config["system_prompt"]
 
 
+def test_anthropic_runner_config_thinking_lifts_max_tokens(client):
+    pytest.importorskip("anthropic")
+    config = client.anthropic_runner_config(
+        model="claude-sonnet-4-5",
+        thinking={"type": "enabled", "budget_tokens": 10000})
+    assert config["max_tokens"] == 10000 + 8192
+    assert config["thinking"] == {"type": "enabled", "budget_tokens": 10000}
+    assert "thinking" not in client.anthropic_runner_config(
+        model="claude-sonnet-4-5")
+
+
+def test_bridge_invoker_reraises_auth_failures():
+    class Revoked:
+        def call_tool(self, name, arguments):
+            raise PageIndexAPIError("HTTP 401", status_code=401)
+
+    invoke = agent_tools_module._bridge_invoker(Revoked(), "get_document", {})
+    with pytest.raises(PageIndexAPIError, match="401"):
+        invoke({})
+    # Transport blips stay contained in the retryable envelope.
+
+    class Down:
+        def call_tool(self, name, arguments):
+            raise PageIndexAPIError("HTTP 503", status_code=503)
+
+    text, is_error = agent_tools_module._bridge_invoker(
+        Down(), "get_document", {})({})
+    assert is_error
+    assert json.loads(text)["errorCode"] == "INTERNAL_ERROR"
+
+
+def test_cloud_bridge_gates_the_endpoint(monkeypatch):
+    """Instructions come from the same endpoint the tools register: the
+    read-gated URL by default, the full one with include_management."""
+    created = []
+
+    class FakeBridge:
+        def __init__(self, url, headers):
+            created.append(url)
+
+        def instructions(self):
+            return "SERVED"
+
+    monkeypatch.setattr("pageindex.mcp_bridge.McpBridge", FakeBridge)
+    cloud = PageIndexCloudClient(api_key="pi-k")
+    assert cloud.agent_instructions() == "SERVED"
+    assert created == [f"{cloud.BASE_URL}/mcp?tools=read"]
+    cloud.agent_instructions(include_management=True)
+    assert created[1:] == [f"{cloud.BASE_URL}/mcp"]
+    cloud.agent_instructions()
+    cloud.agent_instructions(include_management=True)
+    assert len(created) == 2  # cached per gate
+
+
 def test_doc_scope_rejected_on_cloud_openai():
     pytest.importorskip("agents")
     cloud = PageIndexCloudClient(api_key="pi-test-key")
@@ -1139,7 +1262,9 @@ def test_cloud_agent_tools_discover_live_tool_set(cloud_with_fake_bridge):
     cloud, created = cloud_with_fake_bridge
     tools = cloud.agent_tools()
     bridge = created["bridge"]
-    assert bridge.url == "https://api.pageindex.ai/mcp"
+    # Default discovery rides the read-gated endpoint, matching the
+    # instructions fetch and the hosted/MCP registrations.
+    assert bridge.url == "https://api.pageindex.ai/mcp?tools=read"
     assert bridge.headers == {"Authorization": "Bearer pi-test-key"}
     # Default: only tools the server marks read-only; unannotated tools are
     # treated as non-read-only.
@@ -1184,13 +1309,12 @@ def test_cloud_agent_tools_null_description_survives():
     """A server may send description: null — .get(key, default) does not
     apply the default to it, and agent_tools() died with a TypeError while
     the _tool_specs path handled the same payload fine."""
-    from pageindex.agent_tools import _make_bridge_function
 
     class _Bridge:
         def call_tool(self, name, arguments):
             return json.dumps({"success": True}), False
 
-    tool = _make_bridge_function(_Bridge(), {
+    tool = _synth(_Bridge(), {
         "name": "search_documents",
         "description": None,
         "inputSchema": {"type": "object",
@@ -1227,6 +1351,35 @@ def test_cloud_agent_tools_list_failure_raises(monkeypatch):
         cloud.agent_tools()
 
 
+def test_local_description_edit_actually_removes_the_image_sentence():
+    """The local get_page_content description edits the contract text by
+    exact string replace — a contract wording change must fail here, not
+    silently ship the image-tool sentence to local models."""
+    from pageindex.agent_tools import TOOL_CONTRACT, _LOCAL_DESCRIPTIONS
+    local = _LOCAL_DESCRIPTIONS["get_page_content"]
+    assert "get_document_image" not in local
+    assert len(local) < len(TOOL_CONTRACT["get_page_content"]["description"])
+
+
+def test_list_tools_pagination_is_bounded():
+    """A server echoing its nextCursor terminates (no-progress guard);
+    a cycling one hits the page cap instead of hanging forever."""
+    from pageindex.mcp_bridge import McpBridge
+
+    bridge = McpBridge("http://x/mcp", {})
+    pages = {None: {"tools": [{"name": "a"}], "nextCursor": "c1"},
+             "c1": {"tools": [{"name": "b"}], "nextCursor": "c1"}}
+    bridge._request = lambda method, params=None: pages[
+        (params or {}).get("cursor")]
+    assert [t["name"] for t in bridge.list_tools()] == ["a", "b"]
+
+    bridge._request = lambda method, params=None: {
+        "tools": [],
+        "nextCursor": {"c1": "c2"}.get((params or {}).get("cursor"), "c1")}
+    with pytest.raises(PageIndexAPIError, match="did not terminate"):
+        bridge.list_tools()
+
+
 def test_mcp_bridge_protocol(monkeypatch):
     import requests as requests_mod
     from pageindex.mcp_bridge import McpBridge
@@ -1284,7 +1437,8 @@ def test_mcp_bridge_protocol(monkeypatch):
     # Replace the module's own `requests` binding — patching the shared
     # requests module would leak the fake process-wide.
     monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace(
-        post=fake_post, RequestException=requests_mod.RequestException))
+        Session=lambda: types.SimpleNamespace(post=fake_post),
+        RequestException=requests_mod.RequestException))
     bridge = McpBridge("https://api.pageindex.ai/mcp",
                        {"Authorization": "Bearer k"})
 
@@ -1347,7 +1501,8 @@ def test_mcp_bridge_400_is_an_error_not_session_expiry(monkeypatch):
         return _Resp(400, text="unknown tool")
 
     monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace(
-        post=fake_post, RequestException=requests_mod.RequestException))
+        Session=lambda: types.SimpleNamespace(post=fake_post),
+        RequestException=requests_mod.RequestException))
     bridge = McpBridge("https://api.pageindex.ai/mcp",
                        {"Authorization": "Bearer k"})
     with pytest.raises(PageIndexAPIError, match="HTTP 400"):
@@ -1408,7 +1563,8 @@ def test_mcp_bridge_init_notification_bars_concurrent_requests(monkeypatch):
         return resp
 
     monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace(
-        post=fake_post, RequestException=requests_mod.RequestException))
+        Session=lambda: types.SimpleNamespace(post=fake_post),
+        RequestException=requests_mod.RequestException))
     bridge = McpBridge("https://api.pageindex.ai/mcp",
                        {"Authorization": "Bearer k"})
 
@@ -1447,20 +1603,71 @@ def test_mcp_bridge_blob_blocks_become_stubs():
     bridge._request = lambda method, params: {"content": [
         {"type": "text", "text": "Page 3 of report.pdf"},
         {"type": "image", "mimeType": "image/png", "data": blob},
+        {"type": "resource",
+         "resource": {"mimeType": "image/jpeg", "blob": blob}},
     ]}
     text, is_error = bridge.call_tool("get_document_image", {})
     assert not is_error
     assert "Page 3 of report.pdf" in text
     assert "AAAA" not in text
     assert "[image/png content omitted: ~6 KB]" in text
+    assert "[image/jpeg content omitted: ~6 KB]" in text
 
 
 # ── review-round regressions ──
 
+def _synth(bridge, meta):
+    """Build a tool function the way the cloud lane does: signature
+    synthesis over a bridge invoker."""
+    from pageindex.agent_tools import _bridge_invoker, _make_tool_function
+    name = meta["name"]
+    return _make_tool_function(name, meta.get("description"),
+                               meta["inputSchema"],
+                               _bridge_invoker(bridge, name,
+                                               meta["inputSchema"]))
+
+
+def test_synth_binding_error_names_the_tool():
+    """Binding TypeErrors quote the function's __qualname__; the model used
+    to see \"_synthesized() got an unexpected keyword argument\" and had no
+    tool name to correct against."""
+    from pageindex.agent_tools import TOOL_CONTRACT
+
+    class _Bridge:
+        def call_tool(self, name, args):
+            return json.dumps(args), False
+
+    meta = {"name": "browse_documents", "description": "d",
+            "inputSchema": TOOL_CONTRACT["browse_documents"]["schema"]}
+    fn = _synth(_Bridge(), meta)
+    payload = json.loads(fn(bogus_param=1))
+    assert "browse_documents" in payload["error"]
+    assert "_synthesized" not in payload["error"]
+
+
+def test_bridge_invoker_coerces_string_booleans():
+    """Identical model output must behave the same on both dispatch paths:
+    call_tool coerced "false" but the cloud bridge forwarded it verbatim,
+    turning "don't wait" into a 3-minute wait on lenient servers."""
+    from pageindex.agent_tools import TOOL_CONTRACT, _bridge_invoker
+
+    seen = {}
+
+    class _Bridge:
+        def call_tool(self, name, args):
+            seen.update(args)
+            return "{}", False
+
+    invoke = _bridge_invoker(_Bridge(), "get_document",
+                             TOOL_CONTRACT["get_document"]["schema"])
+    invoke({"doc_name": "q.pdf", "wait_for_completion": "false"})
+    assert seen["wait_for_completion"] is False
+
+
 def test_synth_optional_no_default_param_is_nullable():
     """A non-required, no-default schema param must annotate Optional, or
     strict schemas force the model to always send a value (browse.query)."""
-    from pageindex.agent_tools import _make_bridge_function, TOOL_CONTRACT
+    from pageindex.agent_tools import TOOL_CONTRACT
     from typing import get_args
 
     class _Bridge:
@@ -1470,7 +1677,7 @@ def test_synth_optional_no_default_param_is_nullable():
     meta = {"name": "browse_documents",
             "description": "d",
             "inputSchema": TOOL_CONTRACT["browse_documents"]["schema"]}
-    fn = _make_bridge_function(_Bridge(), meta)
+    fn = _synth(_Bridge(), meta)
     assert type(None) in get_args(fn.__annotations__["query"])
 
 
@@ -1479,7 +1686,6 @@ def test_synth_array_params_keep_their_item_type():
     function_tool then emits {"type": "array", "items": {}}, which strict
     function calling rejects."""
     from typing import Optional
-    from pageindex.agent_tools import _make_bridge_function
 
     class _Bridge:
         def call_tool(self, name, args):
@@ -1498,7 +1704,7 @@ def test_synth_array_params_keep_their_item_type():
                 },
                 "required": ["doc_ids", "mixed"],
             }}
-    fn = _make_bridge_function(_Bridge(), meta)
+    fn = _synth(_Bridge(), meta)
     assert fn.__annotations__["doc_ids"] == list[str]
     assert fn.__annotations__["tags"] == Optional[list[int]]
     # A type-array in items (nullable elements) degrades to bare list —
@@ -1507,7 +1713,6 @@ def test_synth_array_params_keep_their_item_type():
 
 
 def test_synth_escape_hatches():
-    from pageindex.agent_tools import _make_bridge_function
 
     calls = []
 
@@ -1517,7 +1722,7 @@ def test_synth_escape_hatches():
             return "ok", False
 
     # Tool named "_invoke" must not recurse into itself.
-    invoke_named = _make_bridge_function(_Bridge(), {
+    invoke_named = _synth(_Bridge(), {
         "name": "_invoke", "description": "d",
         "inputSchema": {"type": "object", "properties": {"x": {"type": "string"}},
                         "required": ["x"]}})
@@ -1525,7 +1730,7 @@ def test_synth_escape_hatches():
     assert calls[-1] == ("_invoke", {"x": "v"})
 
     # Param named "dict" must not shadow the builtin.
-    dict_param = _make_bridge_function(_Bridge(), {
+    dict_param = _synth(_Bridge(), {
         "name": "t", "description": "d",
         "inputSchema": {"type": "object", "properties": {"dict": {"type": "string"}},
                         "required": ["dict"]}})
@@ -1534,7 +1739,7 @@ def test_synth_escape_hatches():
 
     # Non-identifier tool name still gets a real signature.
     import inspect
-    dashed = _make_bridge_function(_Bridge(), {
+    dashed = _synth(_Bridge(), {
         "name": "page-content.v2", "description": "d",
         "inputSchema": {"type": "object", "properties": {"a": {"type": "string"}},
                         "required": ["a"]}})
@@ -1603,7 +1808,8 @@ def test_bridge_call_tool_surfaces_iserror(monkeypatch):
             "content": [{"type": "text", "text": '{"error": "denied"}'}]}})
 
     monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace(
-        post=fake_post, RequestException=requests_mod.RequestException))
+        Session=lambda: types.SimpleNamespace(post=fake_post),
+        RequestException=requests_mod.RequestException))
     bridge = McpBridge("https://api.pageindex.ai/mcp", {})
     assert bridge.call_tool("t", {}) == ('{"error": "denied"}', True)
 
@@ -1640,7 +1846,8 @@ def test_bridge_rejects_mismatched_reply_id(monkeypatch):
                                                    "text": "old"}]}})
 
     monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace(
-        post=fake_post, RequestException=requests_mod.RequestException))
+        Session=lambda: types.SimpleNamespace(post=fake_post),
+        RequestException=requests_mod.RequestException))
     bridge = McpBridge("https://api.pageindex.ai/mcp", {})
     with pytest.raises(PageIndexAPIError, match="no reply matching"):
         bridge.call_tool("t", {})
@@ -1664,7 +1871,8 @@ def test_bridge_transport_error_is_pageindex_error(monkeypatch):
         raise requests_mod.ConnectionError("dns down")
 
     monkeypatch.setattr(mcp_bridge, "requests", types.SimpleNamespace(
-        post=dead_post, RequestException=requests_mod.RequestException))
+        Session=lambda: types.SimpleNamespace(post=dead_post),
+        RequestException=requests_mod.RequestException))
     bridge = McpBridge("https://api.pageindex.ai/mcp", {})
     with pytest.raises(PageIndexAPIError, match="Could not reach"):
         bridge.list_tools()
@@ -1737,6 +1945,17 @@ def test_all_documents_survives_short_pages_and_missing_total():
     assert _all_documents(exact) == docs
     assert type(exact).calls == 2   # ...total still saves the empty page
 
+    # stop_ids ends the walk once every wanted id has been seen — a
+    # doc-scoped chat turn must not page the whole library...
+    early = make_client(120, 100)
+    listed = _all_documents(early, stop_ids=frozenset({"pi-3"}))
+    assert type(early).calls == 1
+    assert any(doc["id"] == "pi-3" for doc in listed)
+    # ...while an id the listing lacks still costs the full sweep.
+    full = make_client(120, 100)
+    assert _all_documents(full, stop_ids=frozenset({"pi-missing"})) == docs
+    assert type(full).calls == 2
+
 
 def test_null_arguments_mean_omitted(client, store_path):
     """Adapters that forward the model's null values verbatim (the Claude
@@ -1759,17 +1978,6 @@ def test_page_spec_span_bomb_rejected(client, store_path):
     assert "Too many pages" in payload["error"]
 
 
-def test_agent_instructions_shadowed_doc_id_raises(client, store_path):
-    seed_doc(store_path, "pi-old", "report.pdf",
-             created_at="2026-08-01T10:00:00.000000")
-    seed_doc(store_path, "pi-new", "report.pdf",
-             created_at="2026-08-02T10:00:00.000000")
-    with pytest.raises(PageIndexAPIError, match="shadowed"):
-        client.agent_instructions(doc_id="pi-old")
-    text = client.agent_instructions(doc_id="pi-new")
-    assert "report.pdf" in text
-
-
 def test_wait_tolerates_transient_network_failures(fake_cloud_client, monkeypatch):
     import requests as requests_mod
     cloud = fake_cloud_client(["processing", "completed"])
@@ -1826,6 +2034,7 @@ def test_wait_tolerates_transient_poll_failures(fake_cloud_client, monkeypatch):
     assert cloud.submit_document("x.pdf", wait=True) == {"doc_id": "pi-fake"}
 
 
+import pageindex.utils  # noqa: F401 — its import loads .env
 LIVE_KEY = os.getenv("PAGEINDEX_API_KEY")
 
 
@@ -2008,6 +2217,31 @@ def test_cloud_bridge_cache_threadsafe_and_pickle_clean(monkeypatch):
     pickle.dumps(cloud)
 
 
+def test_cloud_bridge_rebuilds_on_credential_change(monkeypatch):
+    """CloudAPI re-reads client.api_key on every REST call; the MCP half
+    must not keep authenticating with a rotation-stale snapshot."""
+    import pageindex.mcp_bridge as mcp_bridge
+    from pageindex.agent_tools import _cloud_bridge
+    built = []
+
+    class _Bridge(_FakeBridge):
+        def __init__(self, url, headers):
+            super().__init__(url, headers)
+            built.append((url, dict(headers)))
+
+    monkeypatch.setattr(mcp_bridge, "McpBridge", _Bridge)
+    cloud = PageIndexCloudClient(api_key="pi-old")
+    first = _cloud_bridge(cloud)
+    assert _cloud_bridge(cloud) is first  # unchanged credentials: cached
+    cloud.api_key = "pi-new"
+    second = _cloud_bridge(cloud)
+    assert second is not first
+    assert built[-1][1]["Authorization"] == "Bearer pi-new"
+    cloud.BASE_URL = "https://alt.example"
+    assert _cloud_bridge(cloud) is not second
+    assert built[-1][0] == "https://alt.example/mcp"
+
+
 def test_cloud_agent_instructions_blank_or_nonstring_raises(monkeypatch):
     """Whitespace-only or non-string initialize.instructions must hit the
     same honest error as a missing one — never a blank system prompt."""
@@ -2141,6 +2375,28 @@ def test_config_helpers_reject_empty_doc_id_on_cloud():
         cloud.claude_agent_config(doc_id=[])
 
 
+def test_doc_targeting_keeps_transport_errors_out_of_not_found():
+    """A cloud 429/5xx during the doc_id fetch is an outage, not a missing
+    document — only a definite not-found/denied batches into the
+    "Documents not found" message; anything else propagates raw."""
+    class Stub:
+        def __init__(self, status):
+            self.status = status
+
+        def get_document(self, doc_id):
+            raise PageIndexAPIError(
+                f"Failed to get document metadata: {self.status}",
+                status_code=self.status)
+
+    for status in (429, 500):
+        with pytest.raises(PageIndexAPIError, match=f"metadata: {status}"):
+            agent_tools_module.doc_targeting_block(Stub(status), "pi-a")
+    for status in (403, 404):
+        with pytest.raises(PageIndexAPIError,
+                           match="Documents not found or access denied: pi-a"):
+            agent_tools_module.doc_targeting_block(Stub(status), "pi-a")
+
+
 def test_call_tool_coerces_string_booleans(client, store_path, monkeypatch):
     """Models routinely send booleans as JSON strings — "false" must not
     read as True (a full wait_for_completion stall)."""
@@ -2221,3 +2477,54 @@ def test_agent_instructions_carry_user_metadata(client, store_path):
              metadata={"quarter": "Q3", "year": 2025})
     text = client.agent_instructions(doc_id="pi-1")
     assert '"quarter": "Q3"' in text and '"year": 2025' in text
+
+
+# ── wait-poll resilience, instruction scoping, agent_tools doc_id ──
+
+def test_await_completion_polls_through_transient_refetch_failure(monkeypatch):
+    """A refetch that fails once must not end the wait early — the caller
+    would report that 5-second exit as the full 3-minute timeout."""
+    monkeypatch.setattr(agent_tools_module, "_TOOL_WAIT_INTERVAL", 0.0)
+    calls = {"n": 0}
+
+    class Flaky:
+        def get_document(self, doc_id):
+            calls["n"] += 1
+            if calls["n"] == 1:
+                raise PageIndexAPIError("transient listing failure")
+            return {"id": doc_id, "status": "completed"}
+
+    result = agent_tools_module._await_completion(
+        Flaky(), {"id": "pi-x", "status": "processing"}, wait=True)
+    assert result["status"] == "completed"
+    assert calls["n"] == 2
+
+
+def test_agent_instructions_doc_id_shadow_check(client, store_path):
+    seed_doc(store_path, "pi-old", "report.pdf",
+             created_at="2026-08-01T10:00:00.123000")
+    seed_doc(store_path, "pi-new", "report.pdf",
+             created_at="2026-08-05T10:00:00.123000")
+    # standalone instructions get the strict check; only the *_agent_config
+    # bundles (which build the tools too) relax it
+    with pytest.raises(PageIndexAPIError, match="shadowed"):
+        client.agent_instructions(doc_id="pi-old")
+
+
+def test_agent_tools_doc_id_scopes_the_functions(client, store_path):
+    seed_doc(store_path, "pi-a", "alpha.pdf")
+    seed_doc(store_path, "pi-b", "secret.pdf")
+    funcs = {fn.__name__: fn for fn in client.agent_tools(doc_id="pi-a")}
+    blocked = json.loads(funcs["get_page_content"](doc_name="secret.pdf",
+                                                   pages="1"))
+    assert "success" not in blocked
+    assert blocked["errorCode"] == "NOT_FOUND"
+    allowed = json.loads(funcs["get_page_content"](doc_name="alpha.pdf",
+                                                   pages="1"))
+    assert allowed["success"] is True
+
+
+def test_agent_tools_doc_id_refused_on_cloud():
+    cloud = PageIndexCloudClient(api_key="pi-test-key")
+    with pytest.raises(PageIndexAPIError, match="local tools only"):
+        cloud.agent_tools(doc_id="pi-a")

tests/test_client.py

@@ -68,14 +68,19 @@ def test_local_client_does_not_touch_disk(tmp_path):
     assert not storage.exists()
 
 
-def test_retrieve_model_carries_agents_sdk_prefix(tmp_path):
+def test_retrieve_model_stays_as_configured(tmp_path):
+    """The public attribute keeps the caller's spelling — ``litellm/`` is
+    Agents SDK routing grammar, applied at the config door
+    (openai_agent_config), never baked into ``chat_model``: handed to the
+    Anthropic SDK or a raw request, the prefixed form is a 404."""
     def resolved(retrieve_model):
         return PageIndexClient(retrieve_model=retrieve_model,
                                storage_path=str(tmp_path / "s")).retrieve_model
 
-    assert resolved("anthropic/claude-sonnet-4-6") == "litellm/anthropic/claude-sonnet-4-6"
-    for already_routable in ("gpt-4o", "openai/gpt-4o", "litellm/anthropic/claude-sonnet-4-6"):
-        assert resolved(already_routable) == already_routable
+    for as_configured in ("anthropic/claude-sonnet-4-6", "gpt-4o",
+                          "openai/gpt-4o",
+                          "litellm/anthropic/claude-sonnet-4-6"):
+        assert resolved(as_configured) == as_configured
 
 
 def test_model_resolution_covers_every_generation(tmp_path):
@@ -173,7 +178,7 @@ def test_get_page_content(local_client, indexed_doc):
 
     assert local_client.get_page_content(indexed_doc, "99") == []
 
-    with pytest.raises(ValueError):
+    with pytest.raises(PageIndexAPIError):
         local_client.get_page_content(indexed_doc, "abc")
 
 
@@ -181,7 +186,7 @@ def test_get_page_content_span_bomb_rejected(local_client, indexed_doc):
     """An absurd range must be rejected arithmetically, not expanded into
     a billion integers in the caller's process (the tool layer already
     refused; the public client method did not)."""
-    with pytest.raises(ValueError, match="spans more than 10000"):
+    with pytest.raises(PageIndexAPIError, match="spans more than 10000"):
         local_client.get_page_content(indexed_doc, "1-1000001")
     # At the bound itself the spec still parses.
     assert local_client.get_page_content(indexed_doc, "5-10004") == []
@@ -285,18 +290,36 @@ def test_page_index_flash_rejects_unknown_optimize():
 
 def test_llm_completion_missing_key_raises_immediately(monkeypatch):
     import openai
+    import litellm  # first import may load a .env; delenv after it
     monkeypatch.delenv("OPENAI_API_KEY", raising=False)
-    monkeypatch.setattr(pageindex.utils, "_openai_sync_client", None)
-    monkeypatch.setattr(pageindex.utils, "_openai_async_client", None)
-    with pytest.raises(openai.OpenAIError):
+    with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"):
         pageindex.utils.llm_completion("gpt-4o", "probe")
-    with pytest.raises(openai.OpenAIError):
+    with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"):
         asyncio.run(pageindex.utils.llm_acompletion("gpt-4o", "probe"))
+    # unknown bare names are OpenAI shorthand, so the same check applies
+    with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"):
+        pageindex.utils.llm_completion("my-finetune-v2", "probe")
+    # a blank exported key is as missing as no key (litellm's
+    # validate_environment reports it present)
+    monkeypatch.setenv("OPENAI_API_KEY", "   ")
+    with pytest.raises(openai.OpenAIError, match="OPENAI_API_KEY"):
+        pageindex.utils.llm_completion("gpt-4o", "probe")
+
+
+def test_llm_completion_refuses_unknown_provider(monkeypatch):
+    """A first segment LiteLLM does not know (a HuggingFace repo id like
+    Qwen/...) is refused with the openai/ escape before the retry loop,
+    instead of burning it on per-call 400s."""
+    import litellm
+    monkeypatch.setattr(litellm, "completion",
+                        lambda **kw: pytest.fail("reached the wire"))
+    with pytest.raises(Exception, match="not a LiteLLM provider"):
+        pageindex.utils.llm_completion("Qwen/my-model", "probe")
 
 
 def test_submit_missing_llm_key_fails_loud(local_client, sample_pdf, monkeypatch):
+    import litellm  # first import may load a .env; delenv after it
     monkeypatch.delenv("OPENAI_API_KEY", raising=False)
-    monkeypatch.setattr(pageindex.utils, "_openai_sync_client", None)
     def first_llm_call(*args, **kwargs):
         return pageindex.utils.llm_completion("gpt-4o", "probe")
     monkeypatch.setattr(page_index_module, "page_index_main", first_llm_call)
@@ -652,6 +675,70 @@ def test_generate_summaries_partial_failure_absorbed(monkeypatch):
     assert summaries == {"A": "", "B": "ok"}
 
 
+def test_generate_summaries_unrecoverable_raises(monkeypatch):
+    """A per-node 401 must abort, not store a blank node as completed."""
+    class Denied(Exception):
+        status_code = 401
+
+    async def deny_t1(model, prompt):
+        if "t1" in prompt:
+            raise Denied("key rejected")
+        return "ok"
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", deny_t1)
+    structure = [{"title": "A", "text": "t1",
+                  "nodes": [{"title": "B", "text": "t2"}]}]
+    with pytest.raises(Denied):
+        asyncio.run(pageindex.utils.generate_summaries_for_structure(structure))
+
+
+def test_summarize_tree_child_unrecoverable_raises(monkeypatch):
+    """A 401 on a leaf must abort the run, not store a blank subtree as
+    completed: the child gather's exceptions are checked, not discarded."""
+    class Denied(Exception):
+        status_code = 401
+
+    async def deny_alpha(model, prompt):
+        if "alpha" in prompt:
+            raise Denied("key rejected")
+        return '{"points": [], "summary": "ok"}'
+    monkeypatch.setattr(pageindex.utils, "llm_acompletion", deny_alpha)
+    pdf_pages = [("alpha " * 5, 5), ("beta " * 5, 5)]
+    structure = [{"title": "R", "start_index": 1, "end_index": 2,
+                  "nodes": [
+                      {"title": "A", "start_index": 1, "end_index": 1},
+                      {"title": "B", "start_index": 2, "end_index": 2}]}]
+    with pytest.raises(Denied):
+        asyncio.run(pageindex.utils.summarize_tree(
+            structure, pdf_pages, small_node_tokens=0))
+
+
+def test_llm_completion_suppresses_litellm_cache_seeding(monkeypatch):
+    """Indexing prompts are single-shot: without an explicit injection
+    point litellm 1.97 seeds its own cache marks and every call pays the
+    write premium for nothing. Backend keys still override ours."""
+    import litellm
+    captured = {}
+
+    def fake_completion(**kwargs):
+        captured.clear()
+        captured.update(kwargs)
+        message = types.SimpleNamespace(content="ok")
+        choice = types.SimpleNamespace(message=message, finish_reason="stop")
+        return types.SimpleNamespace(choices=[choice])
+    monkeypatch.setattr(litellm, "completion", fake_completion)
+    monkeypatch.setenv("OPENAI_API_KEY", "k")
+    assert pageindex.utils.llm_completion("gpt-4o", "probe") == "ok"
+    assert captured["cache_control_injection_points"] == [
+        {"location": "message", "role": "system"}]
+    token = pageindex.utils._llm_backend.set(
+        {"api_key": "x", "cache_control_injection_points": []})
+    try:
+        pageindex.utils.llm_completion("gpt-4o", "probe")
+    finally:
+        pageindex.utils._llm_backend.reset(token)
+    assert captured["cache_control_injection_points"] == []
+
+
 def test_delete_survives_marker_tamper(local_client, tmp_path):
     tampered = tmp_path / "store" / "docs" / "tampered" / "doc.json"
     tampered.mkdir(parents=True)
@@ -833,25 +920,28 @@ def test_parse_pages_overlap_counts_union():
     from pageindex.client import _parse_pages
     pages = _parse_pages("1-5000,2000-9000")
     assert len(pages) == 9000 and pages[0] == 1 and pages[-1] == 9000
-    with pytest.raises(ValueError, match="spans more than"):
+    with pytest.raises(PageIndexAPIError, match="spans more than"):
         _parse_pages("1-10001")
+    # one parser with the tool layer now: page 0 is rejected, not passed
+    # on — surfaced as the documented SDK error type
+    with pytest.raises(PageIndexAPIError, match="positive"):
+        _parse_pages("0-3")
 
 
 # ── backend: the indexing lane ──
 
 def test_backend_scopes_the_index_lane(tmp_path, monkeypatch):
-    """index_backend reaches both gateway paths — LiteLLM's call kwargs
-    and a fresh openai-SDK client — scoped to the operation, with the
-    endpoint spelling normalized for the openai SDK."""
+    """index_backend reaches the indexing lane's LiteLLM call kwargs
+    verbatim — bare and provider-prefixed models alike — scoped to the
+    operation, bypassing the env pre-check."""
     pytest.importorskip("litellm")
     import litellm
-    import openai
     from types import SimpleNamespace
     from pageindex.local_api import LocalAPI
     from pageindex.utils import _llm_backend, llm_completion
 
     api = LocalAPI(storage_path=str(tmp_path / "s"), model="m",
-                   summary_model="s", retrieve_model="r",
+                   summary_model="s",
                    index_backend={"api_key": "ik", "api_base": "http://b"})
     assert api._with_backend(_llm_backend.get) == {"api_key": "ik",
                                                    "api_base": "http://b"}
@@ -862,41 +952,64 @@ def test_backend_scopes_the_index_lane(tmp_path, monkeypatch):
     captured = {}
     monkeypatch.setattr(litellm, "completion",
                         lambda **kw: (captured.update(kw), reply)[1])
-    api._with_backend(lambda: llm_completion("anthropic/claude-x", "p"))
-    assert captured["api_key"] == "ik"
-    assert captured["api_base"] == "http://b"
-
-    seen = {}
-
-    class _FakeOpenAI:
-        def __init__(self, **kw):
-            seen.update(kw)
-            self.chat = SimpleNamespace(completions=SimpleNamespace(
-                create=lambda **_: reply))
-
-    monkeypatch.setattr(openai, "OpenAI", _FakeOpenAI)
-    api._with_backend(lambda: llm_completion("gpt-4o", "p"))
-    assert seen["api_key"] == "ik"
-    assert seen["base_url"] == "http://b"
-
-
-def test_index_backend_skips_the_env_precheck(tmp_path, monkeypatch):
-    """The missing-key pre-check reads environment variables only, so a
-    backend-supplied key must bypass it instead of being refused."""
-    pytest.importorskip("litellm")
-    import litellm
-    from pageindex.local_api import LocalAPI
-
-    api = LocalAPI(storage_path=str(tmp_path / "s"), model="m",
-                   summary_model="s", retrieve_model="r",
-                   index_backend={"api_key": "ik"})
     monkeypatch.setattr(litellm, "validate_environment",
                         lambda *a, **k: pytest.fail("env pre-check ran"))
-    monkeypatch.setattr(pageindex.flash, "page_index_flash",
-                        lambda *a, **k: (_ for _ in ()).throw(
-                            RuntimeError("reached-flash")))
-    with pytest.raises(RuntimeError, match="reached-flash"):
-        api._index_flash("f.pdf", ["text"])
+    for model, wire in (("anthropic/claude-x", "anthropic/claude-x"),
+                        ("gpt-4o", "openai/gpt-4o"),
+                        ("my-finetune-v2", "openai/my-finetune-v2")):
+        captured.clear()
+        api._with_backend(lambda: llm_completion(model, "p"))
+        assert captured["model"] == wire
+        assert captured["api_key"] == "ik"
+        assert captured["api_base"] == "http://b"
+
+
+def test_index_precheck_covers_only_openai_shaped(monkeypatch):
+    """The missing-key pre-check fires only for OpenAI-shaped names — other
+    providers resolve credentials at call time (IAM chains, ADC, Ollama's
+    localhost default), invisible to env inspection, so the lane must not
+    block them up front."""
+    pytest.importorskip("litellm")
+    import litellm
+    from types import SimpleNamespace
+    from pageindex.utils import llm_completion
+
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    with pytest.raises(litellm.AuthenticationError, match="missing API key"):
+        llm_completion("my-finetune-v2", "p")
+
+    reply = SimpleNamespace(choices=[SimpleNamespace(
+        message=SimpleNamespace(content="ok"), finish_reason="stop")])
+    monkeypatch.setattr(litellm, "completion", lambda **kw: reply)
+    monkeypatch.setattr(litellm, "validate_environment",
+                        lambda *a, **k: pytest.fail("env pre-check ran"))
+    assert llm_completion("ollama/llama3", "p") == "ok"
+    assert llm_completion("bedrock/anthropic.claude-sonnet", "p") == "ok"
+
+
+def test_litellm_routing_prefix_skips_key_precheck(monkeypatch):
+    """litellm/-prefixed names are an explicit routing choice: litellm
+    resolves credentials beyond the environment (litellm.api_key, a
+    keyless OPENAI_BASE_URL server), so the env pre-check stands aside."""
+    pytest.importorskip("litellm")
+    import litellm  # noqa: F401 — first import may load a .env; delenv after
+    from pageindex.utils import _litellm_model, _openai_missing_keys
+
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    assert _openai_missing_keys("litellm/gpt-4o") == []
+    assert _litellm_model("litellm/gpt-4o", None) == "openai/gpt-4o"
+
+
+def test_custom_provider_map_passes_provider_precheck(monkeypatch):
+    """litellm appends custom_provider_map providers to provider_list only
+    at completion time, so the pre-check must consult the map itself."""
+    pytest.importorskip("litellm")
+    import litellm
+    from pageindex.utils import _litellm_model
+
+    monkeypatch.setattr(litellm, "custom_provider_map",
+                        [{"provider": "my-llm", "custom_handler": object()}])
+    assert _litellm_model("my-llm/model-a", None) == "my-llm/model-a"
 
 
 def test_backend_args_are_local_only():
@@ -904,3 +1017,65 @@ def test_backend_args_are_local_only():
         PageIndexClient(api_key="pi-k", chat_backend={"api_key": "x"})
     with pytest.raises(PageIndexAPIError, match="index_backend"):
         PageIndexClient(api_key="pi-k", index_backend={"api_key": "x"})
+
+
+def test_chat_wraps_answerless_cloud_reply(monkeypatch):
+    """A cloud reply without choices (filtered / malformed) surfaces as
+    the SDK's error, not a bare IndexError/KeyError."""
+    client = PageIndexClient(api_key="pi-k")
+    for reply in ({"id": "x", "object": "chat.completion", "choices": []},
+                  {"id": "x"}):
+        monkeypatch.setattr(client, "chat_completions",
+                            lambda *a, _r=reply, **k: _r)
+        with pytest.raises(PageIndexAPIError, match="carries no answer"):
+            client.chat("hi")
+
+
+def test_retrieve_model_assignment_still_works(local_client):
+    """0.2.9 allowed `client.retrieve_model = ...`; the legacy property
+    keeps the write path as an alias for chat_model."""
+    local_client.retrieve_model = "gpt-x"
+    assert local_client.chat_model == "gpt-x"
+    assert local_client.retrieve_model == "gpt-x"
+
+
+def test_concurrent_same_name_submits_store_unique_names(local_client,
+                                                         sample_pdf,
+                                                         monkeypatch):
+    """Name uniquing runs under the store lock at save time, so two
+    clients indexing the same filename concurrently cannot both store it
+    — a stored duplicate would shadow the older doc_id forever."""
+    import threading
+    import time as time_mod
+
+    def slow_flash(pdf, **kwargs):
+        time_mod.sleep(0.1)  # both threads index before either saves
+        return {"structure": [{"title": "T", "start_index": 1,
+                               "end_index": 2, "summary": "s", "nodes": []}]}
+
+    monkeypatch.setattr(pageindex.flash, "page_index_flash", slow_flash)
+    monkeypatch.setattr(pageindex.utils, "llm_completion",
+                        lambda *a, **k: "d")
+    results = []
+    workers = [threading.Thread(
+        target=lambda: results.append(local_client.submit_document(sample_pdf)))
+        for _ in range(2)]
+    for worker in workers:
+        worker.start()
+    for worker in workers:
+        worker.join()
+    assert {r["name"] for r in results} == {"sample.pdf", "sample_1.pdf"}
+
+
+
+def test_format_tree_node_keeps_key_items():
+    """key_items from the merge optimization survive the get_tree formatter."""
+    from pageindex.local_api import _format_tree_node
+
+    node = {"title": "Chapter 1", "node_id": "0000", "start_index": 1,
+            "summary": "s",
+            "key_items": ["1.1 Alpha", "1.2 Beta", "1.3 Gamma"]}
+    out = _format_tree_node(node, node_summary=True)
+    assert out["key_items"] == ["1.1 Alpha", "1.2 Beta", "1.3 Gamma"]
+    assert "key_items" not in _format_tree_node(
+        {"title": "t", "node_id": "0001", "start_index": 1}, False)

tests/test_flash_extraction.py

@@ -0,0 +1,170 @@
+"""Pins the pdfium 5.x text-extraction semantics the parser is calibrated to.
+
+pdfium split FPDFFont_GetFontName into GetBaseFontName (/BaseFont, per-face)
+and GetFamilyName (old family semantics); the parser uses the per-face names,
+which changes word joining and figure-label pickup. These sentinels come from
+a 4.30-vs-5.13 corpus A/B and fail if the semantics move again.
+"""
+from importlib.metadata import version
+from pathlib import Path
+
+import pytest
+
+PDF = Path(__file__).parent.parent / "examples" / "documents" / "earthmover.pdf"
+
+
+@pytest.mark.skipif(int(version("pypdfium2").split(".")[0]) < 5,
+                    reason="extraction is pinned to pdfium 5.x font-name semantics")
+def test_page_text_pins_pdfium5_semantics():
+    from pageindex.flash.main import extract_toc
+
+    page7 = extract_toc(str(PDF))["page_texts"][6]
+    assert "p5\nEMD\n1.0" in page7          # figure axis label pdfium 4.x dropped
+    assert "break loop\n5: if lbp" in page7  # pseudocode lines no longer glued
+
+
+def test_page_mode_walk_uses_merged_surrogate_census():
+    """The page-mode unicode walk must consume the char census char_extract
+    built (astral chars merged to one entry at the high-surrogate slot).
+    Re-reading the textpage split them back into two lone-surrogate slots,
+    desynced the walk against their one-char cmap targets, and silently
+    dropped every patch on any page containing an astral char."""
+    from pageindex.flash.parser_pdfium_charlevel.unicode_apply import (
+        _apply_font_unicode)
+
+    astral = {"i": 0, "ch": "\U0001d44e", "is_gen": False}   # slots 0-1 merged
+    unmapped = {"i": 2, "ch": "\x00", "is_gen": False}       # PDFium found no unicode
+    raw_chars = [astral, unmapped]
+    show_codes = [(7, (5, 6), 100.0)]
+    map_cache = {7: (1, {5: "\U0001d44e", 6: "β"})}
+
+    # objects vs show ops count differs -> page mode.
+    _apply_font_unicode(raw_chars, [], show_codes, None, map_cache)
+
+    assert astral["ch"] == "\U0001d44e"
+    assert unmapped["ch"] == "β"
+
+
+def test_optimize_full_fails_fast_without_a_key(tmp_path, monkeypatch):
+    """optimize='full' runs LLM expand: with no key configured it must be
+    an instant, guided PageIndexAPIError — raised before any PDF work (a
+    bogus path proves the ordering) — while the LLM-free spellings and a
+    backend-carrying indexing scope stay untouched."""
+    from conftest import build_pdf
+    from pageindex import PageIndexAPIError
+    from pageindex.flash import page_index_flash
+    from pageindex.utils import _llm_backend
+    import litellm  # noqa: F401 — first import may load a .env; delenv after it
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    monkeypatch.delenv("CHATGPT_API_KEY", raising=False)
+
+    with pytest.raises(PageIndexAPIError, match="optimize='merge'"):
+        page_index_flash(str(tmp_path / "missing.pdf"), summary=False)
+
+    pdf = tmp_path / "doc.pdf"
+    pdf.write_bytes(build_pdf(["1 Introduction", "Body text"]))
+    result = page_index_flash(str(pdf), summary=False, optimize="merge")
+    assert "structure" in result
+    result = page_index_flash(str(pdf), summary=False, optimize=False)
+    assert "structure" in result
+
+    token = _llm_backend.set({"api_key": "k"})
+    try:
+        with pytest.raises(FileNotFoundError):
+            page_index_flash(str(tmp_path / "missing.pdf"), summary=False)
+    finally:
+        _llm_backend.reset(token)
+
+
+def test_bootstrap_reimport_is_not_swallowed(monkeypatch):
+    # An unguarded caller script re-imported by a spawn worker must die loudly,
+    # not fall back to a silent full sequential rerun in every worker.
+    import multiprocessing
+    import sys
+
+    from pageindex.flash import parser_pdfium_parallel as mod
+
+    class BoomExecutor:
+        def __init__(self, *a, **k):
+            pass
+
+        def map(self, *a, **k):
+            raise RuntimeError("start a new process before bootstrapping")
+
+        def shutdown(self, *a, **k):
+            pass
+
+    monkeypatch.setattr(mod, "ProcessPoolExecutor", BoomExecutor)
+    cur = multiprocessing.current_process()
+
+    monkeypatch.setattr(cur, "_inheriting", True, raising=False)
+    with pytest.raises(RuntimeError):
+        mod.parse_charlevel_meta_parallel(str(PDF), workers=2, min_pages=1)
+    assert hasattr(sys.modules["__main__"], "__file__")  # window restored on error
+
+    monkeypatch.delattr(cur, "_inheriting")
+    out, meta = mod.parse_charlevel_meta_parallel(str(PDF), workers=2, min_pages=1)
+    assert len(out) == len(meta) > 0  # normal failures still fall back sequentially
+
+
+def test_submit_document_refuses_during_bootstrap(tmp_path, monkeypatch):
+    import multiprocessing
+
+    from pageindex import PageIndexAPIError, PageIndexLocalClient
+
+    c = PageIndexLocalClient(storage_path=str(tmp_path))
+    monkeypatch.setattr(
+        multiprocessing.current_process(), "_inheriting", True, raising=False
+    )
+    with pytest.raises(PageIndexAPIError, match="__main__"):
+        c.submit_document("whatever.pdf")
+
+
+def test_unguarded_script_parses_parallel_without_reexecution(tmp_path):
+    # spawn workers must not re-run an unguarded caller script: one completion,
+    # no dead-worker noise (dying workers would trip the sequential fallback).
+    import os
+    import subprocess
+    import sys
+
+    marker = tmp_path / "runs.txt"
+    script = tmp_path / "unguarded.py"
+    script.write_text(
+        "from pageindex.flash.parser_pdfium_parallel import parse_charlevel_meta_parallel\n"
+        f"out, meta = parse_charlevel_meta_parallel({str(PDF)!r}, workers=2, min_pages=1)\n"
+        "assert len(out) == len(meta) > 0\n"
+        f"open({str(marker)!r}, 'a').write('ran\\n')\n"
+    )
+    env = {**os.environ, "PYTHONPATH": str(Path(__file__).parent.parent)}
+    res = subprocess.run(
+        [sys.executable, str(script)], capture_output=True, env=env, timeout=120
+    )
+    assert res.returncode == 0, res.stderr.decode()
+    assert marker.read_text() == "ran\n"
+    assert b"Traceback" not in res.stderr
+
+
+def test_optimize_wins_over_deprecated_optimize_expand(tmp_path, monkeypatch):
+    """Explicit optimize= beats optimize_expand; legacy True still honors it."""
+    from conftest import build_pdf
+    from pageindex.flash import page_index_flash
+    import litellm  # noqa: F401 — first import may load a .env; delenv after it
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    monkeypatch.delenv("CHATGPT_API_KEY", raising=False)
+
+    pdf = tmp_path / "doc.pdf"
+    pdf.write_bytes(build_pdf(["1 Introduction", "Body text"]))
+    # resolved to "full" before the precedence fix (keyless → fail-fast)
+    with pytest.warns(DeprecationWarning):
+        result = page_index_flash(str(pdf), summary=False,
+                                  optimize="merge", optimize_expand=True)
+    assert "structure" in result
+    with pytest.warns(DeprecationWarning):
+        result = page_index_flash(str(pdf), summary=False,
+                                  optimize=True, optimize_expand=False)
+    assert "structure" in result
+    # optimize=None means unset ("full"), not off
+    from pageindex import PageIndexAPIError
+    with pytest.raises(PageIndexAPIError, match="optimize='merge'"):
+        page_index_flash(str(tmp_path / "missing.pdf"), summary=False,
+                         optimize=None)

tests/test_local_chat.py

@@ -124,6 +124,15 @@ class FakeModel(Model):
         self._record(system_instructions, input)
         output = self.turns.pop(0)
         sequence = 0
+        if getattr(self, "emit_created", False):
+            from openai.types.responses import ResponseCreatedEvent
+            sequence += 1
+            yield ResponseCreatedEvent(
+                type="response.created", sequence_number=sequence,
+                response=Response(
+                    id="resp_backend_turn", created_at=0.0, model="fake",
+                    object="response", output=[], parallel_tool_calls=False,
+                    tool_choice="auto", tools=[], status="in_progress"))
         for item in output:
             if item.type == "message":
                 pieces = getattr(self, "pieces", ("The ", "answer"))
@@ -306,7 +315,8 @@ def test_anthropic_routed_models_mark_managed_prefix_for_cache(
     fake_model([[_msg_item("ok")]])
     from pageindex.local_chat import _openai_agent
     marked = {"cache_control_injection_points": [
-        {"location": "message", "role": "system"}]}
+        {"location": "message", "role": "system"},
+        {"location": "message", "index": -1}]}
     for name in ("anthropic/claude-x", "litellm/anthropic/claude-x",
                  "bedrock/us.anthropic.claude-sonnet-5",
                  "vertex_ai/claude-sonnet-4-5"):
@@ -525,8 +535,8 @@ def test_doc_id_conversations_get_distinct_cache_keys(client, store_path,
     keys = []
     real = local_chat._conversation_cache_key
 
-    def spy(model_name, instructions, items):
-        key = real(model_name, instructions, items)
+    def spy(model_name, instructions, doc_id, items):
+        key = real(model_name, instructions, doc_id, items)
         keys.append(key)
         return key
 
@@ -551,6 +561,11 @@ def test_doc_id_conversations_get_distinct_cache_keys(client, store_path,
     client.chat_completions("Summarize section 3.", doc_id="pi-a")
     assert keys[3] != keys[4]  # same property on the chat surface
 
+    seed_doc(store_path, "pi-b", "contract.pdf")
+    fake_model([[_msg_item("f")]])
+    client.responses("What is the CAGR?", doc_id="pi-b")
+    assert keys[5] != keys[0]  # same opener, different doc: no pooling
+
 
 @needs_agents
 def test_responses_stream_passthrough(client, store_path, fake_model):
@@ -580,6 +595,30 @@ def test_responses_stream_passthrough(client, store_path, fake_model):
             .get("type", "message") == "message")
 
 
+@needs_agents
+def test_responses_stream_opens_with_created(client, store_path, fake_model):
+    """N per-turn openings collapse to one response.created, not zero —
+    the logical stream must open with a response object carrying the same
+    id the terminal event reports, and the terminal envelope reports the
+    backend's tool-param echo, not assumed values."""
+    seed_doc(store_path, "pi-a", "report.pdf")
+    fake = fake_model([
+        [_call_item("get_document", {"doc_name": "report.pdf"})],
+        [_msg_item("The answer")],
+    ])
+    fake.emit_created = True  # two turns emit two; one must pass through
+    events = list(client.responses("q", stream=True))
+    created = [e for e in events if e["type"] == "response.created"]
+    assert len(created) == 1 and events[0] is created[0]
+    assert events[0]["sequence_number"] == 1
+    terminal = events[-1]
+    assert terminal["type"] == "response.completed"
+    assert created[0]["response"]["id"] == terminal["response"]["id"]
+    assert (created[0]["response"]["created_at"]
+            == terminal["response"]["created_at"])  # one timestamp, not two
+    assert terminal["response"]["parallel_tool_calls"] is False  # echo
+
+
 @needs_agents
 def test_responses_envelope_validates_as_official_response(client, store_path,
                                                            fake_model):
@@ -824,6 +863,10 @@ def test_max_turns_rejects_non_positive(client, store_path, fake_model):
     with pytest.raises(PageIndexAPIError, match="positive integer"):
         client.chat_completions([{"role": "user", "content": "q"}],
                                 max_turns=0)
+    # every door that takes max_turns validates it, the runner config too
+    with pytest.raises(PageIndexAPIError, match="positive integer"):
+        client.anthropic_runner_config(model="claude-sonnet-4-5",
+                                       max_turns=-1)
 
 
 def test_enable_citations_rejected_before_framework_check(client, monkeypatch):
@@ -874,6 +917,7 @@ def test_responses_envelope_fields_and_cache_group(client, store_path,
                      "get_document_structure", "get_page_content"}
     assert all(tool["type"] == "function" for tool in result["tools"])
     assert result["instructions"].startswith(CHAT_HEADER)
+    # No transport echo attached here, so these are the fallbacks.
     assert result["parallel_tool_calls"] is True
     assert result["tool_choice"] == "auto"
 
@@ -893,17 +937,22 @@ def test_sol_class_refusal_names_its_exits():
 def test_conversation_cache_key_stable_per_conversation():
     """Cache-routing key, sent as the OpenAI prompt_cache_key. A
     conversation's continuations must share one key (same model /
-    instructions / first item), and unrelated conversations must not pool
-    under it."""
+    instructions / doc targeting / first item), and unrelated
+    conversations must not pool under it."""
     turn1 = [{"role": "user", "content": "q"}]
     continuation = turn1 + [{"role": "assistant", "content": "a"},
                             {"role": "user", "content": "and?"}]
-    key = local_chat._conversation_cache_key("m", "sys", turn1)
-    assert key == local_chat._conversation_cache_key("m", "sys", continuation)
+    key = local_chat._conversation_cache_key("m", "sys", "d1", turn1)
+    assert key == local_chat._conversation_cache_key(
+        "m", "sys", "d1", continuation)
+    assert key == local_chat._conversation_cache_key(
+        "m", "sys", ["d1"], turn1)  # str and one-item list: same targeting
     assert key != local_chat._conversation_cache_key(
-        "m", "sys", [{"role": "user", "content": "other"}])
-    assert key != local_chat._conversation_cache_key("m2", "sys", turn1)
-    assert key != local_chat._conversation_cache_key("m", "sys2", turn1)
+        "m", "sys", "d1", [{"role": "user", "content": "other"}])
+    assert key != local_chat._conversation_cache_key("m2", "sys", "d1", turn1)
+    assert key != local_chat._conversation_cache_key("m", "sys2", "d1", turn1)
+    assert key != local_chat._conversation_cache_key("m", "sys", "d2", turn1)
+    assert key != local_chat._conversation_cache_key("m", "sys", None, turn1)
 
 
 @needs_agents
@@ -924,6 +973,9 @@ def test_agent_carries_prompt_cache_key_in_extra_body(monkeypatch):
     settings = agent.model_settings
     assert settings.extra_body is None
     assert "cache_control_injection_points" in settings.extra_args
+    assert settings.extra_args["cache_control_injection_points"] == [
+        {"location": "message", "role": "system"},
+        {"location": "message", "index": -1}]
     for name in ("gpt-test", "openai/gpt-test", "litellm/openai/gpt-test"):
         agent = local_chat._openai_agent(
             None, "chat", name, "sys", None, None,
@@ -1070,18 +1122,18 @@ def test_doc_id_scopes_tools_to_targeted_documents(client, store_path,
 
 
 @needs_agents
-def test_empty_doc_id_is_an_empty_allowlist(client, store_path, fake_model):
-    """doc_id=[] scopes the agent to nothing; `or None` used to wash it
-    into unscoped full-library access."""
+def test_empty_doc_id_is_refused(client, store_path):
+    """doc_id=[] fails loud on every local surface, like cloud already
+    did: washing it to None would mean "everything", and the empty
+    allowlist meant "nothing" — an agent confidently reporting the
+    documents don't exist, with no signal the scope was empty."""
     seed_doc(store_path, "pi-a", "report.pdf")
-    fake = fake_model([
-        [_call_item("browse_documents", {})],
-        [_msg_item("done")],
-    ])
-    client.chat_completions("q", doc_id=[])
-    outputs = [item["output"] for item in fake.inputs[1]
-               if item.get("type") == "function_call_output"]
-    assert json.loads(outputs[-1])["documents"] == []
+    with pytest.raises(PageIndexAPIError, match="doc_id is empty"):
+        client.chat_completions("q", doc_id=[])
+    with pytest.raises(PageIndexAPIError, match="doc_id is empty"):
+        client.as_openai_tools(doc_id=[])
+    with pytest.raises(PageIndexAPIError, match="doc_id is empty"):
+        client.agent_instructions(doc_id=[])
 
 
 @needs_agents
@@ -1108,6 +1160,11 @@ def test_openai_model_resolves_provider_prefixes():
     model = local_chat._openai_model("responses", "openai/gpt-5.2")
     assert isinstance(model, OpenAIResponsesModel)
     assert str(model.model) == "gpt-5.2"
+    # litellm/ is routing grammar, not a provider: it strips before the
+    # provider-prefix guard, so an OpenAI model stays reachable.
+    model = local_chat._openai_model("responses", "litellm/gpt-5.2")
+    assert isinstance(model, OpenAIResponsesModel)
+    assert str(model.model) == "gpt-5.2"
 
 
 @needs_agents
@@ -1141,7 +1198,7 @@ def test_envelope_model_strips_litellm_routing_prefix(store_path, fake_model):
     seed_doc(store_path, "pi-a", "report.pdf")
     client = PageIndexLocalClient(storage_path=store_path,
                                   retrieve_model="anthropic/claude-x")
-    assert client.retrieve_model == "litellm/anthropic/claude-x"
+    assert client.retrieve_model == "anthropic/claude-x"
     fake_model([[_msg_item("ok")]])
     result = client.chat_completions("q")
     assert result["model"] == "anthropic/claude-x"
@@ -1170,6 +1227,7 @@ def test_envelope_model_strips_openai_routing_prefix(store_path, fake_model):
 def test_chat_missing_openai_key_fails_loud(monkeypatch):
     """A missing backend credential surfaces as the SDK's own error type,
     like every other precondition on the chat surfaces."""
+    pytest.importorskip("litellm")  # first import may load a .env; delenv after
     monkeypatch.delenv("OPENAI_API_KEY", raising=False)
     for name in ("gpt-4o", "openai/gpt-4o"):
         with pytest.raises(PageIndexAPIError, match="OPENAI_API_KEY"):
@@ -1224,6 +1282,25 @@ def test_responses_envelope_reports_backend_truncation(client, store_path,
     assert result["error"] is None
 
 
+@needs_agents
+def test_responses_envelope_reports_backend_tool_params(client, store_path,
+                                                        fake_model):
+    """tool_choice / parallel_tool_calls come from the backend's echo —
+    the request sends neither, so the envelope must not assume values."""
+    seed_doc(store_path, "pi-a", "report.pdf")
+    fake = fake_model([[_msg_item("ok")]])
+
+    async def create(*args, **kwargs):
+        return types.SimpleNamespace(status=None, tool_choice="none",
+                                     parallel_tool_calls=False)
+
+    fake._client = types.SimpleNamespace(
+        responses=types.SimpleNamespace(create=create))
+    result = client.responses("q")
+    assert result["tool_choice"] == "none"
+    assert result["parallel_tool_calls"] is False
+
+
 @needs_agents
 def test_chat_completions_wraps_framework_errors(client, store_path,
                                                  fake_model, monkeypatch):
@@ -1410,14 +1487,24 @@ def test_chat_stream_close_at_opening_chunk_cancels_run(client, store_path,
 
 @needs_agents
 def test_stream_abandonment_cancels_pending_turn(client, store_path,
-                                                 fake_model):
+                                                 fake_model, monkeypatch):
     """Closing the iterator cancels the run even while it is awaiting the
     backend: the blocked turn is torn down (pump thread exits) instead of
-    running — and billing — to completion in the background."""
+    running — and billing — to completion in the background. The pump
+    thread is tracked directly — a process-global thread count would be
+    flaky against litellm's background threads."""
     import threading
-    import time as time_mod
     seed_doc(store_path, "pi-a", "report.pdf")
-    baseline = threading.active_count()
+    pumps = []
+    real_thread = threading.Thread
+
+    class _Tracking(real_thread):
+        def __init__(self, *args, **kwargs):
+            super().__init__(*args, **kwargs)
+            if getattr(kwargs.get("target"), "__name__", "") == "pump":
+                pumps.append(self)
+
+    monkeypatch.setattr(threading, "Thread", _Tracking)
     fake = fake_model([
         [_call_item("get_document", {"doc_name": "report.pdf"})],
         [_msg_item("The answer")],
@@ -1427,11 +1514,9 @@ def test_stream_abandonment_cancels_pending_turn(client, store_path,
                                      stream=True, stream_metadata=True)
     next(stream)  # the opening role chunk
     stream.close()
-    deadline = time_mod.monotonic() + 3.0
-    while (threading.active_count() > baseline
-           and time_mod.monotonic() < deadline):
-        time_mod.sleep(0.05)
-    assert threading.active_count() <= baseline
+    assert len(pumps) == 1
+    pumps[0].join(timeout=3.0)
+    assert not pumps[0].is_alive()
     assert fake.deltas_emitted == 0  # turn 2 never produced output
 
 
@@ -1638,6 +1723,38 @@ def test_backend_connection_reaches_each_engine(monkeypatch):
     agent = local_chat._openai_agent(None, "responses", "gpt-test", "sys",
                                      None, None, backend={"api_key": "k2"})
     assert agent.model._client.api_key == "k2"
+    # the LiteLLM endpoint spelling works on the SDK-constructed door too
+    agent = local_chat._openai_agent(None, "responses", "gpt-test", "sys",
+                                     None, None,
+                                     backend={"api_key": "k3",
+                                              "api_base": "http://rb"})
+    assert str(agent.model._client.base_url).rstrip("/") == "http://rb"
+    # both endpoint spellings on one merged dict: normalization keeps the
+    # later (per-call) key instead of the eager nested pop discarding it
+    agent = local_chat._openai_agent(
+        None, "chat", "anthropic/claude-x", "sys", None, None,
+        backend=local_chat._merged_backend(
+            types.SimpleNamespace(chat_backend={"base_url": "http://client"}),
+            {"api_base": "http://call"}))
+    assert agent.model.base_url == "http://call"
+
+
+@needs_anthropic
+def test_messages_top_level_cache_control(client, store_path, fake_anthropic):
+    """The moving breakpoint rides every request so each turn re-reads the
+    growing conversation; it stands down when the caller's own marks fill
+    the four-breakpoint budget (a fifth is a live-verified 400)."""
+    seed_doc(store_path, "pi-a", "report.pdf")
+    calls = fake_anthropic([
+        _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")])
+    client.messages("q", model="claude-test", max_tokens=50)
+    assert calls[0]["cache_control"] == {"type": "ephemeral"}
+    marked = [{"type": "text", "text": f"b{i}",
+               "cache_control": {"type": "ephemeral"}} for i in range(3)]
+    calls = fake_anthropic([
+        _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")])
+    client.messages("q", model="claude-test", max_tokens=50, system=marked)
+    assert "cache_control" not in calls[0]
 
 
 def test_merged_backend_precedence():
@@ -1655,6 +1772,9 @@ def test_messages_backend_merges_and_reaches_the_client(client, fake_anthropic,
                                          "base_url": "http://x"})
     assert real.api_key == "kk"
     assert str(real.base_url).rstrip("/") == "http://x"
+    real = local_chat._anthropic_client({"api_key": "kk",
+                                         "api_base": "http://y"})
+    assert str(real.base_url).rstrip("/") == "http://y"
 
     calls = fake_anthropic([
         _anthropic_message([{"type": "text", "text": "ok"}], "end_turn")])
@@ -1714,3 +1834,207 @@ def test_messages_extra_headers_reach_the_wire(client, monkeypatch):
     client.messages("q", model="claude-sonnet-4-5",
                     extra_headers={"anthropic-beta": "context-1m-2025"})
     assert seen["beta"] == "context-1m-2025"
+
+
+@needs_agents
+def test_chat_model_settings_request_stream_usage(monkeypatch):
+    """Without include_usage the streamed run carries no usage at all and
+    the terminal chunk reports zeros (agents forwards it as
+    stream_options only on streaming calls)."""
+    monkeypatch.setattr(local_chat, "_openai_model", lambda *a: None)
+    agent = local_chat._openai_agent(None, "chat", "gpt-test", "sys",
+                                     None, None)
+    assert agent.model_settings.include_usage is True
+
+
+@needs_anthropic
+def test_messages_default_max_tokens_clears_thinking_budget(client, fake_anthropic):
+    calls = fake_anthropic([
+        _anthropic_message([{"type": "text", "text": "a"}], "end_turn"),
+    ])
+    client.messages("q", model="claude-test",
+                    thinking={"type": "enabled", "budget_tokens": 10000})
+    assert calls[0]["max_tokens"] == 10000 + 8192
+    assert calls[0]["thinking"] == {"type": "enabled",
+                                    "budget_tokens": 10000}
+    calls = fake_anthropic([  # fresh fake: each run closes its client
+        _anthropic_message([{"type": "text", "text": "b"}], "end_turn"),
+    ])
+    client.messages("q", model="claude-test", max_tokens=11000,
+                    thinking={"type": "enabled", "budget_tokens": 10000})
+    assert calls[0]["max_tokens"] == 11000  # explicit value passes through
+
+
+def test_record_chat_finish_records_and_delegates():
+    recorded = {}
+    closed = {"n": 0}
+
+    class Stream:
+        def __init__(self):
+            self.chunks = [
+                types.SimpleNamespace(choices=[]),
+                types.SimpleNamespace(choices=[types.SimpleNamespace(
+                    finish_reason="content_filter")]),
+            ]
+
+        def __aiter__(self):
+            return self
+
+        async def __anext__(self):
+            if not self.chunks:
+                raise StopAsyncIteration
+            return self.chunks.pop(0)
+
+        async def aclose(self):
+            closed["n"] += 1
+
+    async def fetch(*args, **kwargs):
+        return "shell", Stream()
+
+    model = types.SimpleNamespace(_fetch_response=fetch)
+    local_chat._record_chat_finish(types.SimpleNamespace(model=model),
+                                   recorded)
+
+    async def drive():
+        shell, tee = await model._fetch_response()
+        assert shell == "shell"
+        async for _chunk in tee:
+            pass
+        await tee.aclose()
+
+    asyncio.run(drive())
+    assert recorded == {"finish_reason": "content_filter"}
+    assert closed["n"] == 1
+    # A model without the seam: silently a no-op.
+    local_chat._record_chat_finish(
+        types.SimpleNamespace(model=types.SimpleNamespace()), {})
+
+
+@needs_agents
+def test_chat_completions_reports_native_finish_reason(client, store_path,
+                                                       monkeypatch):
+    seed_doc(store_path, "pi-a", "report.pdf")
+
+    class TruncatingModel(FakeModel):
+        async def _fetch_response(self, *args, **kwargs):
+            return types.SimpleNamespace(choices=[
+                types.SimpleNamespace(finish_reason="length")])
+
+        async def get_response(self, *args, **kwargs):
+            await self._fetch_response()
+            return await super().get_response(*args, **kwargs)
+
+        async def stream_response(self, *args, **kwargs):
+            await self._fetch_response()
+            async for event in super().stream_response(*args, **kwargs):
+                yield event
+
+    fake = TruncatingModel([[_msg_item("cut ")], [_msg_item("cut ")]])
+    monkeypatch.setattr(local_chat, "_openai_model", lambda *a: fake)
+    result = client.chat_completions("q")
+    assert result["choices"][0]["finish_reason"] == "length"
+    chunks = list(client.chat_completions("q", stream=True,
+                                          stream_metadata=True))
+    assert chunks[-2]["choices"][0]["finish_reason"] == "length"
+
+
+@needs_agents
+def test_chat_gate_honors_litellm_routing_and_custom_providers(monkeypatch):
+    """Mirrors the indexing lane: an explicit litellm/ prefix skips the env
+    key pre-check, and custom_provider_map providers pass the allowlist;
+    a name LiteLLM cannot route is still refused up front."""
+    pytest.importorskip("litellm")
+    import litellm  # first import may load a .env; delenv after it
+    from agents.extensions.models.litellm_model import LitellmModel
+
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    model = local_chat._openai_model("chat", "litellm/gpt-4o")
+    assert isinstance(model, LitellmModel) and model.model == "openai/gpt-4o"
+
+    monkeypatch.setattr(litellm, "custom_provider_map",
+                        [{"provider": "my-llm", "custom_handler": object()}])
+    model = local_chat._openai_model("chat", "my-llm/model-a")
+    assert isinstance(model, LitellmModel) and model.model == "my-llm/model-a"
+
+    with pytest.raises(PageIndexAPIError, match="not a LiteLLM provider"):
+        local_chat._openai_model("chat", "Qwen/my-model")
+
+
+@needs_agents
+def test_litellm_model_still_has_the_fetch_response_seam():
+    # Guards the private seam _record_chat_finish rides (LitellmModel
+    # ._fetch_response): a vendor rename turns the recorder into a silent
+    # no-op and every truncated turn reports finish_reason "stop".
+    pytest.importorskip("litellm")
+    from agents.extensions.models.litellm_model import LitellmModel
+
+    assert hasattr(LitellmModel, "_fetch_response")
+
+
+def test_openai_protocol_predicate_follows_litellm_routing():
+    pytest.importorskip("litellm")
+    for name in ("gpt-5", "openai/gpt-4o", "litellm/gpt-4o",
+                 "azure/gpt-4o", "openrouter/openai/gpt-4o",
+                 "deepseek/deepseek-chat", "groq/llama-3.3-70b-versatile",
+                 "xai/grok-3"):
+        assert local_chat._openai_protocol(name), name
+    for name in ("anthropic/claude-sonnet-4-5", "gemini/gemini-2.5-pro",
+                 "bedrock/us.anthropic.claude-sonnet-5",
+                 "vertex_ai/claude-sonnet-4-5"):
+        assert not local_chat._openai_protocol(name), name
+
+
+@needs_agents
+def test_chat_backend_without_key_stands_aside_like_index_lane(monkeypatch):
+    """Any non-empty backend dict suppresses the key pre-check (utils rule)."""
+    pytest.importorskip("litellm")
+    import litellm  # noqa: F401 — first import may load a .env; delenv after it
+    from agents.extensions.models.litellm_model import LitellmModel
+
+    monkeypatch.delenv("OPENAI_API_KEY", raising=False)
+    model = local_chat._openai_model(
+        "chat", "gpt-test", {"base_url": "http://localhost:9"})
+    assert isinstance(model, LitellmModel)
+
+
+@needs_agents
+def test_responses_model_marks_caller_owned_transport(monkeypatch):
+    monkeypatch.setenv("OPENAI_API_KEY", "sk-test")
+    shared = httpx.AsyncClient()
+    caller = local_chat._openai_model("responses", "gpt-test",
+                                      {"http_client": shared})
+    assert caller._client._pageindex_caller_http is True
+    owned = local_chat._openai_model("responses", "gpt-test")
+    assert owned._client._pageindex_caller_http is False
+
+    async def run():
+        await local_chat._aclose_backend(types.SimpleNamespace(model=caller))
+        assert not shared.is_closed  # caller-owned transport survives
+        await local_chat._aclose_backend(types.SimpleNamespace(model=owned))
+        await shared.aclose()
+
+    asyncio.run(run())
+
+
+@needs_anthropic
+def test_messages_keeps_caller_owned_http_client_open(client):
+    body = _anthropic_message([{"type": "text", "text": "a"}], "end_turn")
+    shared = httpx.Client(transport=httpx.MockTransport(
+        lambda request: httpx.Response(200, json=body)))
+    out = client.messages("q", model="claude-test",
+                          backend={"api_key": "t", "http_client": shared})
+    assert out["content"][0]["text"] == "a"
+    assert not shared.is_closed
+    client.messages("q", model="claude-test",
+                    backend={"api_key": "t", "http_client": shared})
+    shared.close()
+
+
+@needs_anthropic
+def test_messages_without_credentials_raises_contract_error(client,
+                                                            monkeypatch):
+    monkeypatch.delenv("ANTHROPIC_API_KEY", raising=False)
+    monkeypatch.delenv("ANTHROPIC_AUTH_TOKEN", raising=False)
+    with pytest.raises(PageIndexAPIError,
+                       match="Anthropic backend is not configured"):
+        client.messages("q", model="claude-test")

tests/test_package_surface.py

@@ -1,4 +1,5 @@
 """What `pip install pageindex` exposes: 0.2.8 helper compat and import cost."""
+import os
 import subprocess
 import sys
 
@@ -101,3 +102,21 @@ def test_classic_compat_surface_still_reachable():
     out = subprocess.run([sys.executable, "-c", probe],
                          capture_output=True, text=True, check=True)
     assert out.stdout.strip() == "ok"
+
+
+def test_import_leaves_litellm_env_untouched(tmp_path):
+    """Importing the package must not configure litellm for the host
+    process; constructing a local client (which will use litellm) does."""
+    env = {k: v for k, v in os.environ.items()
+           if k != "LITELLM_LOCAL_MODEL_COST_MAP"}
+    probe = (
+        "import os, pageindex\n"
+        "assert 'LITELLM_LOCAL_MODEL_COST_MAP' not in os.environ, "
+        "'stamped at import'\n"
+        f"pageindex.PageIndexLocalClient(storage_path={str(tmp_path / 's')!r})\n"
+        "assert os.environ['LITELLM_LOCAL_MODEL_COST_MAP'] == 'True'\n"
+        "print('ok')\n"
+    )
+    out = subprocess.run([sys.executable, "-c", probe], env=env,
+                         capture_output=True, text=True, check=True)
+    assert out.stdout.strip() == "ok"

tests/test_page_index_md.py

@@ -19,5 +19,33 @@ class ExtractNodesFromMarkdownTest(unittest.TestCase):
         )
 
 
+class MarkdownCliTest(unittest.TestCase):
+    def test_md_cli_runs_without_llm_or_key(self):
+        """--md_path with no flags makes zero LLM calls: config.yaml's PDF
+        summary default must not leak in, so the run completes without any
+        provider key and writes the structure file."""
+        import json
+        import os
+        import subprocess
+        import sys
+        import tempfile
+        from pathlib import Path
+
+        script = Path(__file__).resolve().parent.parent / "run_pageindex.py"
+        with tempfile.TemporaryDirectory() as tmp:
+            md = Path(tmp) / "notes.md"
+            md.write_text("# Title\n\nIntro.\n\n## Section\n\nBody.\n")
+            env = {k: v for k, v in os.environ.items()
+                   if k not in ("OPENAI_API_KEY", "CHATGPT_API_KEY")}
+            env["PYTHONPATH"] = str(script.parent)
+            res = subprocess.run(
+                [sys.executable, str(script), "--md_path", str(md)],
+                capture_output=True, cwd=tmp, env=env, timeout=180)
+            self.assertEqual(res.returncode, 0, res.stderr.decode())
+            out = Path(tmp) / "results" / "notes_structure.json"
+            self.assertTrue(out.exists(), res.stdout.decode())
+            json.loads(out.read_text())
+
+
 if __name__ == "__main__":
     unittest.main()