Let the record say whether a question is still being answered

Resumption was inferred from the transcript, and no shape says it. A missing
prompt is how UI adapters ask their first question as much as how pydantic-ai
resumes one, so every AG-UI host failed on its first message. Reading the tail
instead moved the error rather than fixing it: a settled structured answer ends
with an output tool's return, indistinguishable from results delivered to a
question still in progress, so following questions inherited the first one's
identity.

`CapabilityEvidenceRecord.in_progress` is now the authority. `begin_question`
sets it, `after_run` clears it unless the run is only pausing for deferred work,
and a run that raised never reaches `after_run`, which is what leaves an
interrupted question claimable. The history is consulted only to catch a host
that dropped the state of a question the model is unmistakably still owed.

CHANGELOG.md

@@ -6,7 +6,7 @@
 - `CitationPolicyCapability` (`haiku.rag.capabilities.policy.create_capability`): registering it requires every answer to declare its grounding, in any conversation that has something to declare — this question retrieved evidence, or something was cited earlier. A question that ends undeclared is sent back once to record what grounded the answer already given, and is recorded in `CitationPolicyState.violations` if it finishes undeclared regardless. A conversation with neither a current-question evidence outcome nor any earlier citation is not enforced.
 - `haiku.rag.capabilities.evidence.discover_evidence()` and `DiscoveredEvidence`, moved out of `compaction` so both optional capabilities share them. `RAGCapabilityBase.cite_available`.
 - `EvidenceCompactionCapability` (`haiku.rag.capabilities.compaction.create_capability`): registering it replaces earlier questions' evidence on the model request with the evidence that was cited, grouped by the question that cited it, cited page images re-attached, other earlier evidence returns reduced to a receipt. Requests only; `all_messages()` is untouched. No configuration.
-- `RAGState.evidence` / `AnalysisState.evidence` (`CapabilityEvidenceRecord`): which evidence a capability retrieved and cited, per question, keyed by message-count question identities and epochs. `haiku.rag.capabilities.ledger.citation_status(records, question=...)` derives `missing` / `grounded` / `ungrounded` across capabilities.
+- `RAGState.evidence` / `AnalysisState.evidence` (`CapabilityEvidenceRecord`): which evidence a capability retrieved and cited, per question, keyed by message-count question identities and epochs, and whether that question is still being answered. `haiku.rag.capabilities.ledger.citation_status(records, question=...)` derives `missing` / `grounded` / `ungrounded` across capabilities.
 - `RAGCapabilityBase.evidence_tool_names()` and `get_picture_bytes()`.
 - `haiku.rag.tools.search.decode_picture()`.
 
@@ -14,7 +14,7 @@
 
 - `rag_cite` / `analysis_cite` accept an empty `chunk_ids`, recording the answer as ungrounded rather than failing the call, and the instructions no longer exempt a refusal or a corpus-level computation from citing.
 - `RAGCapability` and `AnalysisCapability` no longer rewrite the model request. Register `create_capability()` from `haiku.rag.capabilities.compaction` alongside them to keep earlier questions compacted.
-- Resuming a run (no prompt, deferred tool results, an unfinished history tail) raises `RuntimeError` unless the host carries the capability state from the run being resumed.
+- Resuming a run (deferred tool results, an unfinished history tail) raises `RuntimeError` unless the host carries the capability state from the run being resumed.
 
 ### Removed
 

haiku_rag_slim/haiku/rag/capabilities/_base.py

@@ -6,13 +6,19 @@ from pathlib import Path
 from typing import Any, cast
 
 from pydantic import BaseModel
-from pydantic_ai import ModelRetry, RunContext, ToolFailed
+from pydantic_ai import (
+    DeferredToolRequests,
+    ModelRetry,
+    RunContext,
+    ToolFailed,
+)
 from pydantic_ai.capabilities import AbstractCapability
 from pydantic_ai.messages import (
     InstructionPart,
     ModelMessage,
     ModelRequest,
     ModelResponse,
+    RetryPromptPart,
     ToolCallPart,
     ToolReturn,
 )
@@ -82,29 +88,21 @@ def _clear_invocation_state(state: BaseModel) -> None:
             value.clear()
 
 
