From d1e3b828f1bbc4508379383ff5f826942e5d713e Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 13 Aug 2026 15:57:33 +0300 Subject: [PATCH] Judge each capability's evidence by its own record The guard accepted one carried record as covering the whole request, so a host retaining only the RAG namespace had its earlier analysis evidence replaced by receipts and lost, with the RAG record making the loss look accounted for. Each capability is now judged on its own, and only when its own evidence is at stake: another capability's record says nothing about this one's, and requiring every record would stop a host that registers both capabilities and only ever uses one. --- .../haiku/rag/capabilities/compaction.py | 40 ++++++--- tests/capabilities/test_evidence_wire.py | 88 +++++++++++++++++++ 2 files changed, 118 insertions(+), 10 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index 6b1decbb..05b57667 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -264,6 +264,32 @@ def compact_history( return compacted +def _require_a_record_of_what_was_cited( + evidence: Sequence[DiscoveredEvidence], + messages: list[ModelMessage], + boundary: int, +) -> None: + """Refuse to compact a capability's evidence when its record was not carried. + + Judged per capability, and only for one whose own evidence is actually at + stake: another capability's carried record says nothing about this one's, and a + capability the model never used has nothing to lose. Without the record there + is no capsule to put in the evidence's place, so compacting would drop it and + leave the citations the host already displayed as the only trace. + """ + for found in evidence: + if found.state_carried: + continue + if _newest_owned_return(messages, boundary, found.tool_names) is None: + continue + raise RuntimeError( + f"Evidence compaction found {found.capability} 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, " + f"alongside the message history: {found.capability} state was missing." + ) + + def _newest_owned_return( messages: list[ModelMessage], boundary: int, owned_tools: frozenset[str] ) -> tuple[int, int] | None: @@ -322,16 +348,10 @@ 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: + _require_a_record_of_what_was_cited( + evidence, request_context.messages, boundary + ) if boundary > 0: await self._build_once(ctx, evidence) request_context.messages = compact_history( diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py index 7e588ca1..ef6a1cef 100644 --- a/tests/capabilities/test_evidence_wire.py +++ b/tests/capabilities/test_evidence_wire.py @@ -19,6 +19,7 @@ from pydantic_ai.models.function import FunctionModel from pydantic_ai.models.test import TestModel from pydantic_ai.usage import RunUsage +from haiku.rag.capabilities.analysis import create_capability as create_analysis from haiku.rag.capabilities.compaction import ( RECEIPT, Capsule, @@ -789,3 +790,90 @@ async def test_compaction_proceeds_for_a_host_that_carries_state(temp_db_path): await agent.run("a follow-up", deps=carried, message_history=history) assert returns_of(wire[-1]) == [RECEIPT] + + +@pytest.mark.asyncio +async def test_compaction_refuses_when_one_capability_of_two_lost_its_record( + temp_db_path, +): + """One carried record does not vouch for the other capability's evidence. + + A host retaining only the RAG namespace leaves the analysis record empty, and + its earlier evidence would be replaced by receipts retaining nothing while the + RAG record made the loss look accounted for. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + analysis = create_analysis( + db_path=temp_db_path, config=AppConfig(), defer_loading=False + ) + compactor = create_compaction() + + async def model(_messages, _info): # pragma: no cover - never reached + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent( + FunctionModel(model), + deps_type=Deps, + capabilities=[rag, analysis, compactor], + ) + history: list[Any] = [ + ModelRequest(parts=[UserPromptPart("an earlier question")]), + ModelResponse( + parts=[ToolCallPart("analysis_search", {"query": "q"}, "call-1")] + ), + ModelRequest( + parts=[ToolReturnPart("analysis_search", "ANALYSIS EVIDENCE", "call-1")] + ), + ModelResponse(parts=[TextPart("an answer")]), + ] + # Only the RAG namespace comes back, as a host whitelisting fields would send. + rag_only = Deps( + state={ + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=0, in_progress=False) + ).model_dump(mode="json") + } + ) + + with pytest.raises(RuntimeError, match="analysis"): + await agent.run("a follow-up", deps=rag_only, message_history=history) + + +@pytest.mark.asyncio +async def test_compaction_proceeds_when_the_capability_without_a_record_has_no_evidence( + temp_db_path, +): + """A capability the earlier question never used has nothing to lose. + + Refusing whenever any record is missing would stop a host that registers both + capabilities and only ever uses one, which is the composition the docs + recommend against but hosts still have. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + analysis = create_analysis( + db_path=temp_db_path, config=AppConfig(), defer_loading=False + ) + compactor = create_compaction() + 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, analysis, compactor], + ) + history = answered_question("an earlier question", evidence="RAG EVIDENCE") + rag_only = Deps( + state={ + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=0, in_progress=False) + ).model_dump(mode="json") + } + ) + + await agent.run("a follow-up", deps=rag_only, message_history=history) + + assert returns_of(wire[-1]) == [RECEIPT]