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

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

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 = [

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"{names}; further calls to them fail. Answer from the evidence "
-                f"from the evidence already gathered and call "
+                f"already gathered and call {self._cite_tool_name} with the "
-                f"{self._cite_tool_name} with the chunk_ids supporting it."
+                "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:

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:

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