-def _is_resumption(prompt: Any, messages: list[ModelMessage]) -> bool:
-    """Whether this run continues a question rather than asking a new one.
+def _awaits_the_model(messages: list[ModelMessage]) -> bool:
+    """Whether the history unmistakably leaves the model something to answer.
 
-    Two signals, either of which is enough, because getting this wrong hands the
-    model a notice where its own evidence should be:
-
-    - no prompt: how pydantic-ai resumes for interruptions and suspensions.
-    - an unfinished tail: the history ends with a request the model has not
-      answered, or with a response whose tool calls have no returns yet. Deferred
-      tool results may arrive *with* a prompt, so the prompt alone is not enough.
-
-    A settled history ends with the previous answer, so a genuinely new question
-    is not mistaken for a continuation. The framework's own first-new-message
-    index would be better than either signal, but it is not public here.
+    Used to validate what the record already says, never to decide it. Only two
+    shapes are unambiguous: a response whose tool calls have no returns, and a
+    retry the model has not answered. A trailing tool return is not one of them,
+    being both how a settled structured answer ends and how results reach a
+    question still in progress.
     """
-    if prompt is None:
-        return True
     if not messages:
         return False
     last = messages[-1]
-    if isinstance(last, ModelRequest):
-        return True
-    return any(isinstance(part, ToolCallPart) for part in last.parts)
+    if isinstance(last, ModelResponse):
+        return any(isinstance(part, ToolCallPart) for part in last.parts)
+    return any(isinstance(part, RetryPromptPart) for part in last.parts)
 
 
 def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) -> bool:
@@ -154,11 +152,10 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
         outer = getattr(ctx.deps, "state", None)
         outer_state = outer if isinstance(outer, dict) else None
         raw_state = outer_state.get(self.state_namespace) if outer_state else None
-        resuming = _is_resumption(ctx.prompt, ctx.messages)
-        continuing = resuming and bool(ctx.messages)
         state = self.state_type.model_validate(raw_state or {})
         record = cast(CapabilityEvidenceRecord, cast(Any, state).evidence)
-        if continuing and record.question is None:
+        continuing = record.in_progress
+        if not continuing and _awaits_the_model(ctx.messages):
             raise RuntimeError(
                 f"The {self.state_namespace} capability is resuming a question with "
                 "no stored question identity. Capabilities cannot be added, removed "
@@ -321,6 +318,16 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
     async def after_run(
         self, ctx: RunContext[Any], *, result: AgentRunResult[Any]
     ) -> AgentRunResult[Any]:
+        """Close the question, unless the run is only pausing for deferred results.
+
+        A run that raised never arrives here, which is what leaves an interrupted
+        question in progress for the resumption to claim.
+        """
+        if self.state is not None and not isinstance(
+            result.output, DeferredToolRequests
+        ):
+            self._evidence_record().end_question()
+            self._sync_state()
         await self._close()
         return result
 

haiku_rag_slim/haiku/rag/capabilities/ledger.py

@@ -51,6 +51,11 @@ class CapabilityEvidenceRecord(BaseModel):
     overwritten by whichever of them synced its state last; merging happens in the
     transient views built by ``citation_status`` and the optional capabilities.
 
+    ``in_progress`` is whether the question this record names is still being
+    answered. It is the only authority on that: a transcript ending in a tool
+    return is a settled structured answer and a question awaiting its model
+    equally, so the shape of the history cannot decide it.
+
     ``question`` is the number of messages that existed when the question arrived,
     and ``epoch`` the number when an outcome occurred. Both are derived from the
     conversation rather than counted locally, so every participant computes the
@@ -61,6 +66,7 @@ class CapabilityEvidenceRecord(BaseModel):
 
     occurrences: dict[str, EvidenceOccurrence] = Field(default_factory=dict)
     question: int | None = None
+    in_progress: bool = False
     latest_evidence_epoch: int = 0
     declaration: CitationDeclaration | None = None
 
@@ -110,9 +116,14 @@ class CapabilityEvidenceRecord(BaseModel):
                 "one question from the next and are compared as recency."
             )
         self.question = identity
+        self.in_progress = True
         self.latest_evidence_epoch = 0
         self.declaration = None
 
+    def end_question(self) -> None:
+        """Mark the question answered, so the next run knows it is a new one."""
+        self.in_progress = False
+
     def note_evidence(self, epoch: int) -> None:
         """Record that the model has seen an evidence outcome.
 

