Trim earlier-question evidence off the wire, keeping the images

_compact_old_tool_returns ran in before_model_request, whose result core
assigns back onto ctx.state.message_history, so the trim reached
all_messages() and every host that persists a thread. It now runs in
wrap_model_request, which operates on a detached list: the model sees the
trimmed history, the host keeps what it retrieved.

Content attached to a ToolReturn arrives as its own UserPromptPart in the
same ModelRequest as the ToolReturnPart, so the turn-boundary scan read an
image-bearing search result as a new user turn and discarded evidence
retrieved earlier in the same turn. _is_user_turn() now requires a request
with no tool returns.

Page images on a replaced return stay. Dropping them with their text bounds
context growth, but a follow-up about a figure already shown ("what colour
is that box?") carries no terms that could retrieve it again: measured on
gemma4-26b against two ORB figures, removing the image turned both answers
into "I cannot find enough information", and keeping it answers correctly.
Bounding that growth needs to preserve cited figures, which is a separate
change.

The replacement notice no longer claims citations remain in state;
_clear_invocation_state has cleared them by then.

CHANGELOG.md

@@ -7,6 +7,8 @@
 - `haiku-rag` and `haiku-ingester` CLI startup no longer imports `lancedb`, `pyarrow` and `pydantic_ai`.
 - `haiku.rag.store` no longer re-exports `Store`; import it from `haiku.rag.store.engine`.
 - `docling-local` reuses one docling `DocumentConverter` per set of conversion options instead of building one per document, so local layout, table and OCR models are no longer loaded per document. Conversions through a shared converter are serialized.
+- Prior-question tool output is trimmed from the model request only; `all_messages()` retains what the run gathered.
+- A search result carrying page images no longer counts as a turn boundary; evidence retrieved earlier in the same turn survives.
 
 ## [0.73.0] - 2026-08-06
 

haiku_rag_slim/haiku/rag/capabilities/_base.py

@@ -7,7 +7,7 @@ from typing import Any, cast
 
 from pydantic import BaseModel
 from pydantic_ai import ModelRetry, RunContext, ToolFailed
-from pydantic_ai.capabilities import AbstractCapability
+from pydantic_ai.capabilities import AbstractCapability, WrapModelRequestHandler
 from pydantic_ai.messages import (
     InstructionPart,
     ModelMessage,
@@ -76,20 +76,43 @@ def _clear_invocation_state(state: BaseModel) -> None:
             value.clear()
 
 
+PRIOR_TURN_NOTICE = (
+    "[Evidence retrieved for an earlier question, no longer shown. It does not "
+    "count as cited for the current question.]"
+)
+
+
+def _is_user_turn(message: ModelMessage) -> bool:
+    """Whether this message is a user turn rather than tool output.
+
+    Content attached to a ``ToolReturn`` (page images) arrives as its own
+    ``UserPromptPart`` in the same ``ModelRequest`` as the ``ToolReturnPart``,
+    so a bare ``UserPromptPart`` check reads tool output as a new turn.
+    """
+    if not isinstance(message, ModelRequest):
+        return False
+    return any(isinstance(part, UserPromptPart) for part in message.parts) and not any(
+        isinstance(part, ToolReturnPart) for part in message.parts
+    )
+
+
 def _compact_old_tool_returns(
     messages: list[ModelMessage], tool_names: frozenset[str]
 ) -> list[ModelMessage]:
-    """Remove bulky prior-turn evidence while retaining the current turn.
+    """Remove bulky earlier-question evidence while retaining the current one.
 
     Tool call and return parts remain paired; only the old return payload is
-    replaced. This keeps provider histories valid and preserves all evidence
-    gathered since the most recent user prompt.
+    replaced.
+
+    Page images attached to a replaced return are deliberately left in place.
+    Dropping them alongside their text is tempting — they are the bulk, and
+    they accumulate — but a follow-up about a figure already shown ("what
+    colour is that box?") carries no terms that could retrieve it again, so
+    removing the image turns an answerable question into a refusal.
     """
     latest_user_message = -1
     for index, message in enumerate(messages):
-        if isinstance(message, ModelRequest) and any(
-            isinstance(part, UserPromptPart) for part in message.parts
-        ):
+        if _is_user_turn(message):
             latest_user_message = index
 
     if latest_user_message < 0:
@@ -100,10 +123,7 @@ def _compact_old_tool_returns(
         if not isinstance(message, ModelRequest):
             continue
         parts = [
-            replace(
-                part,
-                content="[Prior-turn RAG tool output removed; citations remain in state.]",
-            )
+            replace(part, content=PRIOR_TURN_NOTICE)
             if isinstance(part, ToolReturnPart) and part.tool_name in tool_names
             else part
             for part in message.parts
@@ -168,12 +188,27 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
             return f"{self.config.prompts.domain_preamble}\n\n{self.instruction_text}"
         return self.instruction_text
 
-    async def before_model_request(
-        self, ctx: RunContext[Any], request_context: ModelRequestContext
-    ) -> ModelRequestContext:
+    async def wrap_model_request(
+        self,
+        ctx: RunContext[Any],
+        *,
+        request_context: ModelRequestContext,
+        handler: WrapModelRequestHandler,
+    ) -> ModelResponse:
+        """Trim earlier-question evidence off the wire only.
+
+        Deliberately not ``before_model_request``: that hook's result is
+        assigned back onto the run's message history, so trimming there would
+        destroy the host's record of what was retrieved.
+        """
         request_context.messages = _compact_old_tool_returns(
             request_context.messages, self.tool_names
         )
+        return await handler(request_context)
+
+    async def before_model_request(
+        self, ctx: RunContext[Any], request_context: ModelRequestContext
+    ) -> ModelRequestContext:
         if instruction := self._budget_notice():
             current_request = request_context.messages[-1]
             if isinstance(current_request, ModelRequest):

tests/capabilities/test_capabilities.py

@@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, patch
 import pytest
 from pydantic_ai import Agent, ModelRetry, RunContext, ToolFailed
 from pydantic_ai.messages import (
+    BinaryContent,
     ModelRequest,
     ModelResponse,
     TextPart,
@@ -19,6 +20,7 @@ from pydantic_ai.usage import RunUsage
 
 from haiku.rag.capabilities._base import (
     CITATION_GRACE_REQUESTS,
+    PRIOR_TURN_NOTICE,
     _called_own_tool,
     _compact_old_tool_returns,
 )
@@ -806,7 +808,7 @@ def test_prior_turn_tool_results_are_compacted_but_current_evidence_is_kept():
     old_return = compacted[2].parts[0]
     current_return = compacted[5].parts[0]
     assert isinstance(old_return, ToolReturnPart)
-    assert "removed" in str(old_return.content)
+    assert old_return.content == PRIOR_TURN_NOTICE
     assert isinstance(current_return, ToolReturnPart)
     assert current_return.content == "current evidence"
 
@@ -825,3 +827,116 @@ def test_tool_results_are_unchanged_when_history_has_no_user_prompt():
     current_return = compacted[1].parts[0]
     assert isinstance(current_return, ToolReturnPart)
     assert current_return.content == "current evidence"
+
+
+PAGE_IMAGE = BinaryContent(data=b"\x89PNG" + b"\x00" * 64, media_type="image/png")
+
+
+def _search_exchange(call_id: str, evidence: str, *, images: bool):
+    """One search round-trip in the shape pydantic-ai produces.
+
+    Images on a ``ToolReturn`` arrive as a separate ``UserPromptPart`` appended
+    to the same ``ModelRequest`` as the ``ToolReturnPart``.
+    """
+    parts: list[Any] = [ToolReturnPart("rag_search", evidence, call_id)]
+    if images:
+        parts.append(UserPromptPart(content=[PAGE_IMAGE]))
+    return [
+        ModelResponse(parts=[ToolCallPart("rag_search", {}, call_id)]),
+        ModelRequest(parts=parts),
+    ]
+
+
+def test_image_bearing_tool_return_does_not_end_the_current_turn():
+    messages = [
+        ModelRequest(parts=[UserPromptPart("current question")]),
+        *_search_exchange("first-call", "first evidence", images=False),
+        *_search_exchange("second-call", "second evidence", images=True),
+    ]
+
+    compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"}))
+
+    first_return = compacted[2].parts[0]
+    assert isinstance(first_return, ToolReturnPart)
+    assert first_return.content == "first evidence"
+
+
+def test_prior_turn_images_outlive_their_tool_return():
+    """A follow-up about a figure cannot retrieve it again, so keep the image.
+
+    "What colour is that box?" has no terms the search can use, so dropping the
+    image with its text turns an answerable question into a refusal.
+    """
+    messages = [
+        ModelRequest(parts=[UserPromptPart("old question")]),
+        *_search_exchange("old-call", "old evidence", images=True),
+        ModelResponse(parts=[TextPart("an answer")]),
+        ModelRequest(parts=[UserPromptPart("current question")]),
+    ]
+
+    compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"}))
+
+    old_return = compacted[2].parts[0]
+    assert isinstance(old_return, ToolReturnPart)
+    assert old_return.content == PRIOR_TURN_NOTICE
+    assert any(
+        isinstance(part, UserPromptPart) and not isinstance(part.content, str)
+        for message in compacted
+        for part in message.parts
+    )
+
+
+def test_user_attached_image_starts_a_turn_and_is_never_dropped():
+    messages = [
+        ModelRequest(parts=[UserPromptPart("old question")]),
+        *_search_exchange("old-call", "old evidence", images=False),
+        ModelResponse(parts=[TextPart("an answer")]),
+        ModelRequest(parts=[UserPromptPart(content=[PAGE_IMAGE, "what is this?"])]),
+    ]
+
+    compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"}))
+
+    old_return = compacted[2].parts[0]
+    assert isinstance(old_return, ToolReturnPart)
+    assert old_return.content == PRIOR_TURN_NOTICE
+    attached = compacted[-1].parts[0]
+    assert isinstance(attached, UserPromptPart)
+    assert attached.content == [PAGE_IMAGE, "what is this?"]
+
+
+async def test_compaction_never_reaches_the_stored_message_history(temp_db_path):
+    """Trimming is for the wire; hosts keep the evidence they gathered."""
+    capability = create_rag(
+        db_path=temp_db_path, config=AppConfig(), defer_loading=False
+    )
+    turns = iter(
+        [
+            ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
+            ModelResponse(parts=[TextPart("first answer")]),
+            ModelResponse(parts=[TextPart("second answer")]),
+        ]
+    )
+
+    async def model(_messages, _info):
+        return next(turns)
+
+    agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability])
+    deps = Deps(state={"rag": RAGState().model_dump(mode="json")})
+
+    with patch.object(
+        RAGCapability, "_search", AsyncMock(return_value="REAL EVIDENCE")
+    ):
+        first = await agent.run("old question", deps=deps)
+        second = await agent.run(
+            "current question", deps=deps, message_history=first.all_messages()
+        )
+
+    returns = [
+        str(part.content)
+        for message in second.all_messages()
+        if isinstance(message, ModelRequest)
+        for part in message.parts
+        if isinstance(part, ToolReturnPart)
+    ]
+    assert "REAL EVIDENCE" in returns
+    assert PRIOR_TURN_NOTICE not in returns