Merge pull request #524 from ggozad/fix/capability-citations

Fix capability citation handling and add eval diagnostics

CHANGELOG.md

@@ -1,12 +1,24 @@
 # Changelog
 ## [Unreleased]
 
+### Added
+
+- `evaluations run` records `cited_chunk_ids`, `searched_uris`, `n_searches`, `n_search_calls`, `n_rejected_searches`, `n_failed_tools`, `n_executions` and `n_requests` as eval attributes alongside `cited_uris`.
+
+### Changed
+
+- A capability whose search or code-execution budget is spent says so in its instructions on every following request, naming the exhausted tools.
+
 ### Fixed
 
+- A failed `analysis_execute_code` call that iterated a file object reports the `.readlines()` workaround alongside the `TypeError`.
+- A capability that reaches its request limit keeps its cite tool for two further requests that call one of its tools, while its other tools are removed.
+- A cited chunk id that does not match a retrieved id exactly resolves to the closest one above a 0.75 similarity cutoff.
 - Dotfiles are parsed as their actual format instead of a single unstructured text block. Docling ignores the extension of a name starting with a dot, so converters strip leading dots from the name they hand it.
 
 ### Documentation
 
+- `qa.max_searches` is documented as defaulting to 5 in `docs/configuration/qa.md` and `docs/configuration/index.md`, was 3.
 - New "Vacuum Memory Requirements" subsection in `docs/configuration/storage.md` documenting the ~5x peak memory of vacuum compaction relative to the `documents` table, the mitigations, and upstream issue lancedb/lancedb#2325.
 
 ## [0.71.0] - 2026-07-29

docs/capabilities/analysis.md

@@ -4,7 +4,9 @@
 
 It is deferred by default, keeping its substantial instructions and tool schemas out of context until the model chooses to load it.
 
-The default request limit is 30 model requests per question. Override it with `create_capability(request_limit=...)`, or set `request_limit=None` to disable it. As with the RAG capability, `create_capability(vision=...)` overrides the image-attachment gate, defaulting to the configured analysis model's `vision` flag. At the limit, only analysis tools are removed and the model gets one more turn to answer from gathered evidence. Other agent and capability tools remain available, and the budget resets for every agent run.
+The default request limit is 30 model requests per question. Override it with `create_capability(request_limit=...)`, or set `request_limit=None` to disable it. As with the RAG capability, `create_capability(vision=...)` overrides the image-attachment gate, defaulting to the configured analysis model's `vision` flag. At the limit, `analysis_search` and `analysis_execute_code` are removed while `analysis_cite` remains for two further requests that call an analysis tool, so the model can register citations before answering from gathered evidence. Requests spent on other capabilities do not count against that window. Other agent and capability tools remain available, and the budget resets for every agent run.
+
+When `qa.max_searches` or `analysis.max_executions` runs out, the exhausted tool keeps failing rather than disappearing, and the instructions name it on every following request. Searching from inside `analysis_execute_code` does not count against `qa.max_searches`.
 
 ## Tools
 

docs/capabilities/rag.md

@@ -26,7 +26,7 @@ print(result.output)
 
 `create_capability` accepts `db_path`, `config`, `defer_loading`, `request_limit`, and `vision`. Set `defer_loading=False` for a dedicated RAG agent where routing is unnecessary. The default request limit is 20 model requests per question; set `request_limit=None` to disable it. `vision` controls whether picture results are attached to search returns as images and should reflect the model the hosting agent runs; it defaults to the configured QA model's `vision` flag.
 
-When the limit is reached, only the RAG capability's tools are removed. The model gets one more turn to answer from evidence already gathered, while unrelated agent and capability tools remain available. A new agent run starts a fresh limit, so multi-turn chat does not consume one shared budget.
+When the limit is reached, `rag_search` is removed while `rag_cite` remains for two further requests that call a RAG tool, so the model can register citations before answering from evidence already gathered. Requests spent on other capabilities do not count against that window. Unrelated agent and capability tools remain available. A new agent run starts a fresh limit, so multi-turn chat does not consume one shared budget.
 
 ## State
 

docs/configuration/index.md

@@ -107,7 +107,7 @@ qa:
     name: gpt-oss
     enable_thinking: true
     temperature: 0.3
-  max_searches: 3
+  max_searches: 5
 
 search:
   limit: 10                    # Default number of results to return

docs/configuration/qa.md

@@ -30,12 +30,12 @@ qa:
     enable_thinking: true
     temperature: 0.3          # Default: 0.3
     vision: false             # Set true for vision-capable models
-  max_searches: 3       # Maximum search tool calls per question
+  max_searches: 5       # Maximum search tool calls per question
 ```
 
 - **model**: LLM configuration (see [Providers](providers.md#model-settings))
 - **model.vision**: Set to `true` for vision-capable models (`qwen2.5vl`, `qwen3.6`, `gpt-4o`, `claude-sonnet`, …). The capability's `search` tool only attaches picture bytes (`BinaryContent`) to its `ToolReturn` when this is `true`, otherwise picture bytes are withheld. See [Pictures × embedder × QA model](processing.md#pictures-embedder-qa-model-how-the-pieces-compose) for the full matrix.
-- **max_searches**: Maximum number of search tool calls the RAG capability can make per question (default: 3)
+- **max_searches**: Maximum number of search tool calls a capability can make per question (default: 5). Shared by the RAG and analysis capabilities.
 
 !!! note "Thinking on vLLM"
     `enable_thinking` only applies to models with a pydantic-ai reasoning profile (o-series, gpt-5, gpt-oss). For other vLLM-served models such as Qwen3 or the Gemma family, the field is a silent no-op — set the chat template switch via [`extra_body`](providers.md#raw-provider-pass-through) instead.

evaluations/evaluations/benchmark.py

@@ -434,6 +434,14 @@ async def run_qa_benchmark(
             capability_model=resolved_capability_model,
         )
         set_eval_attribute("cited_uris", result.cited_uris)
+        set_eval_attribute("cited_chunk_ids", result.cited_chunk_ids)
+        set_eval_attribute("searched_uris", result.searched_uris)
+        set_eval_attribute("n_searches", result.n_searches)
+        set_eval_attribute("n_search_calls", result.n_search_calls)
+        set_eval_attribute("n_rejected_searches", result.n_rejected_searches)
+        set_eval_attribute("n_failed_tools", result.n_failed_tools)
+        set_eval_attribute("n_executions", result.n_executions)
+        set_eval_attribute("n_requests", result.n_requests)
         return result.answer
 
     report = await _evaluate(answer_question)

evaluations/evaluations/capability_runner.py

@@ -1,9 +1,16 @@
 from collections.abc import Callable
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Protocol, cast
+from typing import Any, NamedTuple, Protocol, cast
 
 from pydantic_ai import Agent
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelResponse,
+    RetryPromptPart,
+    ToolCallPart,
+    ToolReturnPart,
+)
 from pydantic_ai.models import Model
 
 from haiku.rag.capabilities import RAGCapabilityBase
@@ -29,6 +36,70 @@ class CapabilityRunResult:
     searched_uris: list[str] = field(default_factory=list)
     n_searches: int = 0
     n_executions: int = 0
+    n_search_calls: int = 0
+    n_rejected_searches: int = 0
+    n_failed_tools: int = 0
+    n_requests: int = 0
+
+
+class ToolTraffic(NamedTuple):
+    n_search_calls: int
+    n_rejected_searches: int
+    n_failed_tools: int
+    n_requests: int
+
+
+def _count_tool_traffic(
+    messages: list[ModelMessage], namespace: str, tool_names: frozenset[str]
+) -> ToolTraffic:
+    """Count search calls, failed calls and model requests in a run.
+
+    The history is the only source: ``state.searches`` is keyed by query so it
+    hides repeats and refusals, and ``for_run`` hands the run a ``replace()``
+    copy, leaving the outer capability's counters at zero.
+
+    Only search failures mean an exhausted budget. A failed code call may be
+    either the execution budget or any error in model-written Python, so
+    ``n_failed_tools`` covers both without claiming to tell them apart. It counts
+    ``RetryPromptPart`` too, since ``_cite`` rejects with ``ModelRetry`` and only
+    ``ToolFailed`` sets ``outcome="failed"``. Both are restricted to
+    ``tool_names``, excluding host tools and output-validation retries.
+
+    ``n_requests`` counts the run's requests, which matches the capability's own
+    budget only while it stays loaded — a deferred capability skips hooks until
+    it loads.
+    """
+    search_tool = f"{namespace}_search"
+    search_calls = 0
+    rejected_searches = 0
+    failed_tools = 0
+    requests = 0
+    for message in messages:
+        if isinstance(message, ModelResponse):
+            requests += 1
+            search_calls += sum(
+                1
+                for part in message.parts
+                if isinstance(part, ToolCallPart) and part.tool_name == search_tool
+            )
+            continue
+        for part in message.parts:
+            if not isinstance(part, RetryPromptPart | ToolReturnPart):
+                continue
+            if part.tool_name not in tool_names:
+                continue
+            if isinstance(part, RetryPromptPart):
+                failed_tools += 1
+            elif part.outcome == "failed":
+                failed_tools += 1
+                if part.tool_name == search_tool:
+                    rejected_searches += 1
+    return ToolTraffic(
+        n_search_calls=search_calls,
+        n_rejected_searches=rejected_searches,
+        n_failed_tools=failed_tools,
+        n_requests=requests,
+    )
 
 
 @dataclass
@@ -100,11 +171,25 @@ async def run_capability_question(
     executions = getattr(state, "executions", None)
     n_executions = len(executions) if executions is not None else 0
 
+    traffic = _count_tool_traffic(
+        agent_result.all_messages(),
+        capability.state_namespace,
+        capability.tool_names,
+    )
+
     return CapabilityRunResult(
         answer=agent_result.output,
         cited_uris=cited_uris,
         cited_chunk_ids=cited_chunk_ids,
         searched_uris=searched_uris,
+        # Distinct search keys, not searches. Analysis files every in-code
+        # `search()` under one "_sandbox" key, so twenty sandbox searches read
+        # as one here; `n_search_calls` is the true count of search *tool*
+        # calls, and in-code searches are not counted anywhere.
         n_searches=len(typed.searches),
         n_executions=n_executions,
+        n_search_calls=traffic.n_search_calls,
+        n_rejected_searches=traffic.n_rejected_searches,
+        n_failed_tools=traffic.n_failed_tools,
+        n_requests=traffic.n_requests,
     )

evaluations/tests/test_capability_runner.py

@@ -2,13 +2,108 @@ from types import SimpleNamespace
 from unittest.mock import AsyncMock, patch
 
 import pytest
+from pydantic_ai.messages import (
+    ModelRequest,
+    ModelResponse,
+    RetryPromptPart,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+    UserPromptPart,
+)
 from pydantic_ai.models.test import TestModel
 
-from evaluations.capability_runner import run_capability_question
+from evaluations.capability_runner import _count_tool_traffic, run_capability_question
 from haiku.rag.capabilities.analysis import create_capability as create_analysis
 from haiku.rag.capabilities.rag import create_capability as create_rag
 from haiku.rag.config.models import AppConfig
 
+ANALYSIS_TOOLS = frozenset(
+    {"analysis_search", "analysis_execute_code", "analysis_cite"}
+)
+
+
+def test_count_tool_traffic_sees_a_rejected_cite_call():
+    """`_cite` rejects with ModelRetry, which is not a failed ToolReturnPart."""
+    messages = [
+        ModelRequest(parts=[UserPromptPart(content="q")]),
+        ModelResponse(parts=[ToolCallPart("analysis_cite", {"chunk_ids": []})]),
+        ModelRequest(
+            parts=[
+                RetryPromptPart(
+                    tool_name="analysis_cite",
+                    content="No citations registered: chunk_ids was empty.",
+                    tool_call_id="1",
+                )
+            ]
+        ),
+        ModelResponse(parts=[TextPart("done")]),
+    ]
+
+    traffic = _count_tool_traffic(messages, "analysis", ANALYSIS_TOOLS)
+
+    assert traffic.n_failed_tools == 1
+    assert traffic.n_rejected_searches == 0
+
+
+def test_count_tool_traffic_separates_search_rejections_from_code_errors():
+    """A crash in model-written Python must not read as budget exhaustion."""
+    messages = [
+        ModelRequest(parts=[UserPromptPart(content="q")]),
+        ModelResponse(parts=[ToolCallPart("analysis_execute_code", {"code": "1/0"})]),
+        ModelRequest(
+            parts=[
+                ToolReturnPart(
+                    tool_name="analysis_execute_code",
+                    content="ZeroDivisionError",
+                    tool_call_id="1",
+                    outcome="failed",
+                )
+            ]
+        ),
+        ModelResponse(parts=[TextPart("done")]),
+    ]
+
+    traffic = _count_tool_traffic(messages, "analysis", ANALYSIS_TOOLS)
+
+    assert traffic.n_search_calls == 0
+    assert traffic.n_rejected_searches == 0
+    assert traffic.n_failed_tools == 1
+    assert traffic.n_requests == 2
+
+
+def test_count_tool_traffic_counts_attempts_not_distinct_queries():
+    """Rejected and repeated calls both count; `state.searches` hides them."""
+    messages = [
+        ModelRequest(parts=[UserPromptPart(content="q")]),
+        ModelResponse(
+            parts=[
+                ToolCallPart("analysis_search", {"query": "same"}),
+                ToolCallPart("analysis_search", {"query": "same"}),
+            ]
+        ),
+        ModelRequest(
+            parts=[
+                ToolReturnPart(
+                    tool_name="analysis_search", content="results", tool_call_id="1"
+                ),
+                ToolReturnPart(
+                    tool_name="analysis_search",
+                    content="Search limit reached.",
+                    tool_call_id="2",
+                    outcome="failed",
+                ),
+            ]
+        ),
+        ModelResponse(parts=[TextPart("done")]),
+    ]
+
+    traffic = _count_tool_traffic(messages, "analysis", ANALYSIS_TOOLS)
+
+    assert traffic.n_search_calls == 2
+    assert traffic.n_rejected_searches == 1
+    assert traffic.n_requests == 2
+
 
 async def test_runs_rag_capability_without_legacy_capability_layer(tmp_path):
     result = await run_capability_question(
@@ -49,7 +144,7 @@ async def test_analysis_capability_applies_request_limit(tmp_path, override, exp
     with patch(
         "evaluations.capability_runner.Agent.run", new_callable=AsyncMock
     ) as run:
-        run.return_value = SimpleNamespace(output="done")
+        run.return_value = SimpleNamespace(output="done", all_messages=lambda: [])
 
         await run_capability_question(
             lambda **_kwargs: capability,

haiku_rag_slim/haiku/rag/capabilities/_base.py

@@ -1,6 +1,7 @@
 import asyncio
 import os
 from dataclasses import dataclass, field, replace
+from difflib import get_close_matches
 from pathlib import Path
 from typing import Any, cast
 
@@ -11,6 +12,8 @@ from pydantic_ai.messages import (
     InstructionPart,
     ModelMessage,
     ModelRequest,
+    ModelResponse,
+    ToolCallPart,
     ToolReturn,
     ToolReturnPart,
     UserPromptPart,
@@ -27,6 +30,36 @@ from haiku.rag.store.models.chunk import SearchResult
 from haiku.rag.store.models.citation import Citation, resolve_citations
 from haiku.rag.tools.search import build_binary_parts_from_results
 
+CITATION_GRACE_REQUESTS = 2
+"""Requests calling this capability's tools that its cite tool outlives the rest by.
+
+A loop guard, not a budget: cite consumes no retry budget and raises nothing, so
+left available forever a stuck model calls it until the agent's own request limit
+raises ``UsageLimitExceeded`` and the question returns no answer at all. Only
+engagement can loop, which is why other capabilities' turns do not spend it.
+"""
+
+CHUNK_ID_MATCH_CUTOFF = 0.75
+"""Similarity a cited chunk id needs to be treated as a corrupted known id.
+
+Calibration knob. Two unrelated UUID4s reach about 0.5, while dropping or
+duplicating a character or a whole group stays above 0.75, so the gap is wide.
+"""
+
+
+def _nearest_known_id(chunk_id: str, known_ids: list[str]) -> str:
+    """Recover a chunk id the model damaged while transcribing it.
+
+    Models copying opaque UUIDs drop and duplicate characters and whole
+    hyphen-separated groups. Candidates are limited to ids the run actually
+    retrieved, so a wrong match needs both a near miss and a same-run neighbour.
+    Ids that match nothing are returned unchanged for the caller to report.
+    """
+    if not known_ids or chunk_id in known_ids:
+        return chunk_id
+    match = get_close_matches(chunk_id, known_ids, n=1, cutoff=CHUNK_ID_MATCH_CUTOFF)
+    return match[0] if match else chunk_id
+
 
 def resolve_db_path(db_path: Path | None, config: AppConfig) -> Path:
     if db_path is not None:
@@ -80,6 +113,17 @@ def _compact_old_tool_returns(
     return compacted
 
 
+def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) -> bool:
+    """Whether the model's most recent response called one of these tools."""
+    for message in reversed(messages):
+        if isinstance(message, ModelResponse):
+            return any(
+                isinstance(part, ToolCallPart) and part.tool_name in tool_names
+                for part in message.parts
+            )
+    return False
+
+
 @dataclass
 class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
     db_path: Path
@@ -97,6 +141,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
     resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
     search_count: int = field(default=0, repr=False)
     request_count: int = field(default=0, repr=False)
+    grace_requests_used: int = field(default=0, repr=False)
 
     async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]":
         outer = getattr(ctx.deps, "state", None)
@@ -113,6 +158,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
             resource_lock=asyncio.Lock(),
             search_count=0,
             request_count=0,
+            grace_requests_used=0,
         )
         run_capability._sync_state()
         return run_capability
@@ -128,14 +174,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
         request_context.messages = _compact_old_tool_returns(
             request_context.messages, self.tool_names
         )
-        if self._request_limit_reached:
+        if instruction := self._budget_notice():
             current_request = request_context.messages[-1]
             if isinstance(current_request, ModelRequest):
-                instruction = (
-                    f"The {self.state_namespace} capability has reached its request "
-                    "limit. Its tools are no longer available. Give the best answer "
-                    "possible using the evidence already gathered."
-                )
                 current_request.instructions = "\n\n".join(
                     part for part in (current_request.instructions, instruction) if part
                 )
@@ -147,19 +188,90 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
                         InstructionPart(content=instruction, dynamic=True),
                     ],
                 )
-        else:
-            self.request_count += 1
+        if self._request_limit_reached and _called_own_tool(
+            request_context.messages, self.tool_names
+        ):
+            self.grace_requests_used += 1
+        self.request_count += 1
         return request_context
 
+    def _budget_notice(self) -> str | None:
+        """Tell the model which of this capability's budgets just ran out.
+
+        Never names a tool ``prepare_tools`` has already withdrawn: pointing the
+        model at a tool that is gone costs it the agent's unknown-tool retry
+        budget and can abort the run.
+        """
+        if self._citation_grace_expired:
+            return (
+                f"The {self.state_namespace} capability's tools are no longer "
+                "available. Give the best answer possible using the evidence "
+                "already gathered."
+            )
+        if self._request_limit_reached:
+            return (
+                f"The {self.state_namespace} capability has reached its request "
+                f"limit. Only {self._cite_tool_name} remains among its tools: "
+                "register the chunk_ids supporting your answer, then answer from "
+                "the evidence already gathered."
+            )
+        if spent := self._spent_tool_names():
+            names = ", ".join(sorted(spent))
+            if remaining := sorted(self._evidence_tool_names() - spent):
+                return (
+                    f"The {self.state_namespace} capability has spent its budget "
+                    f"for {names}; further calls to them fail. Gather any further "
+                    f"evidence with {', '.join(remaining)}, or call "
+                    f"{self._cite_tool_name} with the chunk_ids you have and "
+                    "answer."
+                )
+            return (
+                f"The {self.state_namespace} capability has spent its budget for "
+                f"{names}; further calls to them fail. Answer from the evidence "
+                f"already gathered and call {self._cite_tool_name} with the "
+                "chunk_ids supporting it."
+            )
+        return None
+
     async def prepare_tools(
         self,
         ctx: RunContext[Any],
         tool_defs: list[ToolDefinition],
     ) -> list[ToolDefinition]:
-        """Remove only this capability's tools after its per-question limit."""
+        """Remove this capability's tools past its limit, cite tool last.
+
+        Tools whose own budget is spent stay declared on purpose. Removing one
+        makes a model that calls it anyway hit ``Unknown tool name``, charged
+        against the agent's unknown-tool retry budget, which kills the run after
+        two attempts. A spent tool that keeps failing only wastes requests.
+        """
+        if self._citation_grace_expired:
+            return [tool for tool in tool_defs if tool.capability_id != self.id]
         if not self._request_limit_reached:
             return tool_defs
-        return [tool for tool in tool_defs if tool.capability_id != self.id]
+        return [
+            tool
+            for tool in tool_defs
+            if tool.capability_id != self.id or tool.name == self._cite_tool_name
+        ]
+
+    def _evidence_tool_names(self) -> set[str]:
+        """Tools that can bring new evidence into the run."""
+        return {f"{self.state_namespace}_search"}
+
+    def _spent_tool_names(self) -> set[str]:
+        """This capability's tools whose own budget is exhausted."""
+        if self.search_count >= self._max_searches:
+            return {f"{self.state_namespace}_search"}
+        return set()
+
+    @property
+    def _cite_tool_name(self) -> str:
+        return f"{self.state_namespace}_cite"
+
+    @property
+    def _max_searches(self) -> int:
+        return self.config.qa.max_searches
 
     @property
     def _request_limit_reached(self) -> bool:
@@ -167,6 +279,12 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
             self.request_limit is not None and self.request_count >= self.request_limit
         )
 
+    @property
+    def _citation_grace_expired(self) -> bool:
+        # No `request_limit is None` guard: the counter only advances under
+        # `_request_limit_reached`, which already requires a limit.
+        return self.grace_requests_used >= CITATION_GRACE_REQUESTS
+
     async def after_run(
         self, ctx: RunContext[Any], *, result: AgentRunResult[Any]
     ) -> AgentRunResult[Any]:
@@ -211,7 +329,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
     async def _search(self, query: str, limit: int | None) -> str | ToolReturn:
         assert self.state is not None
         self.search_count += 1
-        if self.search_count > self.config.qa.max_searches:
+        if self.search_count > self._max_searches:
             raise ToolFailed(
                 "Search limit reached. Answer the question using "
                 "the results you already have."
@@ -241,11 +359,11 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
         state = cast(Any, self.state)
         for results in state.searches.values():
             all_results.extend(results)
-        citations = resolve_citations(chunk_ids, all_results)
+        known_ids = [result.chunk_id for result in all_results if result.chunk_id]
+        requested = [_nearest_known_id(cid.strip("[]"), known_ids) for cid in chunk_ids]
+        citations = resolve_citations(requested, all_results)
         resolved = {citation.chunk_id for citation in citations}
-        missing = [
-            cid.strip("[]") for cid in chunk_ids if cid.strip("[]") not in resolved
-        ]
+        missing = [cid for cid in requested if cid not in resolved]
 
         if missing:
             async with self.rag_lock:

haiku_rag_slim/haiku/rag/capabilities/analysis.py

@@ -37,6 +37,21 @@ def instructions() -> str:
     return _instructions_path.read_text().strip()
 
 
+def _recovery_hint(stderr: str) -> str:
+    """Name the workaround for sandbox limits models trip over repeatedly.
+
+    The instructions already say file objects are not iterable, and models write
+    ``for line in open(...)`` regardless. Carrying the fix in the error gives
+    them something to act on for the retry.
+    """
+    if "TextIOWrapper" in stderr and "not iterable" in stderr:
+        return (
+            "\n\nHint: file objects cannot be iterated here. Read lines with "
+            '.readlines() or .read().split("\\n").'
+        )
+    return ""
+
+
 @dataclass
 class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
     """Deferred capability for sandboxed computation over a RAG corpus."""
@@ -70,6 +85,18 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
             self.sandbox = None
         await super()._close()
 
+    def _evidence_tool_names(self) -> set[str]:
+        # Searching from inside the sandbox does not count against
+        # `qa.max_searches`, so code execution outlives a spent search budget
+        # as a way to reach new evidence.
+        return super()._evidence_tool_names() | {"analysis_execute_code"}
+
+    def _spent_tool_names(self) -> set[str]:
+        spent = super()._spent_tool_names()
+        if self.execute_count >= self.config.analysis.max_executions:
+            spent.add("analysis_execute_code")
+        return spent
+
     async def _execute_code(self, code: str) -> str:
         assert self.state is not None
         self.execute_count += 1
@@ -97,7 +124,10 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
             )
         )
         if not result.success:
-            raise ToolFailed(f"{result.stderr}\n\nOutput: {result.stdout}")
+            raise ToolFailed(
+                f"{result.stderr}{_recovery_hint(result.stderr)}"
+                f"\n\nOutput: {result.stdout}"
+            )
         return result.stdout or "No output."
 
     def get_toolset(self) -> FunctionToolset[Any]:

tests/capabilities/test_capabilities.py

@@ -17,7 +17,11 @@ from pydantic_ai.models.function import FunctionModel
 from pydantic_ai.models.test import TestModel
 from pydantic_ai.usage import RunUsage
 
-from haiku.rag.capabilities._base import _compact_old_tool_returns
+from haiku.rag.capabilities._base import (
+    CITATION_GRACE_REQUESTS,
+    _called_own_tool,
+    _compact_old_tool_returns,
+)
 from haiku.rag.capabilities.analysis import AnalysisCapability, AnalysisState
 from haiku.rag.capabilities.analysis import create_capability as create_analysis
 from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGCapability, RAGState
@@ -191,7 +195,10 @@ async def test_request_limit_removes_only_exhausted_capability_tools_per_run(
     }
     for initial, exhausted in ((0, 1), (2, 3)):
         assert analysis_tools <= seen_tools[initial]
-        assert analysis_tools.isdisjoint(seen_tools[exhausted])
+        assert {"analysis_search", "analysis_execute_code"}.isdisjoint(
+            seen_tools[exhausted]
+        )
+        assert "analysis_cite" in seen_tools[exhausted]
         assert {"host_tool", "rag_search", "rag_cite"} <= seen_tools[exhausted]
         assert (
             "analysis capability has reached its request limit"
@@ -350,6 +357,44 @@ async def test_cite_reports_unresolved_ids_on_partial_success(temp_db_path):
     assert capability.state.citations == ["chunk-1"]
 
 
+@pytest.mark.asyncio
+async def test_cite_repairs_chunk_ids_damaged_in_transcription(temp_db_path):
+    """Models mistype opaque UUIDs; near misses resolve to the retrieved id."""
+    true_id = "b8e25ea1-0bb3-48b1-8fea-2ac1f148bf7c"
+    unrelated = "9c2cd07e-5a3f-45a6-968d-cbd6f06ab57b"
+    capability = create_rag(db_path=temp_db_path, config=AppConfig())
+    capability.state = RAGState(
+        searches={
+            "q": [
+                SearchResult(
+                    content="evidence",
+                    score=1.0,
+                    chunk_id=true_id,
+                    document_id="doc-1",
+                    document_uri="test://document",
+                )
+            ]
+        }
+    )
+    client = AsyncMock()
+    client.get_chunk_by_id.return_value = None
+    capability.rag = client
+
+    dropped_char = "b8e25ea1-0bb3-48b1-8fea-2ac1f148bf7"
+    dropped_group = "0bb3-48b1-8fea-2ac1f148bf7c"
+
+    assert await capability._cite([dropped_char]) == "Registered 1 citation(s)."
+    assert await capability._cite([dropped_group]) == "Registered 1 citation(s)."
+    assert capability.state.citations == [true_id]
+
+    # An unrelated UUID is never attributed to a retrieved neighbour.
+    with pytest.raises(ModelRetry, match=unrelated):
+        await capability._cite([unrelated])
+
+    assert capability.state.citations == [true_id]
+    client.get_chunk_by_id.assert_awaited_once_with(unrelated)
+
+
 @pytest.mark.asyncio
 async def test_analysis_records_new_sandbox_search_results(temp_db_path):
     capability = create_analysis(db_path=temp_db_path, config=AppConfig())
@@ -416,6 +461,251 @@ async def test_analysis_execution_limit_fails_the_tool(temp_db_path):
         await capability._execute_code("print('done')")
 
 
+@pytest.mark.asyncio
+async def test_spent_search_budget_is_announced_but_keeps_the_tool(rag_db):
+    """A spent budget is announced; the tool stays declared to avoid a dead run.
+
+    Withdrawing it would make a model that calls it anyway hit `Unknown tool
+    name`, which exhausts the agent's unknown-tool retries and aborts the run.
+    """
+    config = AppConfig()
+    config.qa.max_searches = 1
+    seen_tools = []
+    seen_instructions = []
+    calls = 0
+
+    def model_function(_messages, info):
+        nonlocal calls
+        calls += 1
+        seen_tools.append({tool.name for tool in info.function_tools})
+        seen_instructions.append(info.instructions or "")
+        if calls == 1:
+            return ModelResponse(
+                parts=[ToolCallPart("rag_search", {"query": "machine learning"})]
+            )
+        return ModelResponse(parts=[TextPart("answered")])
+
+    agent = Agent(
+        FunctionModel(model_function),
+        deps_type=Deps,
+        capabilities=[create_rag(db_path=rag_db, config=config, defer_loading=False)],
+    )
+
+    result = await agent.run("question", deps=Deps())
+
+    assert result.output == "answered"
+    assert {"rag_search", "rag_cite"} <= seen_tools[1]
+    assert "spent its budget for rag_search" in seen_instructions[1]
+
+
+def test_grace_window_ignores_other_capabilities_turns():
+    """Only this capability's own tool calls may spend its cite window.
+
+    A multi-capability agent spends turns elsewhere; those must not expire the
+    window that exists to give this capability a chance to cite.
+    """
+    rag_tools = frozenset({"rag_search", "rag_cite"})
+
+    # Nothing to attribute before the model has responded at all.
+    assert not _called_own_tool(
+        [ModelRequest(parts=[UserPromptPart(content="q")])], rag_tools
+    )
+    assert not _called_own_tool(
+        [ModelResponse(parts=[ToolCallPart("analysis_search", {"query": "x"})])],
+        rag_tools,
+    )
+    assert not _called_own_tool(
+        [ModelResponse(parts=[TextPart("just talking")])], rag_tools
+    )
+    assert _called_own_tool(
+        [ModelResponse(parts=[ToolCallPart("rag_cite", {"chunk_ids": ["a"]})])],
+        rag_tools,
+    )
+    # Only the most recent response counts, not any earlier one.
+    assert not _called_own_tool(
+        [
+            ModelResponse(parts=[ToolCallPart("rag_cite", {"chunk_ids": ["a"]})]),
+            ModelRequest(parts=[UserPromptPart(content="next")]),
+            ModelResponse(parts=[ToolCallPart("analysis_search", {"query": "x"})]),
+        ],
+        rag_tools,
+    )
+
+
+@pytest.mark.asyncio
+async def test_spent_search_notice_points_at_code_while_it_has_budget(temp_db_path):
+    """Analysis must be sent to the sandbox, not told to answer, while it can.
+
+    In-code `search()` bypasses `qa.max_searches`, and the instructions tell the
+    model to escalate to code when search results are insufficient.
+    """
+    config = AppConfig()
+    config.qa.max_searches = 2
+    capability = create_analysis(db_path=temp_db_path, config=config)
+    capability.search_count = 2
+
+    notice = capability._budget_notice()
+
+    assert notice is not None
+    assert "analysis_search" in notice
+    assert "analysis_execute_code" in notice
+
+    # Once the code budget is gone too there is nowhere left to send it.
+    capability.execute_count = config.analysis.max_executions
+    notice = capability._budget_notice()
+    assert notice is not None
+    assert "analysis_execute_code" in notice
+    assert capability._evidence_tool_names() <= capability._spent_tool_names()
+
+
+@pytest.mark.asyncio
+async def test_spent_search_notice_tells_rag_to_answer(temp_db_path):
+    """Search is the RAG capability's only evidence tool, so stopping is right."""
+    config = AppConfig()
+    config.qa.max_searches = 2
+    capability = create_rag(db_path=temp_db_path, config=config)
+    capability.search_count = 2
+
+    notice = capability._budget_notice()
+
+    assert notice is not None
+    assert "rag_search" in notice
+    assert capability._evidence_tool_names() == {"rag_search"}
+
+
+@pytest.mark.asyncio
+async def test_spent_execution_budget_joins_the_notice(temp_db_path):
+    config = AppConfig()
+    config.analysis.max_executions = 3
+    capability = create_analysis(db_path=temp_db_path, config=config)
+
+    assert capability._spent_tool_names() == set()
+
+    capability.execute_count = 3
+
+    assert capability._spent_tool_names() == {"analysis_execute_code"}
+    notice = capability._budget_notice()
+    assert notice is not None
+    assert "analysis_execute_code" in notice
+
+
+@pytest.mark.asyncio
+async def test_exhausted_run_can_still_register_citations(rag_db):
+    """The cite tool outlives the request limit so evidence is not lost.
+
+    Reproduces the measured pathology: the model burns its request budget and
+    reaches the limit, at which point it must still be able to cite what it
+    already found.
+    """
+    config = AppConfig()
+    seen_tools = []
+    calls = 0
+    chunk_id: str | None = None
+
+    def model_function(_messages, info):
+        nonlocal calls
+        calls += 1
+        seen_tools.append({tool.name for tool in info.function_tools})
+        if calls == 1:
+            return ModelResponse(
+                parts=[ToolCallPart("rag_search", {"query": "machine learning"})]
+            )
+        if calls == 2:
+            return ModelResponse(
+                parts=[ToolCallPart("rag_cite", {"chunk_ids": [chunk_id]})]
+            )
+        return ModelResponse(parts=[TextPart("answered from gathered evidence")])
+
+    capability = create_rag(
+        db_path=rag_db,
+        config=config,
+        defer_loading=False,
+        request_limit=1,
+    )
+    agent = Agent(
+        FunctionModel(model_function),
+        deps_type=Deps,
+        capabilities=[capability],
+    )
+    deps = Deps()
+
+    async with agent.iter("question", deps=deps) as run:
+        async for _node in run:
+            if chunk_id is None:
+                searches = deps.state.get("rag", {}).get("searches") or {}
+                for results in searches.values():
+                    if results:
+                        chunk_id = results[0]["chunk_id"]
+                        break
+
+    assert chunk_id is not None
+    # The limit lands on request 2, where cite must still be offered.
+    assert "rag_search" not in seen_tools[1]
+    assert "rag_cite" in seen_tools[1]
+    assert deps.state["rag"]["citations"] == [chunk_id]
+
+
+@pytest.mark.asyncio
+async def test_cite_tool_is_withdrawn_after_the_grace_window(temp_db_path):
+    capability = create_rag(
+        db_path=temp_db_path,
+        config=AppConfig(),
+        defer_loading=False,
+        request_limit=2,
+    )
+    tool_defs = [
+        SimpleNamespace(name=name, capability_id=capability.id)
+        for name in ("rag_search", "rag_cite")
+    ]
+    ctx = make_context(Deps())
+
+    capability.request_count = 2
+    kept = await capability.prepare_tools(ctx, cast(Any, tool_defs))
+    assert {tool.name for tool in kept} == {"rag_cite"}
+    notice = capability._budget_notice()
+    assert notice is not None and "rag_cite" in notice
+
+    capability.grace_requests_used = CITATION_GRACE_REQUESTS
+    kept = await capability.prepare_tools(ctx, cast(Any, tool_defs))
+    assert kept == []
+    # The notice must never point at a tool prepare_tools has withdrawn:
+    # calling a missing tool burns the agent's unknown-tool retries and can
+    # abort the run.
+    notice = capability._budget_notice()
+    assert notice is not None
+    assert "rag_cite" not in notice
+    assert "no longer available" in notice
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    ("stderr", "expect_hint"),
+    [
+        ("TypeError: '_io.TextIOWrapper' object is not iterable", True),
+        ("TypeError: 'list' object is not an iterator", False),
+    ],
+)
+async def test_sandbox_iteration_failure_carries_the_workaround(
+    temp_db_path, stderr, expect_hint
+):
+    """A model that iterates a file object gets told what to do instead."""
+    capability = create_analysis(db_path=temp_db_path, config=AppConfig())
+    capability.state = AnalysisState()
+    sandbox = AsyncMock()
+    sandbox.execute.return_value = SandboxResult(
+        stdout="", stderr=stderr, success=False
+    )
+    sandbox._search_results = []
+    capability.sandbox = cast(Sandbox, sandbox)
+
+    with pytest.raises(ToolFailed) as failure:
+        await capability._execute_code(
+            "for line in open('/documents/x/items.jsonl'): pass"
+        )
+
+    assert (".readlines()" in str(failure.value)) is expect_hint
+
+
 @pytest.mark.asyncio
 async def test_analysis_sandbox_failure_records_execution_and_fails_the_tool(
     temp_db_path,