Price qa.max_searches in search units

Searches a model emits in one response share a budget unit, up to
FREE_SIBLINGS_PER_ROUND (3) per unit; sequential searches pay one unit
each, as before. Grouping keys on RunContext.run_step, which pydantic-ai
increments once per model request. A budget-rejected round fails all its
remaining siblings, and tracking resets per run.

Glimmer opens most questions with a burst of ~3 rephrasings in a single
response (95.8% of its three-search ORB cases are one-response bursts),
spending 3 of 5 searches before reading anything. Pass rate at 3 calls
equals 1 call, so a burst is priced as one probe.

CHANGELOG.md

@@ -2,6 +2,12 @@
 
 ## [Unreleased]
 
+### Changed
+
+- `qa.max_searches` counts search units: searches a model emits in one
+  response share a unit, up to 3 per unit; sequential searches pay one unit
+  each.
+
 ## [0.81.0] - 2026-09-01
 
 ### Added

haiku_rag_slim/haiku/rag/capabilities/_base.py

@@ -61,6 +61,15 @@ 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.
 """
 
+FREE_SIBLINGS_PER_ROUND = 3
+"""Searches one budget unit covers when emitted in the same model response.
+
+Calibration knob, sized to the measured modal burst. ``qa.max_searches``
+counts units, so a model rephrasing its query a few times in one response
+spends one unit, while every search of a sequential searcher is a unit of its
+own.
+"""
+
 
 def _ambiguous_retry(error: AmbiguousCitationError) -> ModelRetry:
     """The only way out of an id that names a chunk in two databases.
@@ -177,6 +186,10 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
     rag_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
     resource_lock: asyncio.Lock = field(default_factory=asyncio.Lock, repr=False)
     search_count: int = field(default=0, repr=False)
+    search_step: int = field(default=0, repr=False)
+    """The run_step whose searches are being priced and deduplicated."""
+    step_searches: int = field(default=0, repr=False)
+    step_rejected: bool = field(default=False, repr=False)
     request_count: int = field(default=0, repr=False)
     grace_requests_used: int = field(default=0, repr=False)
     epoch: int = field(default=0, repr=False)
@@ -226,6 +239,9 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
             rag_lock=asyncio.Lock(),
             resource_lock=asyncio.Lock(),
             search_count=0,
+            search_step=0,
+            step_searches=0,
+            step_rejected=False,
             request_count=0,
             grace_requests_used=0,
             epoch=0,
@@ -493,10 +509,19 @@ class RAGCapabilityBase[StateT: EvidenceState](AbstractCapability[Any]):
             retrieved_now=retrieved,
         )
 
-    async def _search(self, query: str, limit: int | None) -> str | ToolReturn:
+    async def _search(
+        self, query: str, limit: int | None, run_step: int
+    ) -> str | ToolReturn:
         assert self.state is not None
-        self.search_count += 1
-        if self.search_count > self._max_searches:
+        if run_step != self.search_step:
+            self.search_step = run_step
+            self.step_searches = 0
+            self.step_rejected = False
+        self.step_searches += 1
+        if (self.step_searches - 1) % FREE_SIBLINGS_PER_ROUND == 0:
+            self.search_count += 1
+        if self.step_rejected or self.search_count > self._max_searches:
+            self.step_rejected = True
             raise ToolFailed(
                 "Search limit reached. Answer the question using "
                 "the results you already have."

haiku_rag_slim/haiku/rag/capabilities/analysis.py

@@ -172,7 +172,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]):
             ctx: RunContext[Any], query: str, limit: int | None = None
         ) -> str | ToolReturn:
             """Search the knowledge base for evidence to analyze."""
-            return await self._with_state(self._search(query, limit))
+            return await self._with_state(self._search(query, limit, ctx.run_step))
 
         async def analysis_execute_code(ctx: RunContext[Any], code: str) -> Any:
             """Execute Python against the sandboxed document filesystem."""

haiku_rag_slim/haiku/rag/capabilities/rag.py

