emit StateDeltaEvent with client-aware baseline

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

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.

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
 

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")
                             }

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