Enforce a declaration wherever there is something to declare

Requiring an evidence outcome from the current question exempted the case enforcement
exists for: a follow-up about evidence already cited needs no new search, since that
evidence is still on the wire — in a capsule when a compactor is registered, in full
when not. The condition is now that the conversation has something to declare, either
an outcome in this question or evidence it has already cited, which is independent of
whether anything compacts. A conversation that has neither is still left alone.

Citing again cannot narrow a question at any epoch. Declarations merged only within one
epoch, so an empty second thought a request later replaced the refs with nothing and
reported a grounded question ungrounded. They merge while no evidence outcome has
followed the standing declaration, and only genuinely newer evidence starts one afresh.

Whether a question has already been asked to declare is read from the message history
rather than remembered on the run instance, which a resumption's `for_run` discarded —
the same question was asked twice. Reading the history also makes the right call when a
redirect was enqueued but the run ended before it reached the model: nothing is in the
history, so it is asked again. Violations are recorded once per question for the same
reason.

haiku_rag_slim/haiku/rag/capabilities/ledger.py

@@ -133,10 +133,11 @@ class CapabilityEvidenceRecord(BaseModel):
     ) -> None:
         """Record validated citations for the current question.
 
-        Repeated calls at the same epoch merge, so citing again cannot narrow what
-        was already declared: an empty call after a grounded one leaves it
-        grounded. A call at a later epoch declares afresh, because evidence the
-        model saw in between may be what it is now citing.
+        Citing again cannot narrow what a question already declared: calls merge
+        while no evidence outcome has followed the standing declaration, whatever
+        epoch they arrive at, so an empty second thought leaves a grounded question
+        grounded. Only genuinely newer evidence starts a declaration afresh, since
+        what the model saw in between may be what it is now citing.
         """
         if self.question is None:
             raise ValueError(
@@ -145,14 +146,16 @@ class CapabilityEvidenceRecord(BaseModel):
             )
         self._reject_regression(epoch, "A declaration")
         current = self.declaration