tests/capabilities/test_capabilities.py

@@ -4,7 +4,16 @@ from typing import Any, cast
 from unittest.mock import AsyncMock, patch
 
 import pytest
-from pydantic_ai import Agent, DeferredToolResults, ModelRetry, RunContext, ToolFailed
+from pydantic import BaseModel
+from pydantic_ai import (
+    Agent,
+    CallDeferred,
+    DeferredToolRequests,
+    DeferredToolResults,
+    ModelRetry,
+    RunContext,
+    ToolFailed,
+)
 from pydantic_ai.messages import (
     ModelRequest,
     ModelResponse,
@@ -920,9 +929,9 @@ async def test_a_resumption_keeps_the_identity_of_the_question_in_progress(
     agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
     deps = Deps(
         state={
-            "rag": RAGState(evidence=CapabilityEvidenceRecord(question=7)).model_dump(
-                mode="json"
-            )
+            "rag": RAGState(
+                evidence=CapabilityEvidenceRecord(question=7, in_progress=True)
+            ).model_dump(mode="json")
         }
     )
     history = [
@@ -1159,10 +1168,12 @@ async def test_citing_without_searching_grounds_the_question(temp_db_path):
 
 @pytest.mark.asyncio
 async def test_a_host_seeded_record_does_not_pass_for_a_resumption(temp_db_path):
-    """A default record is truthy, so its presence cannot stand in for identity.
+    """A seeded record says nothing about a question, so the history has to.
 
-    Seeding one is what a host does when it has no state to send, and taking it
-    at face value would silently answer as question zero.
+    Seeding one is what a host does when it has no state to send. Its flag is
+    unset, so a history that unmistakably awaits the model means the host dropped
+    the state of a question in progress, and answering as question zero would
+    silently relabel it.
     """
     rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
 
@@ -1174,7 +1185,10 @@ async def test_a_host_seeded_record_does_not_pass_for_a_resumption(temp_db_path)
     with pytest.raises(RuntimeError, match="no stored question identity"):
         await agent.run(
             "carry on",
-            message_history=_in_flight_history(),
+            message_history=[
+                *_in_flight_history(),
+                ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]),
+            ],
             deps=Deps(state={"rag": RAGState().model_dump(mode="json")}),
         )
 
@@ -1208,6 +1222,9 @@ async def test_a_resumption_keeps_the_evidence_the_question_already_gathered(
         interrupted = await agent.run("what does the supervisor do?", deps=deps)
         identity = _record(deps, "rag").question
         assert identity is not None
+        # A run that ends awaiting external work leaves the question in progress,
+        # which is what the resumption claims. See the deferred-request test.
+        deps.state["rag"]["evidence"]["in_progress"] = True
         await agent.run(
             deferred_tool_results=DeferredToolResults(
                 calls={"call-2": "external result"}
@@ -1280,3 +1297,147 @@ async def test_citing_nothing_after_citing_something_keeps_it_grounded(temp_db_p
 
     record = capability.state.evidence
     assert citation_status([record], question=0) == "grounded"
+
+
+@pytest.mark.asyncio
+async def test_a_promptless_run_on_a_settled_history_is_a_new_question(temp_db_path):
+    """AG-UI hosts never pass a prompt: the client's message is the history.
+
+    Pydantic AI's UI adapter builds `message_history` from the frontend messages
+    and calls the agent without a prompt, so the run has no prompt *and* the
+    history ends with the user's own request. Reading either as a continuation
+    fails every AG-UI host on its first message.
+    """
+    rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
+
+    async def model(_messages, _info):
+        return ModelResponse(parts=[TextPart("answer")])
+
+    agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
+    deps = Deps(state={"rag": RAGState().model_dump(mode="json")})
+    history: list[Any] = [
+        ModelRequest(parts=[UserPromptPart("what does the manual say about masks?")])
+    ]
+
+    await agent.run(message_history=history, deps=deps)
+
+    assert _record(deps, "rag").question == len(history)
+
+
+@pytest.mark.asyncio
+async def test_a_promptless_run_on_an_unfinished_tail_is_still_a_continuation(
+    temp_db_path,
+):
+    """A suspended run resumes without a prompt, and must keep its question."""
+    rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
+
+    async def model(_messages, _info):
+        return ModelResponse(parts=[TextPart("answer")])
+
+    agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
+    deps = Deps(
+        state={
+            "rag": RAGState(
+                evidence=CapabilityEvidenceRecord(question=3, in_progress=True)
+            ).model_dump(mode="json")
+        }
+    )
+
+    await agent.run(message_history=_in_flight_history(), deps=deps)
+
+    assert _record(deps, "rag").question == 3
+
+
+@pytest.mark.asyncio
+async def test_a_structured_answer_does_not_leave_the_question_in_progress(
+    temp_db_path,
+):
+    """A settled run ends with a tool return, which says nothing about progress.
+
+    Pydantic AI answers a structured `output_type` by calling an output tool, so the
+    history ends with a request carrying that tool's return. Reading the transcript
+    shape alone, that is indistinguishable from tool results delivered to a question
+    still being answered, and every following question inherited the first one's
+    identity.
+    """
+
+    class Answer(BaseModel):
+        text: str
+
+    rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
+
+    async def model(_messages, info):
+        return ModelResponse(
+            parts=[ToolCallPart(info.output_tools[0].name, {"text": "answer"})]
+        )
+
+    agent = Agent(
+        FunctionModel(model),
+        deps_type=Deps,
+        capabilities=[rag],
+        output_type=Answer,
+    )
+    deps = Deps()
+
+    first = await agent.run("first question", deps=deps)
+    first_identity = _record(deps, "rag").question
+    assert _record(deps, "rag").in_progress is False
+
+    await agent.run(
+        "second question", deps=deps, message_history=list(first.all_messages())
+    )
+
+    second_identity = _record(deps, "rag").question
+    assert first_identity == 0
+    assert second_identity is not None and second_identity > 0
+
+
+@pytest.mark.asyncio
+async def test_a_run_pausing_for_deferred_work_leaves_the_question_in_progress(
+    temp_db_path,
+):
+    """The question is unfinished, so its resumption must find it claimable.
+
+    A deferred tool call ends the run with `DeferredToolRequests` rather than an
+    answer. Closing the question here would let the resumption relabel it.
+    """
+    rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
+
+    async def model(_messages, _info):
+        return ModelResponse(parts=[ToolCallPart("external_tool", {})])
+
+    agent = Agent(
+        FunctionModel(model),
+        deps_type=Deps,
+        capabilities=[rag],
+        output_type=[str, DeferredToolRequests],
+    )
+
+    @agent.tool_plain
+    def external_tool() -> str:
+        raise CallDeferred
+
+    deps = Deps()
+
+    result = await agent.run("a question needing external work", deps=deps)
+
+    assert isinstance(result.output, DeferredToolRequests)
+    assert _record(deps, "rag").in_progress is True
+
+
+@pytest.mark.asyncio
+async def test_an_answered_question_is_no_longer_in_progress(temp_db_path):
+    """The flag is what tells the next run it is asking something new."""
+    rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
+
+    async def model(_messages, _info):
+        return ModelResponse(parts=[TextPart("answer")])
+
+    agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag])
+    deps = Deps()
+
+    await agent.run("a question", deps=deps)
+
+    record = _record(deps, "rag")
+    assert record.question == 0
+    assert record.in_progress is False

tests/capabilities/test_citation_policy.py

@@ -365,7 +365,9 @@ async def test_a_resumed_question_is_not_redirected_twice(temp_db_path):
 
     with patch.object(RAGCapability, "_search", stub_search):
         first = await agent.run("what does the supervisor do?", deps=deps)
-        # The same question again, continued rather than asked anew.
+        # The same question again, continued rather than asked anew: a run that
+        # ends awaiting external work leaves the question in progress.
+        deps.state["rag"]["evidence"]["in_progress"] = True
         await agent.run(
             deps=deps,
             message_history=[

tests/capabilities/test_evidence_wire.py

@@ -299,11 +299,11 @@ def in_flight_history() -> list[Any]:
 
 
 def resuming_deps(question: int = 0) -> Deps:
-    """State as a resumption always finds it: the question already identified."""
+    """State as a resumption always finds it: the question identified and unfinished."""
     return Deps(
         state={
             "rag": RAGState(
-                evidence=CapabilityEvidenceRecord(question=question)
+                evidence=CapabilityEvidenceRecord(question=question, in_progress=True)
             ).model_dump(mode="json")
         }
     )