Keep the cite tool past a capability's request limit

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

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
 

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

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

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,