From 685a7c393dbfe42dfcd7bc098d429a37e839cb88 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Wed, 29 Jul 2026 11:29:31 +0300
Subject: [PATCH 1/9] Keep the cite tool past a capability's request limit

---
 CHANGELOG.md                                  |   5 +
 docs/capabilities/analysis.md                 |   4 +-
 .../haiku/rag/capabilities/_base.py           |  74 ++++++++--
 .../haiku/rag/capabilities/analysis.py        |   6 +
 tests/capabilities/test_capabilities.py       | 137 +++++++++++++++++-
 5 files changed, 213 insertions(+), 13 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index b0e087ed..164d728d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,8 +1,13 @@
 # Changelog
 ## [Unreleased]
 
+### 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 capability that reaches its request limit keeps its cite tool for two further model requests while its other tools are removed, so an exhausted run can still register citations.
 - 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
diff --git a/docs/capabilities/analysis.md b/docs/capabilities/analysis.md
index e28eb4f1..f382868a 100644
--- a/docs/capabilities/analysis.md
+++ b/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, so the model can register citations before answering from gathered evidence. 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
 
diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py
index 96b32c73..4471d162 100644
--- a/haiku_rag_slim/haiku/rag/capabilities/_base.py
+++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py
@@ -27,6 +27,9 @@ 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 the cite tool outlives this capability's other tools by."""
+
 
 def resolve_db_path(db_path: Path | None, config: AppConfig) -> Path:
     if db_path is not None:
@@ -128,14 +131,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 +145,66 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
                         InstructionPart(content=instruction, dynamic=True),
                     ],
                 )
-        else:
-            self.request_count += 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."""
+        if self._request_limit_reached:
+            return (
+                f"The {self.state_namespace} capability has reached its request "
+                f"limit. Only {self._cite_tool_name} remains available: register "
+                "the chunk_ids supporting your answer, then answer from the "
+                "evidence already gathered."
+            )
+        if spent := self._spent_tool_names():
+            return (
+                f"The {self.state_namespace} capability has spent its budget for "
+                f"{', '.join(sorted(spent))}; further calls to them fail. Answer "
+                f"from the evidence already gathered and call "
+                f"{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.
+
+        The cite tool outlives the others by ``CITATION_GRACE_REQUESTS`` so a run
+        that exhausts its budget can still record the evidence it gathered.
+
+        Tools whose own budget is spent stay declared on purpose. Removing one
+        makes a model that calls it anyway hit ``Unknown tool name``, which is
+        charged against the agent's unknown-tool retry budget and kills the run
+        after two attempts. A spent tool that keeps failing wastes requests; a
+        withdrawn one loses the whole answer.
+        """
+        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 _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 +212,13 @@ 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:
+        return (
+            self.request_limit is not None
+            and self.request_count >= self.request_limit + CITATION_GRACE_REQUESTS
+        )
+
     async def after_run(
         self, ctx: RunContext[Any], *, result: AgentRunResult[Any]
     ) -> AgentRunResult[Any]:
@@ -211,7 +263,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."
diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py
index 8c8e133c..2d754840 100644
--- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py
+++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py
@@ -70,6 +70,12 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
             self.sandbox = None
         await super()._close()
 
+    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
diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py
index 601082d1..5f8146bb 100644
--- a/tests/capabilities/test_capabilities.py
+++ b/tests/capabilities/test_capabilities.py
@@ -191,7 +191,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"
@@ -416,6 +419,138 @@ 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]
+
+
+@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"}
+
+    capability.request_count = 4
+    kept = await capability.prepare_tools(ctx, cast(Any, tool_defs))
+    assert kept == []
+
+
 @pytest.mark.asyncio
 async def test_analysis_sandbox_failure_records_execution_and_fails_the_tool(
     temp_db_path,

From 5752f61f2cba7d7469d4f95589d89b5a7172d339 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Wed, 29 Jul 2026 15:30:20 +0300
Subject: [PATCH 2/9] Recover chunk ids mistyped from search results

Resolve a cited id that misses exactly to the nearest id the run
retrieved, above a 0.75 similarity cutoff.
---
 CHANGELOG.md                                  |  1 +
 .../haiku/rag/capabilities/_base.py           | 30 +++++++++++++--
 tests/capabilities/test_capabilities.py       | 38 +++++++++++++++++++
 3 files changed, 65 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 164d728d..c0876c50 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@
 ### Fixed
 
 - A capability that reaches its request limit keeps its cite tool for two further model requests while its other tools are removed, so an exhausted run can still register citations.
+- A cited chunk id that does not match a retrieved id exactly resolves to the closest one above a 0.75 similarity cutoff, recovering ids mistyped from search results.
 - 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
diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py
index 4471d162..66e09efd 100644
--- a/haiku_rag_slim/haiku/rag/capabilities/_base.py
+++ b/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
 
@@ -30,6 +31,27 @@ from haiku.rag.tools.search import build_binary_parts_from_results
 CITATION_GRACE_REQUESTS = 2
 """Requests the cite tool outlives this capability's other tools by."""
 
+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:
@@ -293,11 +315,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:
diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py
index 5f8146bb..45dc36cf 100644
--- a/tests/capabilities/test_capabilities.py
+++ b/tests/capabilities/test_capabilities.py
@@ -353,6 +353,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())

