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.
This commit is contained in:
Yiorgis Gozadinos 2026-08-13 15:46:23 +03:00
parent 53bbf52697
commit e1dd8517f9
No known key found for this signature in database
10 changed files with 163 additions and 27 deletions

View file

@ -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<string, unknown>,
);
}
}, [JSON.stringify(agent.messages), ragState, sessionId]);

View file

@ -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<string, unknown>;
// 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<RAGState>
| 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

View file

@ -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

View file

@ -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.

View file

@ -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

View file

@ -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

View file

@ -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,
)

View file

@ -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)

View file

@ -68,6 +68,7 @@ def discovered(
return DiscoveredEvidence(
capability=capability,
record=record,
state_carried=True,
citations={
chunk_id: citation(
chunk_id,

View file

@ -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]