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

@@ -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.
 

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
 

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)

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,
     )

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

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 "

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