@@ -81,7 +81,7 @@ class RAGCapability(RAGCapabilityBase[RAGState]):
             ctx: RunContext[Any], query: str, limit: int | None = None
         ) -> str | ToolReturn:
             """Search the knowledge base using hybrid vector and full-text search."""
-            return await self._with_state(self._search(query, limit))
+            return await self._with_state(self._search(query, limit, ctx.run_step))
 
         async def rag_cite(ctx: RunContext[Any], chunk_ids: list[str]) -> Any:
             """Register exact search-result chunk IDs as citations for the answer."""

tests/capabilities/test_capabilities.py

@@ -412,7 +412,7 @@ async def test_a_spent_search_budget_fails_the_tool(temp_db_path):
     capability.state = RAGState()
 
     with pytest.raises(ToolFailed, match="Search limit reached"):
-        await capability._search("anything", None)
+        await capability._search("anything", None, 1)
 
 
 def _stub_client(*batches: list[SearchResult]) -> AsyncMock:
@@ -452,7 +452,7 @@ async def _labels_of_search(temp_db_path, *sources: str) -> list[str]:
     client.source_names = sources
 
     with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=client)):
-        returned = await capability._search("cats", None)
+        returned = await capability._search("cats", None, 1)
 
     assert isinstance(returned, ToolReturn)
     assert returned.content is not None
@@ -483,7 +483,9 @@ async def test_a_fruitless_search_says_so(temp_db_path):
     capability.state = RAGState()
     capability.borrowed_rag = _stub_client([])
 
-    assert await capability._search("nothing about this", None) == "No results found."
+    assert (
+        await capability._search("nothing about this", None, 1) == "No results found."
+    )
 
 
 @pytest.mark.asyncio
@@ -500,8 +502,8 @@ async def test_a_narrower_repeat_keeps_what_the_wider_search_returned(temp_db_pa
         [SearchResult(content="first", score=1.0, chunk_id="chunk-1")],
     )
 
-    await capability._search("Figure 3-1", 20)
-    await capability._search("Figure 3-1", None)
+    await capability._search("Figure 3-1", 20, 1)
+    await capability._search("Figure 3-1", None, 2)
 
     stored = capability.state.searches["Figure 3-1"]
     assert [result.chunk_id for result in stored] == [
@@ -525,8 +527,8 @@ async def test_two_databases_holding_one_chunk_id_both_survive(temp_db_path):
         ],
     )
 
