Merge pull request #597 from ggozad/feat/search-fanout-units

Price search budget in units and deduplicate fan-out results

CHANGELOG.md

@@ -2,6 +2,15 @@
 
 ## [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.
+- Searches in one model response deduplicate their results: evidence a sibling
+  search already showed collapses to a reference line, and a picture attaches
+  once per response.
+
 ## [0.81.0] - 2026-09-01
 
 ### Added

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: 5       # Maximum search tool calls per question
+  max_searches: 5       # Maximum search units 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 a capability can make per question (default: 5). Shared by the RAG and analysis capabilities.
+- **max_searches**: Maximum number of search units a capability can spend per question (default: 5). Up to three searches emitted in the same model response share one unit, so a model that rephrases its query in one response spends one unit. A search in a later response starts a new unit, as does each further group of three within one response. Shared by the RAG and analysis capabilities. Searches in one response also deduplicate their returns: evidence a sibling search already showed collapses to a reference line, and each picture attaches once per response.
 
 !!! 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.

haiku_rag_slim/haiku/rag/capabilities/_base.py

@@ -14,6 +14,7 @@ from pydantic_ai import (
 )
 from pydantic_ai.capabilities import AbstractCapability
 from pydantic_ai.messages import (
+    BinaryContent,
     InstructionPart,
     ModelMessage,
     ModelRequest,
@@ -29,6 +30,7 @@ from pydantic_ai.toolsets import AgentToolset
 
 from haiku.rag.capabilities._tools import (
     CodeExecutionEntry,
+    EvidenceKey,
     merge_results,
     search_corpus,
 )
@@ -43,7 +45,7 @@ from haiku.rag.store.models.citation import (
     ambiguous_citation,
     resolve_citations,
 )
-from haiku.rag.tools.search import build_image_content_from_results
+from haiku.rag.tools.search import PictureKey, build_image_content_from_results
 
 CITATION_GRACE_REQUESTS = 2
 """Requests calling this capability's tools that its cite tool outlives the rest by.
@@ -61,6 +63,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 +188,12 @@ 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)
+    step_shown: set[EvidenceKey] = field(default_factory=set, repr=False)
+    step_pictures: set[PictureKey] = field(default_factory=set, 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 +243,11 @@ 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,
+            step_shown=set(),
+            step_pictures=set(),
             request_count=0,
             grace_requests_used=0,
             epoch=0,
@@ -493,32 +515,54 @@ 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 run_step != self.search_step:
-        if self.search_count > self._max_searches:
+            self.search_step = run_step
+            self.step_searches = 0
+            self.step_rejected = False
+            self.step_shown = set()
+            self.step_pictures = set()
+        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."
             )
         async with self.rag_lock:
-            formatted, results, include_collection = await search_corpus(
+            formatted, results, rendered, include_collection = await search_corpus(
                 await self._ensure_rag(),
                 query,
                 limit=limit,
                 document_filter=self.state.document_filter,
                 sources=self.state.sources,
+                shown=self.step_shown,
             )
+        parts: list[str | BinaryContent] = []
+        emitted: set[PictureKey] = set()
+        if self.vision:
+            parts, emitted = build_image_content_from_results(
+                results,
+                include_collection=include_collection,
+                exclude=self.step_pictures,
+            )
+        # Everything the search produced commits together, after formatting and
+        # image construction have both succeeded: a search that raises must not
+        # leave results citable, note evidence the model never received, or
+        # suppress a later sibling's results.
         state = self.state
         # A model can search the same query twice with different limits, and the
         # narrower return must not drop what the wider one already showed it.
         merge_results(state.searches.setdefault(query, []), results)
         self._note_evidence()
-        if self.vision and (
+        self.step_shown |= rendered
-            parts := build_image_content_from_results(
+        self.step_pictures |= emitted
-                results, include_collection=include_collection
+        if parts:
-            )
-        ):
             return ToolReturn(return_value=formatted, content=parts)
         return formatted
 

haiku_rag_slim/haiku/rag/capabilities/_tools.py

@@ -1,9 +1,11 @@
 from collections.abc import Iterable
+from collections.abc import Set as AbstractSet
 
 from pydantic import BaseModel
 
 from haiku.rag.client import HaikuRAG
 from haiku.rag.store.models.chunk import SearchResult, qualified_id
+from haiku.rag.tools.search import picture_keys
 
 
 class CodeExecutionEntry(BaseModel):
@@ -13,14 +15,48 @@ class CodeExecutionEntry(BaseModel):
     success: bool = True
 
 
+EvidenceKey = tuple[tuple[str | None, str | None], tuple[str, frozenset]]
+"""What tells one rendered result from another: qualified id, then signature.
+
+The qualified id comes first because the rendered string alone would conflate
+identical renderings of the same chunk id held by two databases.
+"""
+
+
+def evidence_signature(result: SearchResult, include_collection: bool) -> tuple:
+    """The rendered evidence a result shows the model, as an equivalence key.
+
+    Rank and total are held at neutral values: they vary with a result's
+    position, and position (like score) must not tell two renderings apart.
+    """
+    return (
+        result.format_for_agent(rank=0, total=0, include_collection=include_collection),
+        picture_keys(result),
+    )
+
+
+def evidence_key(result: SearchResult, include_collection: bool) -> EvidenceKey:
+    return (
+        qualified_id(result.source, result.chunk_id),
+        evidence_signature(result, include_collection),
+    )
+
+
 async def search_corpus(
     rag: HaikuRAG,
     query: str,
     limit: int | None = None,
     document_filter: str | None = None,
     sources: list[str] | None = None,
-) -> tuple[str, list[SearchResult], bool]:
+    shown: AbstractSet[EvidenceKey] = frozenset(),
-    """Search and context-expand results, and whether they name their collection."""
+) -> tuple[str, list[SearchResult], set[EvidenceKey], bool]:
+    """Search and context-expand results, eliding evidence already shown.
+
+    Returns the formatted results, the full result list, the evidence keys the
+    formatting rendered in full, and whether results name their collection. A
+    result whose key is in ``shown`` keeps its slot but collapses to one line;
+    the result list is never filtered.
+    """
     results = await rag.search(
         query, limit=limit, filter=document_filter, sources=sources
     )
@@ -29,13 +65,25 @@ async def search_corpus(
     # two collections names them even when everything came back from one.
     selected = rag.source_names if sources is None else sources
     include_collection = len(set(selected)) > 1
-    formatted = "\n\n---\n\n".join(
+    rendered: set[EvidenceKey] = set()
-        result.format_for_agent(
+    parts: list[str] = []
-            rank=index + 1, total=len(results), include_collection=include_collection
+    total = len(results)
-        )
+    for index, result in enumerate(results):
-        for index, result in enumerate(results)
+        key = evidence_key(result, include_collection)
-    )
+        if key in shown or key in rendered:
-    return formatted or "No results found.", list(results), include_collection
+            parts.append(
+                f"Also matched, shown above: [{result.chunk_id}] "
+                f"[rank {index + 1} of {total}]"
+            )
+        else:
+            parts.append(
+                result.format_for_agent(
+                    rank=index + 1, total=total, include_collection=include_collection
+                )
+            )
+            rendered.add(key)
+    formatted = "\n\n---\n\n".join(parts)
+    return formatted or "No results found.", list(results), rendered, include_collection
 
 
 def merge_results(
@@ -56,6 +104,9 @@ def merge_results(
 
 __all__ = [
     "CodeExecutionEntry",
+    "EvidenceKey",
+    "evidence_key",
+    "evidence_signature",
     "merge_results",
     "search_corpus",
 ]

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

haiku_rag_slim/haiku/rag/tools/search.py

@@ -1,5 +1,6 @@
 import base64
 from collections.abc import Callable
+from collections.abc import Set as AbstractSet
 from io import BytesIO
 
 from PIL import Image
@@ -20,6 +21,22 @@ their own picture had it removed, along with their text.
 """
 
 
+PictureKey = tuple[str | None, str | None, str]
+"""Identity of one attached picture: (source, document_id, self_ref).
+
+``self_ref`` alone collides across documents, and a copy of a document in
+another collection carries its own pictures.
+"""
+
+
+def picture_keys(result: SearchResult) -> frozenset[PictureKey]:
+    """The identity of every picture this result carries."""
+    return frozenset(
+        (result.source, result.document_id, self_ref)
+        for self_ref in (result.image_data or {})
+    )
+
+
 def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None:
     """Wrap picture bytes for the wire, or return nothing if they will not decode.
 
@@ -38,11 +55,14 @@ def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None:
 def build_image_content_from_results(
     results: list[SearchResult],
     include_collection: bool = False,
-) -> list[str | BinaryContent]:
+    exclude: AbstractSet[PictureKey] = frozenset(),
+) -> tuple[list[str | BinaryContent], set[PictureKey]]:
     """Decode and validate picture bytes attached to search results, labelled.
 
-    Dedup keyed on ``(source, document_id, self_ref)`` so the same picture in
+    Returns the labelled content and the ``PictureKey`` of every picture it
-    different chunks is sent once, and a copy in another collection is its own. Pictures that fail
+    emitted. Dedup keyed on ``PictureKey`` so the same picture in
+    different chunks is sent once, and a copy in another collection is its
+    own; ``exclude`` seeds that dedup with pictures already sent. Pictures that fail
     ``PIL.Image.verify()`` are skipped — the model adapter renders one
     vision placeholder per ``BinaryContent``, so emitting one for an
     image the server can't decode leaves the processor with an
@@ -59,7 +79,8 @@ def build_image_content_from_results(
     to the vision API.
     """
     collected: list[tuple[str | None, str | None, str, BinaryContent]] = []
-    seen: set[tuple[str | None, str | None, str]] = set()
+    seen: set[PictureKey] = set(exclude)
+    emitted: set[PictureKey] = set()
     for result in results:
         if not result.image_data:
             continue
@@ -72,6 +93,7 @@ def build_image_content_from_results(
                 continue
             collected.append((result.source, result.chunk_id, self_ref, picture))
             seen.add(key)
+            emitted.add(key)
 
     content: list[str | BinaryContent] = []
     total = len(collected)
@@ -83,7 +105,7 @@ def build_image_content_from_results(
             f"Not provided by the user. {RETRIEVED_IMAGE_TAG}"
         )
         content.append(picture)
-    return content
+    return content, emitted
 
 
 def create_search_toolset(
@@ -174,7 +196,7 @@ def create_search_toolset(
         if not config.qa.model.vision:
             return text
 
-        image_content = build_image_content_from_results(
+        image_content, _ = build_image_content_from_results(
             results_list, include_collection=include_collection
         )
         if image_content:

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", 20, 1)
-    await capability._search("Figure 3-1", None)
+    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", 20, 1)
-    await capability._search("cats", None)
+    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(
@@ -515,6 +517,89 @@ async def test_a_picture_that_will_not_decode_emits_neither_image_nor_label(
     assert texts_of(wire[-1]) == []
 
 
+def _burst_result() -> SearchResult:
+    return SearchResult(
+        content="evidence",
+        score=1.0,
+        chunk_id="chunk-1",
+        document_id="doc-1",
+        source="main",
+        doc_item_refs=["#/pictures/0"],
+        image_data={"#/pictures/0": base64.b64encode(REAL_PNG).decode()},
+    )
+
+
+async def _fanout_question_then_another(temp_db_path, cite: bool) -> list[list[Any]]:
+    """Question 1 fans out over one picture chunk; question 2 follows compacted."""
+    rag = create_rag(
+        db_path=temp_db_path, config=AppConfig(), defer_loading=False, vision=True
+    )
+    client = AsyncMock()
+    client.search.side_effect = [[_burst_result()], [_burst_result()]]
+    client.expand_context.side_effect = lambda results: results
+    client.source_names = ["main"]
+    rag.borrowed_rag = client
+    citing = (
+        [[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-3")]]
+        if cite
+        else []
+    )
+    calls = iter(
+        [
+            [
+                ToolCallPart("rag_search", {"query": "figure"}, "call-1"),
+                ToolCallPart("rag_search", {"query": "the figure"}, "call-2"),
+            ],
+            *citing,
+            [TextPart("first answer")],
+            [TextPart("second answer")],
+        ]
+    )
+    wire: list[list[Any]] = []
+
+    async def model(messages, _info):
+        wire.append(list(messages))
+        return ModelResponse(parts=next(calls))
+
+    agent = Agent(
+        FunctionModel(model),
+        deps_type=Deps,
+        capabilities=[rag, create_compaction()],
+    )
+    deps = Deps()
+
+    with patch.object(
+        RAGCapability, "get_picture_bytes", AsyncMock(return_value=REAL_PNG)
+    ):
+        first = await agent.run("what does the figure show?", deps=deps)
+        await agent.run(
+            "and what else?", deps=deps, message_history=first.all_messages()
+        )
+    return wire
+
+
+@pytest.mark.asyncio
+async def test_a_burst_deduplicated_picture_survives_compaction_when_cited(
+    temp_db_path,
+):
+    """Dedup attaches the picture once in its own question; the capsule re-fetches
+    it for the next. Neither pass may leave the model without it."""
+    wire = await _fanout_question_then_another(temp_db_path, cite=True)
+
+    assert [picture.data for picture in images_of(wire[1])] == [REAL_PNG]
+    assert [picture.data for picture in images_of(wire[-1])] == [REAL_PNG]
+
+
+@pytest.mark.asyncio
+async def test_a_burst_deduplicated_picture_is_dropped_by_compaction_uncited(
+    temp_db_path,
+):
+    wire = await _fanout_question_then_another(temp_db_path, cite=False)
+
+    assert [picture.data for picture in images_of(wire[1])] == [REAL_PNG]
+    assert images_of(wire[-1]) == []
+
+
 @pytest.mark.asyncio
 async def test_the_capsule_is_built_once_per_request_and_again_for_the_next(
     temp_db_path,

tests/capabilities/test_search_units.py

@@ -0,0 +1,329 @@
+import base64
+from dataclasses import dataclass, field
+from io import BytesIO
+from typing import Any
+from unittest.mock import AsyncMock
+
+import pytest
+from PIL import Image as PILImage
+from pydantic_ai import Agent
+from pydantic_ai.messages import (
+    BinaryContent,
+    ModelResponse,
+    TextPart,
+    ToolCallPart,
+    ToolReturnPart,
+)
+from pydantic_ai.models.function import FunctionModel
+from pydantic_ai.run import AgentRunResult
+
+from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
+from haiku.rag.capabilities.rag import RAGState
+from haiku.rag.capabilities.rag import create_capability as create_rag
+from haiku.rag.config.models import AppConfig
+from haiku.rag.store.models.chunk import SearchResult
+
+
+@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"]
+
+
+def _png() -> str:
+    buffer = BytesIO()
+    PILImage.new("RGB", (4, 4), "red").save(buffer, format="PNG")
+    return base64.b64encode(buffer.getvalue()).decode()
+
+
+def make_result(**overrides: Any) -> SearchResult:
+    fields: dict[str, Any] = {
+        "content": "body",
+        "score": 0.9,
+        "source": "main",
+        "chunk_id": "c1",
+        "document_id": "d1",
+        "image_data": {"#/pictures/0": _png()},
+    }
+    fields.update(overrides)
+    return SearchResult(**fields)
+
+
+def stub_client(
+    *batches: list[SearchResult], sources: list[str] | None = None
+) -> AsyncMock:
+    client = AsyncMock()
+    client.search.side_effect = list(batches)
+    client.expand_context.side_effect = lambda results: results
+    client.source_names = sources or ["main"]
+    return client
+
+
+def dedup_capability(client: AsyncMock, temp_db_path, *, vision: bool = True):
+    capability = create_rag(db_path=temp_db_path, config=AppConfig(), vision=vision)
+    capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0))
+    capability.borrowed_rag = client
+    return capability
+
+
+def images_of(returned: Any) -> list[BinaryContent]:
+    if isinstance(returned, str):
+        return []
+    return [item for item in returned.content if isinstance(item, BinaryContent)]
+
+
+def text_of(returned: Any) -> str:
+    return returned if isinstance(returned, str) else returned.return_value
+
+
+@pytest.mark.asyncio
+async def test_a_duplicate_sibling_is_elided_and_stays_citable(temp_db_path):
+    duplicate, novel = make_result(), make_result(chunk_id="c2", content="novel")
+    client = stub_client([make_result()], [duplicate, novel])
+    capability = dedup_capability(client, temp_db_path)
+
+    first = await capability._search("q", None, 1)
+    second = await capability._search("q rephrased", None, 1)
+
+    assert len(images_of(first)) == 1
+    assert images_of(second) == []
+    text = text_of(second)
+    assert "Also matched, shown above: [c1] [rank 1 of 2]" in text
+    assert "body" not in text
+    assert "[rank 2 of 2]" in text and "novel" in text
+    assert [r.chunk_id for r in capability.state.searches["q rephrased"]] == [
+        "c1",
+        "c2",
+    ]
+    assert await capability._cite(["c1"]) == "Registered 1 citation(s)."
+
+
+@pytest.mark.asyncio
+async def test_a_new_run_step_formats_shown_results_in_full(temp_db_path):
+    client = stub_client([make_result()], [make_result()])
+    capability = dedup_capability(client, temp_db_path)
+
+    await capability._search("q", None, 1)
+    second = await capability._search("q again", None, 2)
+
+    assert "body" in text_of(second)
+    assert len(images_of(second)) == 1
+
+
+@pytest.mark.asyncio
+async def test_same_chunk_id_from_another_collection_is_not_elided(temp_db_path):
+    client = stub_client(
+        [make_result(source="alpha")],
+        [make_result(source="beta")],
+        sources=["alpha", "beta"],
+    )
+    capability = dedup_capability(client, temp_db_path)
+
+    await capability._search("q", None, 1)
+    second = await capability._search("q rephrased", None, 1)
+
+    assert "body" in text_of(second)
+
+
+@pytest.mark.asyncio
+async def test_same_anchor_with_new_evidence_formats_in_full(temp_db_path):
+    shared, extra = _png(), _png()
+    client = stub_client(
+        [make_result(content="c1 with c2", image_data={"#/pictures/1": shared})],
+        [
+            make_result(
+                content="c1 with c3",
+                image_data={"#/pictures/1": shared, "#/pictures/3": extra},
+            )
+        ],
+    )
+    capability = dedup_capability(client, temp_db_path)
+
+    await capability._search("q", None, 1)
+    second = await capability._search("q rephrased", None, 1)
+
+    assert "c1 with c3" in text_of(second)
+    assert len(images_of(second)) == 1
+    labels = [item for item in second.content if isinstance(item, str)]
+    assert any("#/pictures/3" in label for label in labels)
+
+
+@pytest.mark.parametrize(
+    ("overrides", "elided"),
+    [
+        ({"score": 0.1}, True),
+        ({"content": "different"}, False),
+        ({"document_title": "Other"}, False),
+        ({"headings": ["Heading"]}, False),
+        ({"labels": ["table"]}, False),
+        ({"picture_captions": {"#/pictures/0": "A caption"}}, False),
+        ({"image_data": {"#/pictures/9": _png()}}, False),
+    ],
+)
+@pytest.mark.asyncio
+async def test_equivalence_follows_the_rendered_evidence(
+    temp_db_path, overrides: dict[str, Any], elided: bool
+):
+    """Any rendered field or picture identity defeats elision; score alone does not."""
+    client = stub_client([make_result()], [make_result(**overrides)])
+    capability = dedup_capability(client, temp_db_path)
+
+    await capability._search("q", None, 1)
+    second = await capability._search("q rephrased", None, 1)
+
+    assert ("Also matched, shown above" in text_of(second)) is elided
+
+
+@pytest.mark.asyncio
+async def test_a_failed_sibling_commits_nothing(temp_db_path):
+    client = stub_client(
+        [make_result(image_data={"#/pictures/0": "AAA"})],
+        [make_result()],
+    )
+    capability = dedup_capability(client, temp_db_path)
+    evidence_before = capability.state.evidence.model_dump()
+
+    with pytest.raises(Exception):
+        await capability._search("q", None, 1)
+
+    assert capability.state.searches == {}
+    assert capability.state.evidence.model_dump() == evidence_before
+
+    second = await capability._search("q rephrased", None, 1)
+
+    assert "body" in text_of(second)
+    assert len(images_of(second)) == 1

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()
 
@@ -305,8 +305,10 @@ class TestWhenTheModelIsToldTheCollection:
         async with HaikuRAG(config=config) as rag:
             monkeypatch.setattr(rag, "search", AsyncMock(return_value=only_alpha))
 
-            spanning, _, spans = await search_corpus(rag, "cats")
+            spanning, _, _, spans = await search_corpus(rag, "cats")
-            narrowed, _, narrows = await search_corpus(rag, "cats", sources=["alpha"])
+            narrowed, _, _, narrows = await search_corpus(
+                rag, "cats", sources=["alpha"]
+            )
 
         assert "Collection: alpha" in spanning
         assert "Collection" not in narrowed

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

tests/tools/test_search.py

@@ -311,7 +311,7 @@ class TestBuildImageContentFromResults:
             SearchResult(content="text only", score=0.5, chunk_id="c1", image_data=None)
         ]
 
-        assert build_image_content_from_results(results) == []
+        assert build_image_content_from_results(results) == ([], set())
 
     def test_duplicate_document_and_ref_is_attached_once(self):
         from pydantic_ai.messages import BinaryContent
@@ -336,7 +336,7 @@ class TestBuildImageContentFromResults:
             ),
         ]
 
-        content = build_image_content_from_results(results)
+        content, _ = build_image_content_from_results(results)
 
         images = [item for item in content if isinstance(item, BinaryContent)]
         assert len(images) == 1
@@ -369,7 +369,7 @@ class TestBuildImageContentFromResults:
 
         from haiku.rag.tools.search import build_image_content_from_results
 
-        content = build_image_content_from_results(
+        content, _ = build_image_content_from_results(
             self._one_picture_in_two_collections()
         )
 
@@ -381,7 +381,7 @@ class TestBuildImageContentFromResults:
         reference."""
         from haiku.rag.tools.search import build_image_content_from_results
 
-        content = build_image_content_from_results(
+        content, _ = build_image_content_from_results(
             self._one_picture_in_two_collections(), include_collection=True
         )
 
@@ -392,7 +392,7 @@ class TestBuildImageContentFromResults:
     def test_an_unasked_for_collection_is_not_named_on_an_image(self):
         from haiku.rag.tools.search import build_image_content_from_results
 
-        content = build_image_content_from_results(
+        content, _ = build_image_content_from_results(
             self._one_picture_in_two_collections()
         )
 
@@ -431,7 +431,7 @@ class TestBuildImageContentFromResults:
             ),
         ]
 
-        content = build_image_content_from_results(results)
+        content, _ = build_image_content_from_results(results)
 
         # label, image, label, image — each picture preceded by its own line.
         assert [type(item) is str for item in content] == [True, False, True, False]