Notice the endings a structured answer arrives in

An output tool call is a `ToolCallPart` like any other, and treating every tool call as
intermediate meant a model could search, skip citing, emit its structured answer and
finish with neither a redirect nor a record. A response ends the question when it
carries no tool calls, or when one of its calls names an output tool.

Some endings are not visible from a single response — a host running
`end_strategy="early"` can finish on text beside a function call — so `after_run` is
the backstop: it cannot ask the model for anything by then, but it records a question
that reached the end of its run undeclared. That also covers a question that was asked
once, ignored, and finished anyway, which previously returned early on the redirect
marker and went unrecorded.

The capability documentation and the changelog said a question that gathered no
evidence is left alone. That describes neither the code nor the intent: enforcement
applies wherever there is something to declare, which includes a follow-up that reuses
evidence cited earlier without searching again.

CHANGELOG.md

@@ -3,7 +3,7 @@
 
 ### Added
 
-- `CitationPolicyCapability` (`haiku.rag.capabilities.policy.create_capability`): registering it requires every answer to declare its grounding. A question that ends undeclared is sent back once to record what grounded the answer already given; when the cite tool is no longer available the question is recorded in `CitationPolicyState.violations` instead. A question that gathered no evidence is left alone.
+- `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 that has never cited anything 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.

docs/capabilities/index.md

@@ -90,10 +90,18 @@ a declaration possible without forcing the model to invent grounding.
 
 When a question ends undeclared, the model is asked once to record what grounded the
 answer it already gave. It is not asked to change the answer. If the cite tool is no
-longer available by then, the question is recorded as a violation in
-`CitationPolicyState` under `"citation_policy"` instead, since pointing a model at a
-tool that is gone costs it retries. A question that gathered no evidence at all — a
-greeting, an aside — is left alone.
+longer available by then, or the question finishes undeclared anyway, it is recorded as
+a violation in `CitationPolicyState` under `"citation_policy"`, since pointing a model
+at a tool that is gone costs it retries.
+
+What gets enforced is every answer in a conversation that has something to declare:
+either this question retrieved evidence, or the conversation has already cited
+something, which stays available to later answers. So a follow-up about evidence cited
+earlier is enforced even though it searched nothing — that case is the reason the
+capability exists. It also means that once anything has been cited, later turns are
+enforced too, a greeting included; the model satisfies the policy by citing an empty
+list, at the cost of one extra request. A conversation that has never cited anything is
+not enforced at all.
 
 Exactly one policy capability makes the decision, however many evidence capabilities
 are registered, so two of them cannot each demand a citation for one answer.

haiku_rag_slim/haiku/rag/capabilities/policy.py

@@ -11,6 +11,7 @@ from pydantic_ai.messages import (
     UserPromptPart,
 )
 from pydantic_ai.models import ModelRequestContext
+from pydantic_ai.run import AgentRunResult
 
 from haiku.rag.capabilities.evidence import (
     DiscoveredEvidence,
@@ -80,12 +81,17 @@ class CitationPolicyCapability(AbstractCapability[Any]):
     ) -> ModelResponse:
         """Decide once, at the last moment a question can still be redirected.
 
-        A response carrying no tool calls ends the question, so there is no later
-        opportunity. Citing is unconditional, so an undeclared answer is a protocol
-        breach whether the model answered or refused, and this never has to guess
-        which it was.
+        A response that ends the question is the last opportunity: one carrying no
+        tool calls, or one whose call is an output tool, which is how a structured
+        answer arrives and which finishes the run just the same. Citing is
+        unconditional, so an undeclared answer is a protocol breach whether the model
+        answered or refused, and this never has to guess which it was.
+
+        Endings this cannot see — a host running ``end_strategy="early"`` can finish
+        on text beside a function call — are caught by ``after_run``, which can still
+        record the outcome even though it can no longer ask for a citation.
         """
-        if any(isinstance(part, ToolCallPart) for part in response.parts):
+        if not _ends_the_question(response, request_context):
             return response
         evidence = discover_evidence(ctx)
         question = question_in_progress(evidence)
@@ -117,6 +123,25 @@ class CitationPolicyCapability(AbstractCapability[Any]):
             state.violations.append(question)
         outer[STATE_NAMESPACE] = state.model_dump(mode="json")
 
+    async def after_run(
+        self, ctx: RunContext[Any], *, result: AgentRunResult[Any]
+    ) -> AgentRunResult[Any]:
+        """Record a question that finished undeclared, whatever ended it.
+
+        The backstop for an ending ``after_model_request`` cannot recognise. Nothing
+        can be asked of the model now, so this only records: a question that reached
+        the end of its run without a declaration is a violation, and one that was
+        asked but never answered is the same.
+        """
+        evidence = discover_evidence(ctx)
+        question = question_in_progress(evidence)
+        if not _has_evidence_to_declare(evidence):
+            return result
+        records = [found.record for found in evidence]
+        if citation_status(records, question=question) == "missing":
+            self._record_violation(ctx, question)
+        return result
+
     async def before_run(self, ctx: RunContext[Any]) -> None:
         """Publish an empty outcome, so a host can tell "none" from "not running"."""
         outer = getattr(ctx.deps, "state", None)
@@ -126,6 +151,24 @@ class CitationPolicyCapability(AbstractCapability[Any]):
             )
 
 
+def _ends_the_question(
+    response: ModelResponse, request_context: ModelRequestContext
+) -> bool:
+    """Whether this response finishes the question rather than continuing it.
+
+    An output tool call is a ``ToolCallPart`` like any other, but it carries the
+    final answer and ends the run, so treating every tool call as intermediate let a
+    structured answer finish undeclared.
+    """
+    calls = [part for part in response.parts if isinstance(part, ToolCallPart)]
+    if not calls:
+        return True
+    output_tools = {
+        tool.name for tool in request_context.model_request_parameters.output_tools
+    }
+    return any(call.tool_name in output_tools for call in calls)
+
+
 def _already_asked(messages: list[ModelMessage], question: int) -> bool:
     """Whether this question has already been asked to declare its grounding.
 

tests/capabilities/test_citation_policy.py

@@ -3,6 +3,7 @@ from typing import Any, cast
 from unittest.mock import patch
 
 import pytest
+from pydantic import BaseModel
 from pydantic_ai import Agent, DeferredToolResults
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart
@@ -412,3 +413,85 @@ async def test_a_user_quoting_the_redirect_does_not_suppress_enforcement(temp_db
         )
 
     assert [p for p in prompts_of(sent[-1]) if CITATION_REDIRECT_TAG in p]
+
+
+class Answer(BaseModel):
+    """A structured output, which the model returns through an output tool."""
+
+    text: str
+
+
+@pytest.mark.asyncio
+async def test_a_structured_output_answer_does_not_escape_enforcement(temp_db_path):
+    """An output tool call is a `ToolCallPart` too, and it ends the run.
+
+    Treating every tool call as intermediate let a model search, skip citing, emit
+    its structured answer and finish with neither a redirect nor a violation.
+    """
+    rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
+    turns = iter(
+        [
+            [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
+            [
+                ToolCallPart(
+                    "final_result", {"text": "uncited structured answer"}, "out"
+                )
+            ],
+            [
+                ToolCallPart(
+                    "final_result", {"text": "uncited structured answer"}, "out"
+                )
+            ],
+        ]
+    )
+    sent: list[list[Any]] = []
+
+    async def model(messages, _info):
+        sent.append(list(messages))
+        return ModelResponse(parts=next(turns))
+
+    agent = Agent(
+        FunctionModel(model),
+        deps_type=Deps,
+        output_type=Answer,
+        capabilities=[rag, create_policy()],
+    )
+    deps = Deps()
+
+    with patch.object(RAGCapability, "_search", stub_search):
+        await agent.run("what does the supervisor do?", deps=deps)
+
+    redirected = [p for p in prompts_of(sent[-1]) if CITATION_REDIRECT_TAG in p]
+    violations = deps.state.get("citation_policy", {}).get("violations", [])
+    assert redirected or violations
+
+
+@pytest.mark.asyncio
+async def test_a_question_asked_once_and_still_undeclared_is_recorded(temp_db_path):
+    """Being asked is not an outcome; the question still ended undeclared.
+
+    Returning early on the redirect marker meant a question that was asked, ignored,
+    and then finished — with the cite tool possibly gone by that point — was neither
+    redirected again nor recorded anywhere.
+    """
+    rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
+    turns = iter(
+        [
+            [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
+            [TextPart("uncited")],
+            [TextPart("still uncited after being asked")],
+        ]
+    )
+
+    async def model(_messages, _info):
+        return ModelResponse(parts=next(turns))
+
+    agent = Agent(
+        FunctionModel(model), deps_type=Deps, capabilities=[rag, create_policy()]
+    )
+    deps = Deps()
+
+    with patch.object(RAGCapability, "_search", stub_search):
+        await agent.run("what does the supervisor do?", deps=deps)
+
+    assert deps.state["citation_policy"]["violations"] == [0]