From f4735667df3609de0988871da78feebb7aaba273 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 17 Feb 2026 12:24:09 +0200 Subject: [PATCH] emit StateDeltaEvent with client-aware baseline --- CHANGELOG.md | 4 +++ haiku_rag_slim/haiku/rag/tools/context.py | 15 ++++++++++ haiku_rag_slim/haiku/rag/tools/qa.py | 36 +++++++++++++---------- tests/agents/chat/test_chat_agent.py | 9 ++---- tests/tools/test_context.py | 24 +++++++++++++++ 5 files changed, 67 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index eb65363f..2b68eb21 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ # Changelog ## [Unreleased] +### Changed + +- **AG-UI state sync**: `ask` tool now emits `StateDeltaEvent` (JSON Patch) instead of `StateSnapshotEvent`, consistent with the `search` tool + ## [0.30.0] - 2026-02-16 ### Added diff --git a/haiku_rag_slim/haiku/rag/tools/context.py b/haiku_rag_slim/haiku/rag/tools/context.py index 7f1dc0f8..4296e75d 100644 --- a/haiku_rag_slim/haiku/rag/tools/context.py +++ b/haiku_rag_slim/haiku/rag/tools/context.py @@ -62,6 +62,7 @@ class ToolContext(BaseModel): state_key: str | None = None _namespaces: dict[str, BaseModel] = PrivateAttr(default_factory=dict) + _client_snapshot: dict[str, Any] | None = PrivateAttr(default=None) def register(self, namespace: str, state: BaseModel) -> None: """Register state for a namespace. @@ -121,6 +122,16 @@ class ToolContext(BaseModel): """List all registered namespaces.""" return list(self._namespaces.keys()) + @property + def client_snapshot(self) -> dict[str, Any] | None: + """Snapshot captured after the last restore_state_snapshot call. + + Represents what the client has, before any server-side overrides. + Tools use this as the baseline for delta computation so that + server-side changes (e.g. background summarization) are included. + """ + return self._client_snapshot + def dump_namespaces(self) -> dict[str, dict[str, Any]]: """Serialize all namespace states to a dictionary. @@ -150,6 +161,9 @@ class ToolContext(BaseModel): validates them via the namespace model, and updates the state in place. Fields not present in *data* are left unchanged. + After restoring, captures a snapshot as ``client_snapshot`` so + tools can compute deltas against what the client actually has. + Args: data: Flat dict as produced by build_state_snapshot(). """ @@ -163,6 +177,7 @@ class ToolContext(BaseModel): updated = state.model_validate(current) for field_name in matching: setattr(state, field_name, getattr(updated, field_name)) + self._client_snapshot = self.build_state_snapshot() def load_namespace(self, namespace: str, state_type: type[T], data: dict) -> T: """Deserialize and register state for a namespace. diff --git a/haiku_rag_slim/haiku/rag/tools/qa.py b/haiku_rag_slim/haiku/rag/tools/qa.py index 0637932a..2f9d67c8 100644 --- a/haiku_rag_slim/haiku/rag/tools/qa.py +++ b/haiku_rag_slim/haiku/rag/tools/qa.py @@ -1,7 +1,6 @@ import math from collections.abc import Callable -from ag_ui.core import EventType, StateSnapshotEvent from pydantic import BaseModel, Field from pydantic_ai import FunctionToolset, RunContext, ToolReturn @@ -23,6 +22,7 @@ from haiku.rag.tools.session import ( SESSION_NAMESPACE, SessionContext, SessionState, + compute_combined_state_delta, ) PRIOR_ANSWER_RELEVANCE_THRESHOLD = 0.7 @@ -245,9 +245,14 @@ def create_qa_toolset( tool_context = ctx.deps.tool_context state_key: str | None = None + client_snapshot: dict | None = None if tool_context is not None: state_key = tool_context.state_key + if tool_context.namespaces: + client_snapshot = ( + tool_context.client_snapshot or tool_context.build_state_snapshot() + ) qa_result = await run_qa_core( client=client, @@ -259,21 +264,22 @@ def create_qa_toolset( on_qa_complete=on_ask_complete, ) - if tool_context is not None and tool_context.namespaces: - snapshot = tool_context.build_state_snapshot() - if state_key: - snapshot = {state_key: snapshot} - - answer_text = qa_result.answer - if qa_result.citations: - citation_refs = " ".join(f"[{c.index}]" for c in qa_result.citations) - answer_text = f"{answer_text}\n\nSources: {citation_refs}" - - state_event = StateSnapshotEvent( - type=EventType.STATE_SNAPSHOT, - snapshot=snapshot, + if client_snapshot is not None and tool_context is not None: + new_snapshot = tool_context.build_state_snapshot() + state_event = compute_combined_state_delta( + client_snapshot, + new_snapshot, + state_key=state_key, ) - return ToolReturn(return_value=answer_text, metadata=[state_event]) + + if state_event is not None: + answer_text = qa_result.answer + if qa_result.citations: + citation_refs = " ".join( + f"[{c.index}]" for c in qa_result.citations + ) + answer_text = f"{answer_text}\n\nSources: {citation_refs}" + return ToolReturn(return_value=answer_text, metadata=[state_event]) return qa_result diff --git a/tests/agents/chat/test_chat_agent.py b/tests/agents/chat/test_chat_agent.py index a1ec782f..3a36d80a 100644 --- a/tests/agents/chat/test_chat_agent.py +++ b/tests/agents/chat/test_chat_agent.py @@ -1,7 +1,7 @@ from pathlib import Path import pytest -from ag_ui.core import StateDeltaEvent, StateSnapshotEvent +from ag_ui.core import StateDeltaEvent from haiku.rag.agents.chat import ( AGUI_STATE_KEY, @@ -23,7 +23,7 @@ from haiku.rag.tools.session import SESSION_NAMESPACE, SessionContext, SessionSt def extract_state_from_result(result, state_key: str = AGUI_STATE_KEY) -> dict | None: """Extract emitted state from agent result's tool return metadata. - For deltas, applies the patch to an empty state to get the final state. + Applies the JSON Patch delta to an empty state to get the final state. """ import jsonpatch @@ -32,10 +32,7 @@ def extract_state_from_result(result, state_key: str = AGUI_STATE_KEY) -> dict | for part in message.parts: if hasattr(part, "metadata") and part.metadata: for meta in part.metadata: - if isinstance(meta, StateSnapshotEvent): - return meta.snapshot.get(state_key) - elif isinstance(meta, StateDeltaEvent): - # Apply delta to empty state to get final state + if isinstance(meta, StateDeltaEvent): empty_state = { state_key: ChatSessionState().model_dump(mode="json") } diff --git a/tests/tools/test_context.py b/tests/tools/test_context.py index 14ad4de1..28aa827a 100644 --- a/tests/tools/test_context.py +++ b/tests/tools/test_context.py @@ -444,6 +444,30 @@ def test_restore_state_snapshot_ignores_unknown_fields(): assert ns1.value == 10 +def test_restore_state_snapshot_captures_client_snapshot(): + """restore_state_snapshot stores the restored state as client_snapshot. + + This baseline is used by tools to compute deltas against what the + client actually has, so server-side changes (e.g. background + summarization) appear in the delta. + """ + ctx = ToolContext() + ctx.register("ns1", TestState(value=0)) + ctx.register("ns2", TestStateWithList(items=[])) + + assert ctx.client_snapshot is None + + ctx.restore_state_snapshot({"value": 10, "items": ["a"]}) + + assert ctx.client_snapshot == {"value": 10, "items": ["a"]} + + # Mutating state after restore doesn't affect the captured snapshot + ns1 = ctx.get("ns1", TestState) + assert ns1 is not None + ns1.value = 99 + assert ctx.client_snapshot == {"value": 10, "items": ["a"]} + + # --- prepare_context tests ---