-        if current is not None and (current.question, current.epoch) == (
-            self.question,
-            epoch,
+        if (
+            current is not None
+            and current.question == self.question
+            and self.latest_evidence_epoch <= current.epoch
         ):
             known = {(ref.capability, ref.chunk_id) for ref in current.refs}
             current.refs.extend(
                 ref for ref in refs if (ref.capability, ref.chunk_id) not in known
             )
+            current.epoch = max(current.epoch, epoch)
         else:
             self.declaration = CitationDeclaration(
                 question=self.question, epoch=epoch, refs=list(refs)

haiku_rag_slim/haiku/rag/capabilities/policy.py

@@ -1,10 +1,15 @@
-from dataclasses import dataclass, field, replace
+from dataclasses import dataclass
 from typing import Any
 
 from pydantic import BaseModel, Field
 from pydantic_ai import RunContext
 from pydantic_ai.capabilities import AbstractCapability
-from pydantic_ai.messages import ModelResponse, ToolCallPart
+from pydantic_ai.messages import (
+    ModelMessage,
+    ModelResponse,
+    ToolCallPart,
+    UserPromptPart,
+)
 from pydantic_ai.models import ModelRequestContext
 
 from haiku.rag.capabilities.evidence import (
@@ -58,12 +63,6 @@ class CitationPolicyCapability(AbstractCapability[Any]):
     would share this capability's id.
     """
 
-    redirected: set[int] = field(default_factory=set, repr=False)
-
-    async def for_run(self, ctx: RunContext[Any]) -> "CitationPolicyCapability":
-        """Give the run its own record of what it has already asked for."""
-        return replace(self, redirected=set())
-
     async def after_model_request(
         self,
         ctx: RunContext[Any],
@@ -82,13 +81,14 @@ class CitationPolicyCapability(AbstractCapability[Any]):
             return response
         evidence = discover_evidence(ctx)
         question = question_in_progress(evidence)
-        if question in self.redirected or not _gathered_evidence(evidence):
+        if not _has_evidence_to_declare(evidence):
             return response
         records = [found.record for found in evidence]
         if citation_status(records, question=question) != "missing":
             return response
+        if _already_asked(ctx.messages, question):
+            return response
 
-        self.redirected.add(question)
         if any(found.cite_available for found in evidence):
             ctx.enqueue(REDIRECT, priority="when_idle")
         else:
@@ -96,12 +96,17 @@ class CitationPolicyCapability(AbstractCapability[Any]):
         return response
 
     def _record_violation(self, ctx: RunContext[Any], question: int) -> None:
-        """Note a question that could not be asked to cite, the tool being gone."""
+        """Note a question that could not be asked to cite, the tool being gone.
+
+        Recorded once per question: a resumption of the same question decides
+        again, and one question is one outcome.
+        """
         outer = getattr(ctx.deps, "state", None)
         if not isinstance(outer, dict):
             return
         state = CitationPolicyState.model_validate(outer.get(STATE_NAMESPACE) or {})
-        state.violations.append(question)
+        if question not in state.violations:
+            state.violations.append(question)
         outer[STATE_NAMESPACE] = state.model_dump(mode="json")
 
     async def before_run(self, ctx: RunContext[Any]) -> None:
@@ -113,16 +118,42 @@ class CitationPolicyCapability(AbstractCapability[Any]):
             )
 
 
-def _gathered_evidence(evidence: list[DiscoveredEvidence]) -> bool:
-    """Whether this question produced anything an answer could be grounded on.
+def _already_asked(messages: list[ModelMessage], question: int) -> bool:
+    """Whether this question has already been asked to declare its grounding.
 
-    A question with no evidence outcome has nothing to declare — a greeting, or a
-    conversational aside. Read from the ledger rather than from ``state.searches``,
-    which a new question clears, so an answer grounded on code execution or on a
-    document read counts as well.
+    Read from the history rather than remembered on the instance, which a
+    resumption's ``for_run`` would forget — the same question would then be asked
+    twice. It also makes the right call when a redirect was enqueued but the run
+    ended before it reached the model: nothing is in the history, so it is asked
+    again, which is what the model needs.
+    """
+    return any(
+        isinstance(part, UserPromptPart)
+        and isinstance(part.content, str)
+        and REDIRECT_HINT in part.content
+        for message in messages[question:]
+        for part in message.parts
+    )
+
+
+def _has_evidence_to_declare(evidence: list[DiscoveredEvidence]) -> bool:
+    """Whether anything exists that this answer could have been grounded on.
+
+    Either this question produced an evidence outcome, or the conversation has
+    already cited something — which stays available to a later answer, in a capsule
+    if a compactor is registered and in full if not. Requiring a fresh outcome
+    exempted exactly the follow-up that reuses earlier evidence, which is the case
+    enforcement exists for.
+
+    A conversation that has neither has nothing to declare: a greeting, an aside.
+    Read from the ledger rather than from ``state.searches``, which a new question
+    clears, so an answer grounded on code execution or a document read counts too.
     """
     question = question_in_progress(evidence)
-    return any(found.record.latest_evidence_epoch > question for found in evidence)
+    return any(
+        found.record.latest_evidence_epoch > question or found.record.occurrences
+        for found in evidence
+    )
 
 
 def create_capability() -> CitationPolicyCapability:

tests/capabilities/test_citation_policy.py

@@ -3,7 +3,7 @@ from typing import Any, cast
 from unittest.mock import patch
 
 import pytest
-from pydantic_ai import Agent
+from pydantic_ai import Agent, DeferredToolResults
 from pydantic_ai.exceptions import UserError
 from pydantic_ai.messages import ModelResponse, TextPart, ToolCallPart
 from pydantic_ai.models.function import FunctionModel
@@ -280,3 +280,100 @@ async def test_a_violation_with_nowhere_to_record_it_does_not_fail_the_run(
         result = await agent.run("what does the supervisor do?", deps=StatelessDeps())
 
     assert result.output == "an answer with no citation"
+
+
+@pytest.mark.asyncio
+async def test_a_follow_up_answered_from_retained_evidence_is_enforced(temp_db_path):
+    """The multi-turn case is the one enforcement exists for.
+
+    A follow-up about something already cited needs no new search — the evidence is
+    still on the wire, whether in a capsule or in full — so requiring a fresh
+    evidence outcome let exactly those answers through undeclared.
+    """
+    rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
+    turns = iter(
+        [
+            [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
+            [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
+            [TextPart("first answer")],
+            [TextPart("a follow-up answered from what is already here")],
+            [TextPart("a follow-up answered from what is already here")],
+        ]
+    )
+    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, capabilities=[rag, create_policy()]
+    )
+    deps = Deps()
+
+    with patch.object(RAGCapability, "_search", stub_search):
+        first = await agent.run("what does the supervisor do?", deps=deps)
+        await agent.run(
+            "and what colour is the box in it?",
+            deps=deps,
+            message_history=first.all_messages(),
+        )
+
+    assert [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]
+
+
+@pytest.mark.asyncio
+async def test_a_conversation_that_never_cited_anything_is_still_left_alone(
+    temp_db_path,
+):
+    """A greeting has nothing to declare, and no evidence exists to declare from."""
+    _, _, sent = await run_with_policy(temp_db_path, [[TextPart("hello back")]])
+
+    assert not [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]
+
+
+@pytest.mark.asyncio
+async def test_a_resumed_question_is_not_redirected_twice(temp_db_path):
+    """Once per question has to mean once, across every run of that question.
+
+    Tracking it on the run instance forgot it at the next `for_run`, so resuming an
+    interrupted question asked for the citation again.
+    """
+    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("uncited again")],
+            [TextPart("uncited a third time")],
+            [TextPart("uncited a fourth time")],
+            [TextPart("uncited a fifth time")],
+        ]
+    )
+    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, capabilities=[rag, create_policy()]
+    )
+    deps = Deps()
+
+    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.
+        await agent.run(
+            deps=deps,
+            message_history=[
+                *first.all_messages(),
+                ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-9")]),
+            ],
+            deferred_tool_results=DeferredToolResults(
+                calls={"call-9": "external result"}
+            ),
+        )
+
+    redirects = [p for p in prompts_of(sent[-1]) if REDIRECT_HINT in p]
+    assert len(redirects) == 1

tests/capabilities/test_evidence_ledger.py

@@ -240,3 +240,38 @@ def test_a_question_starts_clear_of_the_one_before_it():
     assert record.latest_evidence_epoch == 0
     assert record.declaration is None
     assert citation_status([record], question=9) == "missing"
+
+
+def test_an_empty_citation_after_a_grounded_one_cannot_narrow_it():
+    """Citing again must not weaken a declaration, at any epoch.
+
+    Merging only within one epoch meant a second thought a request later replaced
+    the refs with nothing and reported the question ungrounded.
+    """
+    record = CapabilityEvidenceRecord(question=0)
+    record.declare([rag_ref()], epoch=3)
+    record.declare([], epoch=5)
+
+    assert record.declaration is not None
+    assert [ref.chunk_id for ref in record.declaration.refs] == ["c1"]
+    assert citation_status([record], question=0) == "grounded"
+
+
+def test_a_declaration_after_newer_evidence_starts_afresh():
+    """Evidence the model has since seen may be what it is now citing."""
+    record = CapabilityEvidenceRecord(question=0)
+    record.declare([rag_ref("first")], epoch=3)
+    record.note_evidence(4)
+    record.declare([rag_ref("second")], epoch=5)
+
+    assert record.declaration is not None
+    assert [ref.chunk_id for ref in record.declaration.refs] == ["second"]
+
+
+def test_an_empty_citation_after_newer_evidence_is_ungrounded():
+    record = CapabilityEvidenceRecord(question=0)
+    record.declare([rag_ref()], epoch=3)
+    record.note_evidence(4)
+    record.declare([], epoch=5)
+
+    assert citation_status([record], question=0) == "ungrounded"