-    await capability._search("cats", 20)
-    await capability._search("cats", None)
+    await capability._search("cats", 20, 1)
+    await capability._search("cats", None, 2)
 
     stored = capability.state.searches["cats"]
     assert [(r.source, r.chunk_id) for r in stored] == [
@@ -1105,7 +1107,7 @@ def _record(deps: Deps, namespace: str) -> CapabilityEvidenceRecord:
     return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"])
 
 
-async def _stub_search(self, query: str, _limit: int | None) -> str:
+async def _stub_search(self, query: str, _limit: int | None, _run_step: int) -> str:
     """Record a result the way the real search does, so citing resolves."""
     cast(Any, self.state).searches[query] = [
         SearchResult(content="evidence", score=1.0, chunk_id="chunk-1")

tests/capabilities/test_citation_policy.py

@@ -29,7 +29,7 @@ class Deps:
     state: dict[str, Any] = field(default_factory=dict)
 
 
-async def stub_search(self, query: str, _limit: int | None) -> str:
+async def stub_search(self, query: str, _limit: int | None, _run_step: int) -> str:
     cast(Any, self.state).searches[query] = [
         SearchResult(content="evidence", score=1.0, chunk_id="chunk-1")
     ]

tests/capabilities/test_evidence_wire.py

@@ -437,7 +437,9 @@ REAL_PNG = base64.b64decode(
 )
 
 
-async def _search_with_a_picture(self, query: str, _limit: int | None) -> str:
+async def _search_with_a_picture(
+    self, query: str, _limit: int | None, _run_step: int
+) -> str:
     """Record a result carrying a page image, the way a real search does."""
     cast(Any, self.state).searches[query] = [
         SearchResult(

tests/capabilities/test_search_units.py

@@ -0,0 +1,156 @@
+from dataclasses import dataclass, field
+from typing import Any
+
+import pytest
+from pydantic_ai import Agent
+from pydantic_ai.messages import (
+    ModelResponse,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+)
+from pydantic_ai.models.function import FunctionModel
+from pydantic_ai.run import AgentRunResult
+
+from haiku.rag.capabilities.rag import create_capability as create_rag
+from haiku.rag.config.models import AppConfig
+
+
+@dataclass
+class Deps:
+    state: dict[str, Any] = field(default_factory=dict)
+
+
+def burst_model(bursts: list[list[str]]) -> FunctionModel:
+    """Emit one `rag_search` call per query in each burst, then answer."""
+    responses = 0
+
+    def model_function(_messages, _info) -> ModelResponse:
+        nonlocal responses
+        responses += 1
+        if responses <= len(bursts):
+            return ModelResponse(
+                parts=[
+                    ToolCallPart("rag_search", {"query": query})
+                    for query in bursts[responses - 1]
+                ]
+            )
+        return ModelResponse(parts=[TextPart("done")])
+
+    return FunctionModel(model_function)
+
+
+def burst_agent(
+    bursts: list[list[str]], db_path, max_searches: int
+) -> Agent[Deps, str]:
+    config = AppConfig()
+    config.qa.max_searches = max_searches
+    return Agent(
+        burst_model(bursts),
+        deps_type=Deps,
+        capabilities=[create_rag(db_path=db_path, config=config, defer_loading=False)],
+    )
+
+
+def search_returns(result: AgentRunResult[Any]) -> list[ToolReturnPart]:
+    return [
+        part
+        for message in result.all_messages()
+        for part in message.parts
+        if isinstance(part, ToolReturnPart) and part.tool_name == "rag_search"
+    ]
+
+
+def outcomes(result: AgentRunResult[Any]) -> list[str]:
+    return [
+        "failed" if part.outcome == "failed" else "ok"
+        for part in search_returns(result)
+    ]
+
+
+@pytest.mark.asyncio
+async def test_a_burst_in_one_response_consumes_one_unit(rag_db):
+    """Three searches emitted together cost one unit and run in emission order."""
+    agent = burst_agent([["ai", "machine learning", "deep learning"]], rag_db, 1)
+
+    result = await agent.run("question", deps=Deps())
+
+    assert outcomes(result) == ["ok", "ok", "ok"]
+    calls = [
+        part
+        for message in result.all_messages()
+        for part in message.parts
+        if isinstance(part, ToolCallPart) and part.tool_name == "rag_search"
+    ]
+    assert [part.tool_call_id for part in search_returns(result)] == [
+        part.tool_call_id for part in calls
+    ]
+
+
+@pytest.mark.asyncio
+async def test_sequential_searches_pay_one_unit_each(rag_db):
+    agent = burst_agent([["ai"], ["machine learning"]], rag_db, 1)
+
+    result = await agent.run("question", deps=Deps())
+
+    assert outcomes(result) == ["ok", "failed"]
+    assert "Search limit reached" in str(search_returns(result)[1].content)
+
+
+@pytest.mark.asyncio
+async def test_max_searches_zero_fails_every_sibling(rag_db):
+    agent = burst_agent([["ai", "machine learning", "deep learning"]], rag_db, 0)
+
+    result = await agent.run("question", deps=Deps())
+
+    assert outcomes(result) == ["failed", "failed", "failed"]
+
+
+@pytest.mark.asyncio
+async def test_a_rejected_round_fails_all_its_siblings(rag_db):
+    agent = burst_agent([["ai"], ["ml", "deep learning", "supervised"]], rag_db, 1)
+
+    result = await agent.run("question", deps=Deps())
+
+    assert outcomes(result) == ["ok", "failed", "failed", "failed"]
+
+
+@pytest.mark.asyncio
+async def test_a_sibling_past_the_allowance_pays_its_own_unit(rag_db):
+    burst = [["ai", "machine learning", "deep learning", "supervised learning"]]
+
+    within = await burst_agent(burst, rag_db, 2).run("question", deps=Deps())
+    over = await burst_agent(burst, rag_db, 1).run("question", deps=Deps())
+
+    assert outcomes(within) == ["ok", "ok", "ok", "ok"]
+    assert outcomes(over) == ["ok", "ok", "ok", "failed"]
+
+
+@pytest.mark.asyncio
+async def test_unit_tracking_resets_between_runs(rag_db):
+    """A second run's opening burst prices like a first run's."""
+
+    def model_function(messages, _info) -> ModelResponse:
+        if any(isinstance(part, ToolReturnPart) for part in messages[-1].parts):
+            return ModelResponse(parts=[TextPart("done")])
+        return ModelResponse(
+            parts=[
+                ToolCallPart("rag_search", {"query": query})
+                for query in ["ai", "machine learning", "deep learning"]
+            ]
+        )
+
+    config = AppConfig()
+    config.qa.max_searches = 1
+    agent = Agent(
+        FunctionModel(model_function),
+        deps_type=Deps,
+        capabilities=[create_rag(db_path=rag_db, config=config, defer_loading=False)],
+    )
+    deps = Deps()
+
+    first = await agent.run("question", deps=deps)
+    second = await agent.run("another", deps=deps, message_history=first.all_messages())
+
+    assert outcomes(first) == ["ok", "ok", "ok"]
+    assert outcomes(second)[-3:] == ["ok", "ok", "ok"]

tests/multi_db/test_capabilities.py

@@ -30,7 +30,7 @@ class TestAskAcrossDatabases:
             capability = create_capability(config=config, rag=rag, defer_loading=False)
             capability.state = RAGState(sources=["alpha"])
 
-            formatted = await capability._search("cats", limit=10)
+            formatted = await capability._search("cats", 10, 1)
 
         assert isinstance(formatted, str)
         assert "alpha" in formatted
@@ -47,7 +47,7 @@ class TestAskAcrossDatabases:
             capability = create_capability(config=config, rag=rag, defer_loading=False)
             capability.state = RAGState()
 
-            formatted = await capability._search("cats", limit=10)
+            formatted = await capability._search("cats", 10, 1)
 
         assert isinstance(formatted, str)
         assert "alpha document" in formatted
@@ -72,7 +72,7 @@ class TestStandaloneCapabilities:
         assert capability.scope.names == ("alpha", "beta")
         run = await capability.for_run(make_context(Deps()))
         try:
-            formatted = await run._search("cats", limit=10)
+            formatted = await run._search("cats", 10, 1)
         finally:
             await run._close()
 
@@ -135,7 +135,7 @@ class TestAnalyzeAcrossDatabases:
             capability = create_analysis(config=config, rag=rag, defer_loading=False)
             capability.state = AnalysisState(sources=["alpha"])
 
-            formatted = await capability._search("cats", limit=10)
+            formatted = await capability._search("cats", 10, 1)
             sandbox = await capability._ensure_sandbox()
             await capability._close()
 

tests/multi_db/test_citations.py

@@ -342,7 +342,7 @@ class TestCiteFallback:
             )
             run = await capability.for_run(make_context(deps))
             # The search returns the cats chunk, never the aardvark one.
-            await run._search("cats", limit=10)
+            await run._search("cats", 10, 1)
 
             await run._cite([aardvark.id])
 
@@ -374,7 +374,7 @@ class TestCiteFallback:
                 state={"rag": RAGState(sources=["alpha"]).model_dump(mode="json")}
             )
             run = await capability.for_run(make_context(deps))
-            await run._search("cats", limit=10)
+            await run._search("cats", 10, 1)
 
             with pytest.raises(ModelRetry, match="None of the supplied chunk_ids"):
                 await run._cite([outside.id])

tests/test_picture_in_context.py

@@ -1022,7 +1022,7 @@ async def test_rag_capability_attaches_images_for_vision_model(temp_db_path):
     capability.state = RAGState()
     capability.rag = fake_client
 
-    result = await capability._search("anything", None)
+    result = await capability._search("anything", None, 1)
 
     assert isinstance(result, ToolReturn)
     assert result.content is not None