Send state deltas updates instead of full state snapshots

app/backend/main.py

@@ -110,7 +110,26 @@ async def stream_chat(request: Request) -> Response:
     # Use AGUIAdapter for streaming
     adapter = AGUIAdapter(agent=chat_agent, run_input=run_input, accept=accept)
     event_stream = adapter.run_stream(deps=deps)
-    sse_event_stream = adapter.encode_stream(event_stream)
+
+    # Wrap to log state events
+    async def logged_event_stream():
+        async for event in event_stream:
+            event_type = getattr(event, "type", None)
+            if event_type and "state" in str(event_type).lower():
+                delta: list[dict[str, str]] | None = getattr(event, "delta", None)
+                snapshot: dict[str, object] | None = getattr(event, "snapshot", None)
+                if delta is not None:
+                    logger.info(f"StateDeltaEvent: {len(delta)} ops")
+                    for op in delta[:3]:  # Log first 3 ops
+                        logger.info(f"  {op['op']} {op['path']}")
+                    if len(delta) > 3:
+                        logger.info(f"  ... and {len(delta) - 3} more ops")
+                elif snapshot is not None:
+                    snapshot_keys = list(snapshot.keys())
+                    logger.info(f"StateSnapshotEvent: keys={snapshot_keys}")
+            yield event
+
+    sse_event_stream = adapter.encode_stream(logged_event_stream())
 
     return StreamingResponse(
         sse_event_stream,

haiku_rag_slim/haiku/rag/agents/chat/agent.py

@@ -1,7 +1,6 @@
 import asyncio
 import math
 
-from ag_ui.core import EventType, StateSnapshotEvent
 from pydantic_ai import Agent, RunContext, ToolReturn
 
 from haiku.rag.agents.chat.context import (
@@ -20,6 +19,7 @@ from haiku.rag.agents.chat.state import (
     build_document_filter,
     build_multi_document_filter,
     combine_filters,
+    emit_state_event,
 )
 from haiku.rag.agents.research.dependencies import ResearchContext
 from haiku.rag.agents.research.graph import build_conversational_graph
@@ -173,19 +173,14 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             line += f"\n    {snippet}"
             result_lines.append(line)
 
-        snapshot = new_state.model_dump(mode="json")
-        if ctx.deps.state_key:
-            snapshot = {ctx.deps.state_key: snapshot}
+        state_event = emit_state_event(
+            ctx.deps.session_state, new_state, ctx.deps.state_key
+        )
 
         return ToolReturn(
             return_value=f"Found {len(results)} results:\n\n"
             + "\n\n".join(result_lines),
-            metadata=[
-                StateSnapshotEvent(
-                    type=EventType.STATE_SNAPSHOT,
-                    snapshot=snapshot,
-                )
-            ],
+            metadata=[state_event] if state_event else None,
         )
 
     @agent.tool
@@ -358,18 +353,13 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
             citation_refs = " ".join(f"[{c.index}]" for c in citation_infos)
             answer_text = f"{answer_text}\n\nSources: {citation_refs}"
 
-        snapshot = new_state.model_dump(mode="json")
-        if ctx.deps.state_key:
-            snapshot = {ctx.deps.state_key: snapshot}
+        state_event = emit_state_event(
+            ctx.deps.session_state, new_state, ctx.deps.state_key
+        )
 
         return ToolReturn(
             return_value=answer_text,
-            metadata=[
-                StateSnapshotEvent(
-                    type=EventType.STATE_SNAPSHOT,
-                    snapshot=snapshot,
-                )
-            ],
+            metadata=[state_event] if state_event else None,
         )
 
     @agent.tool

haiku_rag_slim/haiku/rag/agents/chat/state.py

@@ -3,6 +3,8 @@ from dataclasses import dataclass, field
 from datetime import datetime
 from typing import Any
 
+import jsonpatch
+from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent
 from pydantic import BaseModel, Field
 
 from haiku.rag.agents.research.models import Citation, SearchAnswer
@@ -194,3 +196,32 @@ def combine_filters(filter1: str | None, filter2: str | None) -> str | None:
     if len(filters) == 1:
         return filters[0]
     return f"({filters[0]}) AND ({filters[1]})"
+
+
+def emit_state_event(
+    current_state: ChatSessionState | None,
+    new_state: ChatSessionState,
+    state_key: str | None = None,
+) -> StateSnapshotEvent | StateDeltaEvent | None:
+    """Emit state delta against current state, or full snapshot if no current state."""
+    new_snapshot = new_state.model_dump(mode="json")
+    wrapped_new = {state_key: new_snapshot} if state_key else new_snapshot
+
+    if current_state is None:
+        return StateSnapshotEvent(
+            type=EventType.STATE_SNAPSHOT,
+            snapshot=wrapped_new,
+        )
+
+    current_snapshot = current_state.model_dump(mode="json")
+    wrapped_current = {state_key: current_snapshot} if state_key else current_snapshot
+
+    patch = jsonpatch.make_patch(wrapped_current, wrapped_new)
+
+    if not patch.patch:
+        return None
+
+    return StateDeltaEvent(
+        type=EventType.STATE_DELTA,
+        delta=patch.patch,
+    )

haiku_rag_slim/pyproject.toml

@@ -24,6 +24,7 @@ classifiers = [
 dependencies = [
     "docling-core==2.60.1",
     "httpx>=0.28.1",
+    "jsonpatch>=1.33",
     "lancedb==0.27.0",
     "pathspec>=1.0.3",
     "pydantic>=2.12.5",

tests/agents/chat/test_state.py

@@ -641,3 +641,118 @@ def test_chat_session_state_model_dump_json_serializes_datetime():
     # datetime should be serialized as ISO string, not datetime object
     assert isinstance(snapshot["session_context"]["last_updated"], str)
     assert snapshot["session_context"]["last_updated"] == "2025-01-27T12:00:00"
+
+
+def test_emit_state_event_returns_snapshot_when_no_current_state():
+    """emit_state_event returns StateSnapshotEvent when current_state is None."""
+    from ag_ui.core import EventType, StateSnapshotEvent
+
+    from haiku.rag.agents.chat.state import emit_state_event
+
+    new_state = ChatSessionState(session_id="test-123")
+
+    event = emit_state_event(None, new_state)
+
+    assert isinstance(event, StateSnapshotEvent)
+    assert event.type == EventType.STATE_SNAPSHOT
+    assert event.snapshot["session_id"] == "test-123"
+
+
+def test_emit_state_event_returns_snapshot_with_state_key():
+    """emit_state_event wraps snapshot in state_key namespace."""
+    from ag_ui.core import EventType, StateSnapshotEvent
+
+    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, emit_state_event
+
+    new_state = ChatSessionState(session_id="test-123")
+
+    event = emit_state_event(None, new_state, state_key=AGUI_STATE_KEY)
+
+    assert isinstance(event, StateSnapshotEvent)
+    assert event.type == EventType.STATE_SNAPSHOT
+    assert AGUI_STATE_KEY in event.snapshot
+    assert event.snapshot[AGUI_STATE_KEY]["session_id"] == "test-123"
+
+
+def test_emit_state_event_returns_none_when_no_changes():
+    """emit_state_event returns None when states are identical."""
+    from haiku.rag.agents.chat.state import emit_state_event
+
+    state = ChatSessionState(session_id="test-123", qa_history=[], citations=[])
+
+    event = emit_state_event(state, state)
+
+    assert event is None
+
+
+def test_emit_state_event_returns_delta_with_changes():
+    """emit_state_event returns StateDeltaEvent with JSON Patch ops for changes."""
+    from ag_ui.core import EventType, StateDeltaEvent
+
+    from haiku.rag.agents.chat.state import emit_state_event
+
+    current_state = ChatSessionState(session_id="test-123", qa_history=[], citations=[])
+    new_state = ChatSessionState(
+        session_id="test-123",
+        qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)],
+        citations=[],
+    )
+
+    event = emit_state_event(current_state, new_state)
+
+    assert isinstance(event, StateDeltaEvent)
+    assert event.type == EventType.STATE_DELTA
+    assert len(event.delta) > 0
+    # Delta should contain an "add" operation for the new qa_history entry
+    ops = event.delta
+    qa_history_op = next((op for op in ops if "/qa_history" in op["path"]), None)
+    assert qa_history_op is not None
+
+
+def test_emit_state_event_delta_with_state_key():
+    """emit_state_event wraps delta paths with state_key namespace."""
+    from ag_ui.core import StateDeltaEvent
+
+    from haiku.rag.agents.chat.state import AGUI_STATE_KEY, emit_state_event
+
+    current_state = ChatSessionState(session_id="test-123", qa_history=[])
+    new_state = ChatSessionState(
+        session_id="test-123",
+        qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)],
+    )
+
+    event = emit_state_event(current_state, new_state, state_key=AGUI_STATE_KEY)
+
+    assert isinstance(event, StateDeltaEvent)
+    # Paths should be namespaced under state_key
+    for op in event.delta:
+        assert op["path"].startswith(f"/{AGUI_STATE_KEY}")
+
+
+def test_emit_state_event_delta_produces_valid_patch():
+    """emit_state_event delta can be applied to reproduce new state."""
+    import jsonpatch
+
+    from haiku.rag.agents.chat.state import emit_state_event
+
+    current_state = ChatSessionState(
+        session_id="test-123",
+        qa_history=[QAResponse(question="Q1", answer="A1", confidence=0.9)],
+        citations=[],
+    )
+    new_state = ChatSessionState(
+        session_id="test-123",
+        qa_history=[
+            QAResponse(question="Q1", answer="A1", confidence=0.9),
+            QAResponse(question="Q2", answer="A2", confidence=0.8),
+        ],
+        citations=[],
+    )
+
+    event = emit_state_event(current_state, new_state)
+
+    # Apply patch to current state and verify it produces new state
+    current_snapshot = current_state.model_dump(mode="json")
+    patched = jsonpatch.apply_patch(current_snapshot, event.delta)
+    new_snapshot = new_state.model_dump(mode="json")
+    assert patched == new_snapshot

uv.lock

@@ -1368,6 +1368,7 @@ source = { editable = "haiku_rag_slim" }
 dependencies = [
     { name = "docling-core" },
     { name = "httpx" },
+    { name = "jsonpatch" },
     { name = "lancedb" },
     { name = "pathspec" },
     { name = "pydantic" },
@@ -1429,6 +1430,7 @@ requires-dist = [
     { name = "docling", marker = "extra == 'docling'", specifier = "==2.69.1" },
     { name = "docling-core", specifier = "==2.60.1" },
     { name = "httpx", specifier = ">=0.28.1" },
+    { name = "jsonpatch", specifier = ">=1.33" },
     { name = "lancedb", specifier = "==0.27.0" },
     { name = "mxbai-rerank", marker = "extra == 'mxbai'", specifier = ">=0.1.6" },
     { name = "opencv-python-headless", marker = "extra == 'docling'", specifier = ">=4.13.0.90" },