From e1dd8517f9f05fc0554a321efc7d1fdfe998542d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 13 Aug 2026 15:46:23 +0300 Subject: [PATCH] Refuse to compact evidence the host kept no record of Both optional capabilities read what earlier questions retrieved and cited from the capability's state, so a host that carries only the message history hands every run an empty record. Compaction then replaced the earlier evidence with receipts and retained nothing, and the loss was invisible: the citations the host already displayed were still there. It now refuses when it finds evidence from an earlier question and no record of what that question cited. `state_carried` reaches the optional capabilities through discovery, so the refusal distinguishes a host that never carries state from a question that simply cited nothing. The documentation taught the pattern that breaks: the compose example is now stateful and the requirement is stated where each capability is introduced. The app's browser storage was doing exactly this, keeping only the fields the UI reads. It now persists the whole namespace map, so the citation policy's violations survive a reload as well as the evidence record. --- app/frontend/components/Chat.tsx | 21 +++--- app/frontend/lib/sessionStorage.ts | 29 ++++++-- docs/capabilities/compaction.md | 6 ++ docs/capabilities/index.md | 28 +++++++- docs/capabilities/policy.md | 11 ++- .../haiku/rag/capabilities/_base.py | 9 +++ .../haiku/rag/capabilities/compaction.py | 15 +++- .../haiku/rag/capabilities/evidence.py | 2 + tests/capabilities/test_evidence_capsule.py | 1 + tests/capabilities/test_evidence_wire.py | 68 ++++++++++++++++++- 10 files changed, 163 insertions(+), 27 deletions(-) diff --git a/app/frontend/components/Chat.tsx b/app/frontend/components/Chat.tsx index 30f1f6a7..2d0770c5 100644 --- a/app/frontend/components/Chat.tsx +++ b/app/frontend/components/Chat.tsx @@ -20,6 +20,8 @@ import { import { FilterIcon } from "../lib/icons"; import type { RAGState } from "../lib/sessionStorage"; import { + AGUI_STATE_KEY, + agentStateOf, createSession, getActiveSessionId, getLatestCitations, @@ -32,10 +34,7 @@ import DbInfo from "./DbInfo"; import DocumentFilter from "./DocumentFilter"; import SessionManager from "./SessionManager"; -// Must match RAGCapability.state_namespace. -const AGUI_STATE_KEY = "rag"; - -// AG-UI state is namespaced under AGUI_STATE_KEY +// AG-UI state is namespaced under AGUI_STATE_KEY (see sessionStorage). interface AgentState { [AGUI_STATE_KEY]?: RAGState; } @@ -319,11 +318,10 @@ function ChatContentInner({ useEffect(() => { if (agent.messages.length > 0) return; const session = getSession(sessionId); - // Seed state for the capability; the backend replaces it after each run. - agent.setState({ - ...agent.state, - [AGUI_STATE_KEY]: normalizeRAGState(session?.ragState), - }); + // Seed state for the capabilities; the backend replaces it after each run. + // The whole namespace map goes back, not just the fields this UI reads: + // compaction and the citation policy read what earlier questions recorded. + agent.setState({ ...agent.state, ...agentStateOf(session ?? undefined) }); if (session && session.messages.length > 0) { // biome-ignore lint/suspicious/noExplicitAny: AG-UI Message type is a broad union agent.setMessages(session.messages as any[]); @@ -336,13 +334,10 @@ function ChatContentInner({ // biome-ignore lint/correctness/useExhaustiveDependencies: JSON.stringify tracks content changes useEffect(() => { if (sessionId && agent.messages.length > 0) { - const currentRagState = normalizeRAGState( - (agent.state as AgentState)?.[AGUI_STATE_KEY], - ); updateSessionMessages( sessionId, serializeMessages(agent.messages), - currentRagState, + (agent.state ?? {}) as Record, ); } }, [JSON.stringify(agent.messages), ragState, sessionId]); diff --git a/app/frontend/lib/sessionStorage.ts b/app/frontend/lib/sessionStorage.ts index 69d8ac7c..63642444 100644 --- a/app/frontend/lib/sessionStorage.ts +++ b/app/frontend/lib/sessionStorage.ts @@ -34,11 +34,32 @@ export interface StoredSession { id: string; title: string; messages: StoredMessage[]; - ragState: RAGState; + // The whole AG-UI state. The rag namespace is not the only one a capability + // writes: the citation policy records violations beside it. + agentState: AgentState; + // Sessions stored before agentState existed. + ragState?: RAGState; createdAt: string; updatedAt: string; } +export const AGUI_STATE_KEY = "rag"; + +export type AgentState = Record; + +// Reads the rag namespace out of a stored session, whichever way it was stored. +export function ragStateOf(session?: StoredSession): RAGState { + const namespaced = session?.agentState?.[AGUI_STATE_KEY] as + | Partial + | undefined; + return normalizeRAGState(namespaced ?? session?.ragState); +} + +// The state to seed an agent with when a session is resumed. +export function agentStateOf(session?: StoredSession): AgentState { + return session?.agentState ?? { [AGUI_STATE_KEY]: ragStateOf(session) }; +} + const SESSIONS_KEY = "haiku.rag.sessions"; const ACTIVE_SESSION_KEY = "haiku.rag.activeSession"; @@ -86,7 +107,7 @@ export function createSession(): StoredSession { id: crypto.randomUUID(), title: "New Session", messages: [], - ragState: normalizeRAGState(), + agentState: { [AGUI_STATE_KEY]: normalizeRAGState() }, createdAt: now, updatedAt: now, }; @@ -111,7 +132,7 @@ export function saveSession(session: StoredSession): void { export function updateSessionMessages( id: string, messages: StoredMessage[], - ragState: RAGState, + agentState: AgentState, ): void { const sessions = getAllSessions(); const idx = sessions.findIndex((s) => s.id === id); @@ -119,7 +140,7 @@ export function updateSessionMessages( const session = sessions[idx]; session.messages = messages; - session.ragState = ragState; + session.agentState = agentState; session.updatedAt = new Date().toISOString(); // Derive title from first user message diff --git a/docs/capabilities/compaction.md b/docs/capabilities/compaction.md index 46c6a038..e0baeb25 100644 --- a/docs/capabilities/compaction.md +++ b/docs/capabilities/compaction.md @@ -21,6 +21,12 @@ agent = Agent( It exposes no tools and takes no configuration. Registering it is the only switch: leave it out and the transcript reaches the model untouched. +The host must carry the capability state between runs, alongside the message +history: the capsule is built from what earlier questions recorded there. Given +only a message history, every run starts from an empty record, and compaction +refuses rather than replace evidence it cannot retain. See +[Compose an agent](index.md#compose-an-agent) for the shape. + ## What it does On each request, evidence from earlier questions is replaced by the evidence those diff --git a/docs/capabilities/index.md b/docs/capabilities/index.md index 15636fff..f4b0e105 100644 --- a/docs/capabilities/index.md +++ b/docs/capabilities/index.md @@ -16,11 +16,22 @@ The two evidence capabilities are deferred by default. An agent initially sees o Pick one evidence capability, and add both optional capabilities to it: ```python +from dataclasses import dataclass, field +from typing import Any + from pydantic_ai import Agent +from pydantic_ai.messages import ModelMessage + from haiku.rag.capabilities.compaction import create_capability as compaction from haiku.rag.capabilities.policy import create_capability as citation_policy from haiku.rag.capabilities.rag import create_capability as rag + +@dataclass +class Deps: + state: dict[str, Any] = field(default_factory=dict) + + agent = Agent( "openai:gpt-5", capabilities=[ @@ -28,12 +39,27 @@ agent = Agent( compaction(), citation_policy(), ], + deps_type=Deps, ) -result = await agent.run("What does the knowledge base say about X?") +# One Deps and one history for the conversation: the capabilities read both. +deps = Deps() +history: list[ModelMessage] = [] + +result = await agent.run("What does the knowledge base say about X?", deps=deps, message_history=history) +history = list(result.all_messages()) print(result.output) ``` +!!! warning "Both optional capabilities need the host to carry state" + + They read what earlier questions retrieved and cited from the capability's + state, so the host must expose a `state` dict on its agent dependencies and + hand the same dict back on every run of a conversation, alongside the message + history. With only the message history, every run starts from an empty record: + compaction refuses rather than replace evidence it cannot retain, and the + citation policy cannot enforce a follow-up about evidence cited earlier. + Swap `rag` for `analysis` for an analysis agent. Both optional capabilities work the same way with either one, and neither exposes tools or takes configuration. diff --git a/docs/capabilities/policy.md b/docs/capabilities/policy.md index de0b56df..19e26728 100644 --- a/docs/capabilities/policy.md +++ b/docs/capabilities/policy.md @@ -16,9 +16,14 @@ agent = Agent( ) ``` -It exposes no tools and takes no configuration. 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. +It exposes no tools and takes no configuration. 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. + +The host must carry the capability state between runs, alongside the message +history. Enforcement reads what the conversation has already cited, so without it +a follow-up about evidence cited earlier goes unenforced. See +[Compose an agent](index.md#compose-an-agent) for the shape. ## Declaring nothing is a valid answer diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 85052457..b9f0ff60 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -135,6 +135,13 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): request_count: int = field(default=0, repr=False) grace_requests_used: int = field(default=0, repr=False) epoch: int = field(default=0, repr=False) + state_carried: bool = field(default=False, repr=False) + """Whether the host handed back a record a previous question had stamped. + + False on a first question, and equally on every question of a host that does + not carry state between runs. Capabilities that need the record to mean + anything across questions read it to refuse rather than act on nothing. + """ async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]": """Start a run's own copy, and settle which question it is answering. @@ -155,6 +162,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): state = self.state_type.model_validate(raw_state or {}) record = cast(CapabilityEvidenceRecord, cast(Any, state).evidence) continuing = record.in_progress + state_carried = record.question is not None if not continuing and _awaits_the_model(ctx.messages): raise RuntimeError( f"The {self.state_namespace} capability is resuming a question with " @@ -176,6 +184,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): request_count=0, grace_requests_used=0, epoch=0, + state_carried=state_carried, ) run_capability._sync_state() return run_capability diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index 85e5170e..6b1decbb 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -321,14 +321,23 @@ class EvidenceCompactionCapability(AbstractCapability[Any]): """ evidence = discover_evidence(ctx) boundary = question_in_progress(evidence) + owned_tools = frozenset().union(*(found.tool_names for found in evidence)) + if boundary > 0 and _newest_owned_return( + request_context.messages, boundary, owned_tools + ): + if not any(found.state_carried for found in evidence): + raise RuntimeError( + "Evidence compaction found evidence from an earlier question but " + "no record of what it cited, so replacing it would retain " + "nothing. The host must carry the capability state between runs, " + "alongside the message history, for this capability to work." + ) if boundary > 0: await self._build_once(ctx, evidence) request_context.messages = compact_history( request_context.messages, boundary=boundary, - owned_tools=frozenset().union( - *(found.tool_names for found in evidence) - ), + owned_tools=owned_tools, capsule_text=self.capsule.text, capsule_images=self.images, ) diff --git a/haiku_rag_slim/haiku/rag/capabilities/evidence.py b/haiku_rag_slim/haiku/rag/capabilities/evidence.py index 75a61cc4..49f8f385 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/evidence.py +++ b/haiku_rag_slim/haiku/rag/capabilities/evidence.py @@ -22,6 +22,7 @@ class DiscoveredEvidence: citations: Mapping[str, Citation] tool_names: frozenset[str] cite_available: bool + state_carried: bool def discover_evidence(ctx: RunContext[Any]) -> list[DiscoveredEvidence]: @@ -40,6 +41,7 @@ def discover_evidence(ctx: RunContext[Any]) -> list[DiscoveredEvidence]: citations=cast(Any, capability.state).citation_index, tool_names=frozenset(capability.evidence_tool_names()), cite_available=capability.cite_available, + state_carried=capability.state_carried, ) for capability in ctx.capabilities.values() if isinstance(capability, RAGCapabilityBase) diff --git a/tests/capabilities/test_evidence_capsule.py b/tests/capabilities/test_evidence_capsule.py index 981db4c3..056accf3 100644 --- a/tests/capabilities/test_evidence_capsule.py +++ b/tests/capabilities/test_evidence_capsule.py @@ -68,6 +68,7 @@ def discovered( return DiscoveredEvidence( capability=capability, record=record, + state_carried=True, citations={ chunk_id: citation( chunk_id, diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py index 6ad32aa4..7e588ca1 100644 --- a/tests/capabilities/test_evidence_wire.py +++ b/tests/capabilities/test_evidence_wire.py @@ -287,6 +287,17 @@ def rag_and_compactor(temp_db_path): ) +def settled_deps(question: int = 0) -> Deps: + """State as a host carrying it has it: an earlier question, answered.""" + return Deps( + state={ + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=question, in_progress=False) + ).model_dump(mode="json") + } + ) + + def in_flight_history() -> list[Any]: """A question already asked and searched, still awaiting its answer.""" return [ @@ -322,7 +333,9 @@ async def test_without_the_compactor_the_history_is_untouched(temp_db_path): agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])] - await agent.run("a different question", deps=Deps(), message_history=settled) + await agent.run( + "a different question", deps=settled_deps(), message_history=settled + ) assert returns_of(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"] @@ -341,7 +354,9 @@ async def test_with_the_compactor_a_new_question_compacts_the_previous_one( agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])] - await agent.run("a different question", deps=Deps(), message_history=settled) + await agent.run( + "a different question", deps=settled_deps(), message_history=settled + ) assert returns_of(wire[-1]) == [RECEIPT] @@ -509,7 +524,8 @@ async def test_the_capsule_is_built_once_per_request_and_again_for_the_next( ctx = RunContext( deps=deps, model=TestModel(), usage=RunUsage(), run_id="run-1", run_step=1 ) - run_rag = await rag.for_run(ctx) + # A host carrying state, which is what compaction requires of one. + run_rag = replace(await rag.for_run(ctx), state_carried=True) run_compactor = await compactor.for_run(ctx) cast(Any, run_rag.state).evidence.begin_question(4) ctx = replace(ctx, capabilities={"rag": run_rag, "compaction": run_compactor}) @@ -727,3 +743,49 @@ def test_the_capsule_is_attached_beside_the_newest_return_of_that_request(): assert images_of(compacted) == [fresh] assert returns_of(compacted) == [RECEIPT, "CAPSULE"] + + +@pytest.mark.asyncio +async def test_compaction_refuses_to_strip_evidence_it_cannot_replace(temp_db_path): + """A host that does not carry state has no record to build a capsule from. + + Compacting anyway replaces the earlier evidence with receipts and retains + nothing, so the model loses what it cited and the loss is invisible: the + citations the host already displayed are still there. + """ + rag, compactor = rag_and_compactor(temp_db_path) + + async def model(_messages, _info): # pragma: no cover - never reached + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) + history = answered_question("an earlier question", evidence="EVIDENCE TO LOSE") + + with pytest.raises(RuntimeError, match="carry the capability state"): + await agent.run("a follow-up", deps=Deps(), message_history=history) + + +@pytest.mark.asyncio +async def test_compaction_proceeds_for_a_host_that_carries_state(temp_db_path): + """The same history, with the record the earlier question left behind.""" + rag, compactor = rag_and_compactor(temp_db_path) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) + history = answered_question("an earlier question", evidence="EVIDENCE TO LOSE") + + carried = Deps( + state={ + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=0, in_progress=False) + ).model_dump(mode="json") + } + ) + + await agent.run("a follow-up", deps=carried, message_history=history) + + assert returns_of(wire[-1]) == [RECEIPT]