From a528ab912f590a529246611f30bba1f2ad22ccf1 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Wed, 29 Jul 2026 15:30:26 +0300
Subject: [PATCH 3/9] Record per-case retrieval diagnostics in evaluations

Count tool traffic from the message history, where refused and repeated
calls stay visible, unlike state.searches which is keyed by query.
---
 CHANGELOG.md                                 |  4 ++
 evaluations/evaluations/benchmark.py         |  8 ++++
 evaluations/evaluations/capability_runner.py | 49 ++++++++++++++++++++
 evaluations/tests/test_capability_runner.py  | 45 +++++++++++++++++-
 4 files changed, 104 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index c0876c50..7c775455 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,10 @@
 # Changelog
 ## [Unreleased]
 
+### Added
+
+- `evaluations run` records `cited_chunk_ids`, `searched_uris`, `n_searches`, `n_search_calls`, `n_rejected_calls`, `n_executions`, `n_requests` and `budget_spent` 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.
diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py
index b23d1826..f2339b82 100644
--- a/evaluations/evaluations/benchmark.py
+++ b/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_calls", result.n_rejected_calls)
+        set_eval_attribute("n_executions", result.n_executions)
+        set_eval_attribute("n_requests", result.n_requests)
+        set_eval_attribute("budget_spent", result.budget_spent)
         return result.answer
 
     report = await _evaluate(answer_question)
diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py
index 3957d737..d5a4203c 100644
--- a/evaluations/evaluations/capability_runner.py
+++ b/evaluations/evaluations/capability_runner.py
@@ -4,6 +4,12 @@ from pathlib import Path
 from typing import Any, Protocol, cast
 
 from pydantic_ai import Agent
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelResponse,
+    ToolCallPart,
+    ToolReturnPart,
+)
 from pydantic_ai.models import Model
 
 from haiku.rag.capabilities import RAGCapabilityBase
@@ -29,6 +35,41 @@ 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_calls: int = 0
+    n_requests: int = 0
+    budget_spent: bool = False
+
+
+def _count_tool_traffic(
+    messages: list[ModelMessage], namespace: str
+) -> tuple[int, int, int]:
+    """Count search calls, rejected calls and model requests in a run.
+
+    ``state.searches`` is keyed by query, so it collapses repeated queries and
+    never records a call the capability refused. Counting the message history
+    instead gives the real number of attempts, which is what shows whether a
+    case ran out of budget.
+    """
+    search_tool = f"{namespace}_search"
+    search_calls = 0
+    rejected = 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
+        rejected += sum(
+            1
+            for part in message.parts
+            if isinstance(part, ToolReturnPart) and part.outcome == "failed"
+        )
+    return search_calls, rejected, requests
 
 
 @dataclass
@@ -100,6 +141,10 @@ async def run_capability_question(
     executions = getattr(state, "executions", None)
     n_executions = len(executions) if executions is not None else 0
 
+    n_search_calls, n_rejected_calls, n_requests = _count_tool_traffic(
+        agent_result.all_messages(), capability.state_namespace
+    )
+
     return CapabilityRunResult(
         answer=agent_result.output,
         cited_uris=cited_uris,
@@ -107,4 +152,8 @@ async def run_capability_question(
         searched_uris=searched_uris,
         n_searches=len(typed.searches),
         n_executions=n_executions,
+        n_search_calls=n_search_calls,
+        n_rejected_calls=n_rejected_calls,
+        n_requests=n_requests,
+        budget_spent=n_rejected_calls > 0,
     )
diff --git a/evaluations/tests/test_capability_runner.py b/evaluations/tests/test_capability_runner.py
index bc310ac4..cd1b0dea 100644
--- a/evaluations/tests/test_capability_runner.py
+++ b/evaluations/tests/test_capability_runner.py
@@ -2,14 +2,55 @@ from types import SimpleNamespace
 from unittest.mock import AsyncMock, patch
 
 import pytest
+from pydantic_ai.messages import (
+    ModelRequest,
+    ModelResponse,
+    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
 
 
+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")]),
+    ]
+
+    search_calls, rejected, requests = _count_tool_traffic(messages, "analysis")
+
+    assert search_calls == 2
+    assert rejected == 1
+    assert requests == 2
+
+
 async def test_runs_rag_capability_without_legacy_capability_layer(tmp_path):
     result = await run_capability_question(
         create_rag,
@@ -49,7 +90,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,

From 485a8f901ec3f30c2ae14748ecab1370b32cee57 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 30 Jul 2026 09:39:16 +0300
Subject: [PATCH 4/9] Document the real max_searches default

Both configuration pages said 3; the default has been 5.
---
 docs/configuration/index.md | 2 +-
 docs/configuration/qa.md    | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/docs/configuration/index.md b/docs/configuration/index.md
index ad55be29..7b475081 100644
--- a/docs/configuration/index.md
+++ b/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
diff --git a/docs/configuration/qa.md b/docs/configuration/qa.md
index aa4f4681..bb20f893 100644
--- a/docs/configuration/qa.md
+++ b/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.

From 68238f0b6a2e609a1d42b36425bf1558728f333c Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 30 Jul 2026 11:21:18 +0300
Subject: [PATCH 5/9] Name the readlines workaround when sandbox code iterates
 a file

50 executions across 40 cases in an 822-case run died on
'_io.TextIOWrapper' object is not iterable, and those cases scored 37.5%
judged against 60.1% and cited 12.5% against 55.2%.
---
 CHANGELOG.md                                  |  1 +
 .../haiku/rag/capabilities/analysis.py        | 20 ++++++++++++-
 tests/capabilities/test_capabilities.py       | 29 +++++++++++++++++++
 3 files changed, 49 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7c775455..838472a8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -11,6 +11,7 @@
 
 ### 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 model requests while its other tools are removed, so an exhausted run can still register citations.
 - A cited chunk id that does not match a retrieved id exactly resolves to the closest one above a 0.75 similarity cutoff, recovering ids mistyped from search results.
 - 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.
diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py
index 2d754840..cea8a878 100644
--- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py
+++ b/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."""
@@ -103,7 +118,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]:
diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py
index 45dc36cf..ac99b1d9 100644
--- a/tests/capabilities/test_capabilities.py
+++ b/tests/capabilities/test_capabilities.py
@@ -589,6 +589,35 @@ async def test_cite_tool_is_withdrawn_after_the_grace_window(temp_db_path):
     assert kept == []
 
 
+@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,

From 721acbcf3802a9d19200a4ff0dedc833895eff47 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 30 Jul 2026 15:50:06 +0300
Subject: [PATCH 6/9] Address review on PR #524

- _budget_notice no longer names the cite tool after prepare_tools has
  withdrawn it; the post-grace state gets the plain no-tools text back.
- Split search-budget rejections from any failed tool call: the code tool
  raises ToolFailed for every error in model-written Python, so
  budget_spent was true for a ZeroDivisionError.
- docs/capabilities/rag.md described the old single-turn removal.
- Drop the rationale clause from the CHANGELOG entry.
---
 CHANGELOG.md                                  |  4 +-
 docs/capabilities/rag.md                      |  2 +-
 evaluations/evaluations/benchmark.py          |  3 +-
 evaluations/evaluations/capability_runner.py  | 39 +++++++++++--------
 evaluations/tests/test_capability_runner.py   | 32 ++++++++++++++-
 .../haiku/rag/capabilities/_base.py           | 13 ++++++-
 tests/capabilities/test_capabilities.py       |  9 +++++
 7 files changed, 80 insertions(+), 22 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 838472a8..3e25dbba 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,7 +3,7 @@
 
 ### Added
 
-- `evaluations run` records `cited_chunk_ids`, `searched_uris`, `n_searches`, `n_search_calls`, `n_rejected_calls`, `n_executions`, `n_requests` and `budget_spent` as eval attributes alongside `cited_uris`.
+- `evaluations run` records `cited_chunk_ids`, `searched_uris`, `n_searches`, `n_search_calls`, `n_rejected_searches`, `n_failed_tools`, `n_executions`, `n_requests` and `budget_spent` as eval attributes alongside `cited_uris`.
 
 ### Changed
 
@@ -12,7 +12,7 @@
 ### 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 model requests while its other tools are removed, so an exhausted run can still register citations.
+- A capability that reaches its request limit keeps its cite tool for two further model requests 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, recovering ids mistyped from search results.
 - 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.
 
diff --git a/docs/capabilities/rag.md b/docs/capabilities/rag.md
index af3b2959..6dce2443 100644
--- a/docs/capabilities/rag.md
+++ b/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, so the model can register citations before answering from evidence already gathered. 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
 
diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py
index f2339b82..51cef761 100644
--- a/evaluations/evaluations/benchmark.py
+++ b/evaluations/evaluations/benchmark.py
@@ -438,7 +438,8 @@ async def run_qa_benchmark(
         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_calls", result.n_rejected_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)
         set_eval_attribute("budget_spent", result.budget_spent)
diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py
index d5a4203c..03da8773 100644
--- a/evaluations/evaluations/capability_runner.py
+++ b/evaluations/evaluations/capability_runner.py
@@ -36,24 +36,30 @@ class CapabilityRunResult:
     n_searches: int = 0
     n_executions: int = 0
     n_search_calls: int = 0
-    n_rejected_calls: int = 0
+    n_rejected_searches: int = 0
+    n_failed_tools: int = 0
     n_requests: int = 0
     budget_spent: bool = False
 
 
 def _count_tool_traffic(
     messages: list[ModelMessage], namespace: str
-) -> tuple[int, int, int]:
-    """Count search calls, rejected calls and model requests in a run.
+) -> tuple[int, int, int, int]:
+    """Count search calls, failed calls and model requests in a run.
 
     ``state.searches`` is keyed by query, so it collapses repeated queries and
     never records a call the capability refused. Counting the message history
-    instead gives the real number of attempts, which is what shows whether a
-    case ran out of budget.
+    instead gives the real number of attempts.
+
+    Failures are split by tool. Only the search tool fails for want of budget,
+    whereas the code tool raises ``ToolFailed`` for any error in model-written
+    Python, so counting every failure together would report a ``ZeroDivisionError``
+    as budget exhaustion.
     """
     search_tool = f"{namespace}_search"
     search_calls = 0
-    rejected = 0
+    rejected_searches = 0
+    failed_tools = 0
     requests = 0
     for message in messages:
         if isinstance(message, ModelResponse):
@@ -64,12 +70,12 @@ def _count_tool_traffic(
                 if isinstance(part, ToolCallPart) and part.tool_name == search_tool
             )
             continue
-        rejected += sum(
-            1
-            for part in message.parts
-            if isinstance(part, ToolReturnPart) and part.outcome == "failed"
-        )
-    return search_calls, rejected, requests
+        for part in message.parts:
+            if isinstance(part, ToolReturnPart) and part.outcome == "failed":
+                failed_tools += 1
+                if part.tool_name == search_tool:
+                    rejected_searches += 1
+    return search_calls, rejected_searches, failed_tools, requests
 
 
 @dataclass
@@ -141,8 +147,8 @@ async def run_capability_question(
     executions = getattr(state, "executions", None)
     n_executions = len(executions) if executions is not None else 0
 
-    n_search_calls, n_rejected_calls, n_requests = _count_tool_traffic(
-        agent_result.all_messages(), capability.state_namespace
+    n_search_calls, n_rejected_searches, n_failed_tools, n_requests = (
+        _count_tool_traffic(agent_result.all_messages(), capability.state_namespace)
     )
 
     return CapabilityRunResult(
@@ -153,7 +159,8 @@ async def run_capability_question(
         n_searches=len(typed.searches),
         n_executions=n_executions,
         n_search_calls=n_search_calls,
-        n_rejected_calls=n_rejected_calls,
+        n_rejected_searches=n_rejected_searches,
+        n_failed_tools=n_failed_tools,
         n_requests=n_requests,
-        budget_spent=n_rejected_calls > 0,
+        budget_spent=n_rejected_searches > 0,
     )
diff --git a/evaluations/tests/test_capability_runner.py b/evaluations/tests/test_capability_runner.py
index cd1b0dea..8e3cf317 100644
--- a/evaluations/tests/test_capability_runner.py
+++ b/evaluations/tests/test_capability_runner.py
@@ -18,6 +18,34 @@ from haiku.rag.capabilities.rag import create_capability as create_rag
 from haiku.rag.config.models import AppConfig
 
 
+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")]),
+    ]
+
+    search_calls, rejected_searches, failed_tools, requests = _count_tool_traffic(
+        messages, "analysis"
+    )
+
+    assert search_calls == 0
+    assert rejected_searches == 0
+    assert failed_tools == 1
+    assert requests == 2
+
+
 def test_count_tool_traffic_counts_attempts_not_distinct_queries():
     """Rejected and repeated calls both count; `state.searches` hides them."""
     messages = [
@@ -44,7 +72,9 @@ def test_count_tool_traffic_counts_attempts_not_distinct_queries():
         ModelResponse(parts=[TextPart("done")]),
     ]
 
-    search_calls, rejected, requests = _count_tool_traffic(messages, "analysis")
+    search_calls, rejected, _failed, requests = _count_tool_traffic(
+        messages, "analysis"
+    )
 
     assert search_calls == 2
     assert rejected == 1
diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py
index 66e09efd..9bd51bf9 100644
--- a/haiku_rag_slim/haiku/rag/capabilities/_base.py
+++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py
@@ -171,7 +171,18 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
         return request_context
 
     def _budget_notice(self) -> str | None:
-        """Tell the model which of this capability's budgets just ran out."""
+        """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 "
diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py
index ac99b1d9..e062bd9c 100644
--- a/tests/capabilities/test_capabilities.py
+++ b/tests/capabilities/test_capabilities.py
@@ -583,10 +583,19 @@ async def test_cite_tool_is_withdrawn_after_the_grace_window(temp_db_path):
     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.request_count = 4
     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

From b62d6e920b200f3628d7a47d17a50378d48551af Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 30 Jul 2026 15:57:48 +0300
Subject: [PATCH 7/9] Drop budget_spent from eval attributes

It was n_rejected_searches > 0 recorded next to the integer it derived
from, and the name overclaimed: analysis_execute_code also raises
ToolFailed when execute_count exceeds max_executions, which the flag
never saw. Callers can compare the counters directly.
---
 CHANGELOG.md                                 |  2 +-
 evaluations/evaluations/benchmark.py         |  1 -
 evaluations/evaluations/capability_runner.py | 17 +++++++++--------
 3 files changed, 10 insertions(+), 10 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 3e25dbba..e6059883 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,7 +3,7 @@
 
 ### Added
 
-- `evaluations run` records `cited_chunk_ids`, `searched_uris`, `n_searches`, `n_search_calls`, `n_rejected_searches`, `n_failed_tools`, `n_executions`, `n_requests` and `budget_spent` as eval attributes alongside `cited_uris`.
+- `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
 
diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py
index 51cef761..2541ca05 100644
--- a/evaluations/evaluations/benchmark.py
+++ b/evaluations/evaluations/benchmark.py
@@ -442,7 +442,6 @@ async def run_qa_benchmark(
         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)
-        set_eval_attribute("budget_spent", result.budget_spent)
         return result.answer
 
     report = await _evaluate(answer_question)
diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py
index 03da8773..7fd8e500 100644
--- a/evaluations/evaluations/capability_runner.py
+++ b/evaluations/evaluations/capability_runner.py
@@ -39,7 +39,6 @@ class CapabilityRunResult:
     n_rejected_searches: int = 0
     n_failed_tools: int = 0
     n_requests: int = 0
-    budget_spent: bool = False
 
 
 def _count_tool_traffic(
@@ -48,13 +47,16 @@ def _count_tool_traffic(
     """Count search calls, failed calls and model requests in a run.
 
     ``state.searches`` is keyed by query, so it collapses repeated queries and
-    never records a call the capability refused. Counting the message history
-    instead gives the real number of attempts.
+    never records a call the capability refused. The capability object cannot be
+    read instead: ``for_run`` hands the run a ``replace()`` copy, so the outer
+    instance's counters stay at zero. Counting the message history is the only
+    way to see the real number of attempts.
 
-    Failures are split by tool. Only the search tool fails for want of budget,
-    whereas the code tool raises ``ToolFailed`` for any error in model-written
-    Python, so counting every failure together would report a ``ZeroDivisionError``
-    as budget exhaustion.
+    Search failures are counted apart because the search tool fails only when
+    its budget is spent. A failed code call is ambiguous — an exhausted
+    execution budget and any error in model-written Python both surface as
+    ``ToolFailed``, distinguishable in the history only by message text — so
+    ``n_failed_tools`` covers both without claiming to tell them apart.
     """
     search_tool = f"{namespace}_search"
     search_calls = 0
@@ -162,5 +164,4 @@ async def run_capability_question(
         n_rejected_searches=n_rejected_searches,
         n_failed_tools=n_failed_tools,
         n_requests=n_requests,
-        budget_spent=n_rejected_searches > 0,
     )

From 27ed0b3bb35de621290282dcc1d707041b9aeba5 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 30 Jul 2026 18:17:26 +0300
Subject: [PATCH 8/9] Send analysis to the sandbox when only its search budget
 is spent

The notice told the model to answer from what it had the moment
qa.max_searches ran out, while up to 15 code executions remained and
in-code search() does not count against that budget. It now names the
spent tool and points at whichever evidence tool still has budget,
falling back to answer-and-cite only when none do.

Also count RetryPromptPart in n_failed_tools: _cite rejects with
ModelRetry, so a run whose every cite attempt was refused reported zero
failures. And note that n_requests is the run's request count, which
tracks a capability's own budget only while it stays loaded.
---
 evaluations/evaluations/capability_runner.py  | 15 ++++++-
 evaluations/tests/test_capability_runner.py   | 26 ++++++++++++
 .../haiku/rag/capabilities/_base.py           | 23 +++++++++--
 .../haiku/rag/capabilities/analysis.py        |  6 +++
 tests/capabilities/test_capabilities.py       | 40 +++++++++++++++++++
 5 files changed, 106 insertions(+), 4 deletions(-)

diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py
index 7fd8e500..65331838 100644
--- a/evaluations/evaluations/capability_runner.py
+++ b/evaluations/evaluations/capability_runner.py
@@ -7,6 +7,7 @@ from pydantic_ai import Agent
 from pydantic_ai.messages import (
     ModelMessage,
     ModelResponse,
+    RetryPromptPart,
     ToolCallPart,
     ToolReturnPart,
 )
@@ -57,6 +58,16 @@ def _count_tool_traffic(
     execution budget and any error in model-written Python both surface as
     ``ToolFailed``, distinguishable in the history only by message text — so
     ``n_failed_tools`` covers both without claiming to tell them apart.
+
+    ``n_failed_tools`` counts ``RetryPromptPart`` as well as failed returns.
+    ``_cite`` rejects a call with ``ModelRetry`` rather than ``ToolFailed``, and
+    only the latter sets ``outcome="failed"``, so counting returns alone would
+    leave a run whose every cite attempt was rejected reporting zero failures.
+
+    ``requests`` is the run's model-request count, which equals the capability's
+    own ``request_count`` budget only while the capability is loaded for the
+    whole run. A deferred capability skips hooks until it loads, so a router
+    turn before that is counted here and not against its limit.
     """
     search_tool = f"{namespace}_search"
     search_calls = 0
@@ -73,7 +84,9 @@ def _count_tool_traffic(
             )
             continue
         for part in message.parts:
-            if isinstance(part, ToolReturnPart) and part.outcome == "failed":
+            if isinstance(part, RetryPromptPart):
+                failed_tools += 1
+            elif isinstance(part, ToolReturnPart) and part.outcome == "failed":
                 failed_tools += 1
                 if part.tool_name == search_tool:
                     rejected_searches += 1
diff --git a/evaluations/tests/test_capability_runner.py b/evaluations/tests/test_capability_runner.py
index 8e3cf317..cd8c979e 100644
--- a/evaluations/tests/test_capability_runner.py
+++ b/evaluations/tests/test_capability_runner.py
@@ -5,6 +5,7 @@ import pytest
 from pydantic_ai.messages import (
     ModelRequest,
     ModelResponse,
+    RetryPromptPart,
     TextPart,
     ToolCallPart,
     ToolReturnPart,
@@ -18,6 +19,31 @@ from haiku.rag.capabilities.rag import create_capability as create_rag
 from haiku.rag.config.models import AppConfig
 
 
+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")]),
+    ]
+
+    _search_calls, rejected_searches, failed_tools, _requests = _count_tool_traffic(
+        messages, "analysis"
+    )
+
+    assert failed_tools == 1
+    assert 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 = [
diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py
index 9bd51bf9..603d224d 100644
--- a/haiku_rag_slim/haiku/rag/capabilities/_base.py
+++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py
@@ -191,11 +191,20 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
                 "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"{', '.join(sorted(spent))}; further calls to them fail. Answer "
-                f"from the evidence already gathered and call "
-                f"{self._cite_tool_name} with the chunk_ids supporting it."
+                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
 
@@ -225,6 +234,14 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
             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.
+
+        Drives the spent-budget notice: while one of these still has budget the
+        model must be pointed at it, not told to answer from what it has.
+        """
+        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:
diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py
index cea8a878..95fa155e 100644
--- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py
+++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py
@@ -85,6 +85,12 @@ 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:
diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py
index e062bd9c..715330a3 100644
--- a/tests/capabilities/test_capabilities.py
+++ b/tests/capabilities/test_capabilities.py
@@ -494,6 +494,46 @@ async def test_spent_search_budget_is_announced_but_keeps_the_tool(rag_db):
     assert "spent its budget for rag_search" in seen_instructions[1]
 
 
+@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
+    assert "Answer from the evidence already gathered" not in notice
+
+    # Once the code budget is gone too there is nowhere left to go.
+    capability.execute_count = config.analysis.max_executions
+    notice = capability._budget_notice()
+    assert notice is not None
+    assert "Answer from the evidence already gathered" in notice
+
+
+@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 "Answer from the evidence already gathered" in notice
+
+
 @pytest.mark.asyncio
 async def test_spent_execution_budget_joins_the_notice(temp_db_path):
     config = AppConfig()

From c137305468ef4ba249040a397945cd3ea065d800 Mon Sep 17 00:00:00 2001
From: Yiorgis Gozadinos <ggozadinos@gmail.com>
Date: Thu, 30 Jul 2026 18:50:05 +0300
Subject: [PATCH 9/9] Scope the limit notice and spend the cite window on own
 turns only

The request-limit notice said only the cite tool remained available, but
chat registers rag and analysis in one agent, so exhausting analysis
claimed rag_search was gone too. Scoped to the capability's own tools.

The cite window was counted over every model request once loaded, so
turns spent on another capability expired it before the model was ever
placed where citing was the obvious move. Count only requests whose
preceding response called one of this capability's tools; engagement is
also the only thing that can loop, which is all the bound guards against.

Also: _count_tool_traffic returns a named tuple rather than four bare
ints, and counts failures only for this capability's tools, so host-tool
retries and output-validation retries no longer read as its failures.
---
 CHANGELOG.md                                  |  5 +-
 docs/capabilities/analysis.md                 |  2 +-
 docs/capabilities/rag.md                      |  2 +-
 evaluations/evaluations/capability_runner.py  | 75 +++++++++++--------
 evaluations/tests/test_capability_runner.py   | 34 ++++-----
 .../haiku/rag/capabilities/_base.py           | 56 +++++++++-----
 tests/capabilities/test_capabilities.py       | 53 +++++++++++--
 7 files changed, 148 insertions(+), 79 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e6059883..138cb2f9 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,12 +12,13 @@
 ### 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 model requests 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, recovering ids mistyped from search results.
+- 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
diff --git a/docs/capabilities/analysis.md b/docs/capabilities/analysis.md
index f382868a..19fbe4d6 100644
--- a/docs/capabilities/analysis.md
+++ b/docs/capabilities/analysis.md
@@ -4,7 +4,7 @@
 
 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, `analysis_search` and `analysis_execute_code` are removed while `analysis_cite` remains for two further requests, so the model can register citations before answering 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`.
 
diff --git a/docs/capabilities/rag.md b/docs/capabilities/rag.md
index 6dce2443..00a912bb 100644
--- a/docs/capabilities/rag.md
+++ b/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, `rag_search` is removed while `rag_cite` remains for two further requests, so the model can register citations before answering from evidence already gathered. 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
 
diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py
index 65331838..b87a6c4c 100644
--- a/evaluations/evaluations/capability_runner.py
+++ b/evaluations/evaluations/capability_runner.py
@@ -1,7 +1,7 @@
 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 (
@@ -42,32 +42,32 @@ class CapabilityRunResult:
     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
-) -> tuple[int, int, int, int]:
+    messages: list[ModelMessage], namespace: str, tool_names: frozenset[str]
+) -> ToolTraffic:
     """Count search calls, failed calls and model requests in a run.
 
-    ``state.searches`` is keyed by query, so it collapses repeated queries and
-    never records a call the capability refused. The capability object cannot be
-    read instead: ``for_run`` hands the run a ``replace()`` copy, so the outer
-    instance's counters stay at zero. Counting the message history is the only
-    way to see the real number of attempts.
+    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.
 
-    Search failures are counted apart because the search tool fails only when
-    its budget is spent. A failed code call is ambiguous — an exhausted
-    execution budget and any error in model-written Python both surface as
-    ``ToolFailed``, distinguishable in the history only by message text — so
-    ``n_failed_tools`` covers both without claiming to tell them apart.
+    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_failed_tools`` counts ``RetryPromptPart`` as well as failed returns.
-    ``_cite`` rejects a call with ``ModelRetry`` rather than ``ToolFailed``, and
-    only the latter sets ``outcome="failed"``, so counting returns alone would
-    leave a run whose every cite attempt was rejected reporting zero failures.
-
-    ``requests`` is the run's model-request count, which equals the capability's
-    own ``request_count`` budget only while the capability is loaded for the
-    whole run. A deferred capability skips hooks until it loads, so a router
-    turn before that is counted here and not against its limit.
+    ``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
@@ -84,13 +84,22 @@ def _count_tool_traffic(
             )
             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 isinstance(part, ToolReturnPart) and part.outcome == "failed":
+            elif part.outcome == "failed":
                 failed_tools += 1
                 if part.tool_name == search_tool:
                     rejected_searches += 1
-    return search_calls, rejected_searches, failed_tools, requests
+    return ToolTraffic(
+        n_search_calls=search_calls,
+        n_rejected_searches=rejected_searches,
+        n_failed_tools=failed_tools,
+        n_requests=requests,
+    )
 
 
 @dataclass
@@ -162,8 +171,10 @@ async def run_capability_question(
     executions = getattr(state, "executions", None)
     n_executions = len(executions) if executions is not None else 0
 
-    n_search_calls, n_rejected_searches, n_failed_tools, n_requests = (
-        _count_tool_traffic(agent_result.all_messages(), capability.state_namespace)
+    traffic = _count_tool_traffic(
+        agent_result.all_messages(),
+        capability.state_namespace,
+        capability.tool_names,
     )
 
     return CapabilityRunResult(
@@ -171,10 +182,14 @@ async def run_capability_question(
         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=n_search_calls,
-        n_rejected_searches=n_rejected_searches,
-        n_failed_tools=n_failed_tools,
-        n_requests=n_requests,
+        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,
     )
diff --git a/evaluations/tests/test_capability_runner.py b/evaluations/tests/test_capability_runner.py
index cd8c979e..9256d3d3 100644
--- a/evaluations/tests/test_capability_runner.py
+++ b/evaluations/tests/test_capability_runner.py
@@ -18,6 +18,10 @@ 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."""
@@ -36,12 +40,10 @@ def test_count_tool_traffic_sees_a_rejected_cite_call():
         ModelResponse(parts=[TextPart("done")]),
     ]
 
-    _search_calls, rejected_searches, failed_tools, _requests = _count_tool_traffic(
-        messages, "analysis"
-    )
+    traffic = _count_tool_traffic(messages, "analysis", ANALYSIS_TOOLS)
 
-    assert failed_tools == 1
-    assert rejected_searches == 0
+    assert traffic.n_failed_tools == 1
+    assert traffic.n_rejected_searches == 0
 
 
 def test_count_tool_traffic_separates_search_rejections_from_code_errors():
@@ -62,14 +64,12 @@ def test_count_tool_traffic_separates_search_rejections_from_code_errors():
         ModelResponse(parts=[TextPart("done")]),
     ]
 
-    search_calls, rejected_searches, failed_tools, requests = _count_tool_traffic(
-        messages, "analysis"
-    )
+    traffic = _count_tool_traffic(messages, "analysis", ANALYSIS_TOOLS)
 
-    assert search_calls == 0
-    assert rejected_searches == 0
-    assert failed_tools == 1
-    assert requests == 2
+    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():
@@ -98,13 +98,11 @@ def test_count_tool_traffic_counts_attempts_not_distinct_queries():
         ModelResponse(parts=[TextPart("done")]),
     ]
 
-    search_calls, rejected, _failed, requests = _count_tool_traffic(
-        messages, "analysis"
-    )
+    traffic = _count_tool_traffic(messages, "analysis", ANALYSIS_TOOLS)
 
-    assert search_calls == 2
-    assert rejected == 1
-    assert requests == 2
+    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):
diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py
index 603d224d..ff7a2b0a 100644
--- a/haiku_rag_slim/haiku/rag/capabilities/_base.py
+++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py
@@ -12,6 +12,8 @@ from pydantic_ai.messages import (
     InstructionPart,
     ModelMessage,
     ModelRequest,
+    ModelResponse,
+    ToolCallPart,
     ToolReturn,
     ToolReturnPart,
     UserPromptPart,
@@ -29,7 +31,13 @@ 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 the cite tool outlives this capability's other tools by."""
+"""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.
@@ -105,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
@@ -122,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)
@@ -138,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
@@ -167,6 +188,10 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
                         InstructionPart(content=instruction, dynamic=True),
                     ],
                 )
+        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
 
@@ -186,9 +211,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
         if self._request_limit_reached:
             return (
                 f"The {self.state_namespace} capability has reached its request "
-                f"limit. Only {self._cite_tool_name} remains available: register "
-                "the chunk_ids supporting your answer, then answer from the "
-                "evidence already gathered."
+                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))
@@ -215,14 +240,10 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
     ) -> list[ToolDefinition]:
         """Remove this capability's tools past its limit, cite tool last.
 
-        The cite tool outlives the others by ``CITATION_GRACE_REQUESTS`` so a run
-        that exhausts its budget can still record the evidence it gathered.
-
         Tools whose own budget is spent stay declared on purpose. Removing one
-        makes a model that calls it anyway hit ``Unknown tool name``, which is
-        charged against the agent's unknown-tool retry budget and kills the run
-        after two attempts. A spent tool that keeps failing wastes requests; a
-        withdrawn one loses the whole answer.
+        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]
@@ -235,11 +256,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
         ]
 
     def _evidence_tool_names(self) -> set[str]:
-        """Tools that can bring new evidence into the run.
-
-        Drives the spent-budget notice: while one of these still has budget the
-        model must be pointed at it, not told to answer from what it has.
-        """
+        """Tools that can bring new evidence into the run."""
         return {f"{self.state_namespace}_search"}
 
     def _spent_tool_names(self) -> set[str]:
@@ -264,10 +281,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
 
     @property
     def _citation_grace_expired(self) -> bool:
-        return (
-            self.request_limit is not None
-            and self.request_count >= self.request_limit + CITATION_GRACE_REQUESTS
-        )
+        # 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]
diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py
index 715330a3..6bbf6141 100644
--- a/tests/capabilities/test_capabilities.py
+++ b/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
@@ -494,6 +498,40 @@ async def test_spent_search_budget_is_announced_but_keeps_the_tool(rag_db):
     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.
@@ -511,13 +549,13 @@ async def test_spent_search_notice_points_at_code_while_it_has_budget(temp_db_pa
     assert notice is not None
     assert "analysis_search" in notice
     assert "analysis_execute_code" in notice
-    assert "Answer from the evidence already gathered" not in notice
 
-    # Once the code budget is gone too there is nowhere left to go.
+    # 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 "Answer from the evidence already gathered" in notice
+    assert "analysis_execute_code" in notice
+    assert capability._evidence_tool_names() <= capability._spent_tool_names()
 
 
 @pytest.mark.asyncio
@@ -531,7 +569,8 @@ async def test_spent_search_notice_tells_rag_to_answer(temp_db_path):
     notice = capability._budget_notice()
 
     assert notice is not None
-    assert "Answer from the evidence already gathered" in notice
+    assert "rag_search" in notice
+    assert capability._evidence_tool_names() == {"rag_search"}
 
 
 @pytest.mark.asyncio
@@ -591,7 +630,7 @@ async def test_exhausted_run_can_still_register_citations(rag_db):
     deps = Deps()
 
     async with agent.iter("question", deps=deps) as run:
-        async for node in run:
+        async for _node in run:
             if chunk_id is None:
                 searches = deps.state.get("rag", {}).get("searches") or {}
                 for results in searches.values():
@@ -626,7 +665,7 @@ async def test_cite_tool_is_withdrawn_after_the_grace_window(temp_db_path):
     notice = capability._budget_notice()
     assert notice is not None and "rag_cite" in notice
 
-    capability.request_count = 4
+    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: