Handle snapshot deltas in TUI

haiku_rag_slim/haiku/rag/chat/app.py

@@ -6,6 +6,7 @@ from datetime import datetime
 from pathlib import Path
 from typing import TYPE_CHECKING, Any
 
+import jsonpatch
 from ag_ui.core import EventType
 from pydantic_ai import (
     Agent,
@@ -106,6 +107,7 @@ class ChatApp(App):
         self._current_worker: Worker[None] | None = None
         self._message_history: list[ModelMessage] = []
         self._document_filter: list[str] = []
+        self._agui_state_snapshot: dict[str, Any] = {}
 
     def compose(self) -> "ComposeResult":
         """Compose the UI layout."""
@@ -185,20 +187,33 @@ class ChatApp(App):
                 widget = self._tool_call_widgets[tool_call_id]
                 chat_history.mark_tool_complete(widget)
 
-            # Extract citations from StateSnapshotEvent in tool metadata
+            # Extract citations from state events in tool metadata
             result = getattr(event, "result", None)
             metadata = getattr(result, "metadata", None) if result else None
             if metadata:
                 for meta_event in metadata:
-                    if (
+                    if not hasattr(meta_event, "type"):
-                        hasattr(meta_event, "type")
+                        continue
-                        and meta_event.type == EventType.STATE_SNAPSHOT
+
-                    ):
+                    if meta_event.type == EventType.STATE_SNAPSHOT:
                         snapshot = getattr(meta_event, "snapshot", {})
+                        self._agui_state_snapshot = snapshot
                         chat_state = snapshot.get(AGUI_STATE_KEY, snapshot)
-                        self._last_citations = [
+                        citations = chat_state.get("citations", [])
-                            Citation(**c) for c in chat_state["citations"]
+                        self._last_citations = [Citation(**c) for c in citations]
-                        ]
+
+                    elif meta_event.type == EventType.STATE_DELTA:
+                        delta = getattr(meta_event, "delta", [])
+                        if delta:
+                            patch = jsonpatch.JsonPatch(delta)
+                            self._agui_state_snapshot = patch.apply(
+                                self._agui_state_snapshot
+                            )
+                        chat_state = self._agui_state_snapshot.get(
+                            AGUI_STATE_KEY, self._agui_state_snapshot
+                        )
+                        citations = chat_state.get("citations", [])
+                        self._last_citations = [Citation(**c) for c in citations]
 
     async def _event_stream_handler(
         self,
@@ -252,6 +267,12 @@ class ChatApp(App):
         await chat_history.show_thinking()
 
         try:
+            # Initialize AGUI state snapshot from session state for delta application
+            if self.session_state:
+                self._agui_state_snapshot = {
+                    AGUI_STATE_KEY: self.session_state.model_dump(mode="json")
+                }
+
             deps = ChatDeps(
                 client=self.client,
                 config=self.config,
@@ -304,6 +325,7 @@ class ChatApp(App):
         await chat_history.clear_messages()
         self._last_citations.clear()
         self._message_history.clear()
+        self._agui_state_snapshot = {}
         # Reset context lock and session state (reset to CLI value)
         self._context_locked = False
         self.session_state = ChatSessionState(

tests/chat/test_chat_app.py

@@ -282,6 +282,235 @@ async def test_clear_chat_resets_session(temp_db_path: Path):
             assert app.session_state.session_id != original_session_id
 
 
+@pytest.mark.asyncio
+async def test_handle_stream_event_extracts_citations_from_state_snapshot(
+    temp_db_path: Path,
+):
+    """Test that _handle_stream_event extracts citations from STATE_SNAPSHOT events."""
+    from ag_ui.core import EventType, StateSnapshotEvent
+    from pydantic_ai import FunctionToolResultEvent
+    from pydantic_ai.messages import ToolReturnPart
+
+    from haiku.rag.agents.chat.state import AGUI_STATE_KEY
+    from haiku.rag.chat.app import ChatApp
+
+    mock_client = AsyncMock()
+    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+    mock_client.__aexit__ = AsyncMock(return_value=None)
+
+    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
+        app = ChatApp(temp_db_path, read_only=True)
+
+        async with app.run_test():
+            # Create a STATE_SNAPSHOT event with citations
+            snapshot_event = StateSnapshotEvent(
+                type=EventType.STATE_SNAPSHOT,
+                snapshot={
+                    AGUI_STATE_KEY: {
+                        "session_id": "test",
+                        "citations": [
+                            {
+                                "index": 1,
+                                "document_id": "doc1",
+                                "chunk_id": "chunk1",
+                                "document_uri": "test.pdf",
+                                "document_title": "Test Doc",
+                                "content": "Test content",
+                            }
+                        ],
+                        "qa_history": [],
+                        "citation_registry": {"chunk1": 1},
+                    }
+                },
+            )
+
+            # Create a tool result with the snapshot in metadata
+            tool_return = ToolReturnPart(
+                tool_name="search",
+                content="Found results",
+                tool_call_id="test-call-1",
+                metadata=[snapshot_event],
+            )
+
+            event = FunctionToolResultEvent(result=tool_return)
+
+            # Handle the event
+            await app._handle_stream_event(event)
+
+            # Citations should be extracted
+            assert len(app._last_citations) == 1
+            assert app._last_citations[0].chunk_id == "chunk1"
+
+
+@pytest.mark.asyncio
+async def test_handle_stream_event_extracts_citations_from_state_delta(
+    temp_db_path: Path,
+):
+    """Test that _handle_stream_event extracts citations from STATE_DELTA events.
+
+    This test demonstrates that state deltas need to be applied to extract citations.
+    After the first STATE_SNAPSHOT, subsequent tool calls emit STATE_DELTA events
+    containing JSON Patch operations.
+    """
+    from ag_ui.core import EventType, StateDeltaEvent, StateSnapshotEvent
+    from pydantic_ai import FunctionToolResultEvent
+    from pydantic_ai.messages import ToolReturnPart
+
+    from haiku.rag.agents.chat.state import AGUI_STATE_KEY
+    from haiku.rag.chat.app import ChatApp
+
+    mock_client = AsyncMock()
+    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+    mock_client.__aexit__ = AsyncMock(return_value=None)
+
+    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
+        app = ChatApp(temp_db_path, read_only=True)
+
+        async with app.run_test():
+            # First, handle a STATE_SNAPSHOT to establish initial state
+            initial_snapshot = StateSnapshotEvent(
+                type=EventType.STATE_SNAPSHOT,
+                snapshot={
+                    AGUI_STATE_KEY: {
+                        "session_id": "test",
+                        "citations": [],
+                        "qa_history": [],
+                        "citation_registry": {},
+                    }
+                },
+            )
+            tool_return1 = ToolReturnPart(
+                tool_name="search",
+                content="Initial search",
+                tool_call_id="test-call-1",
+                metadata=[initial_snapshot],
+            )
+            event1 = FunctionToolResultEvent(result=tool_return1)
+            await app._handle_stream_event(event1)
+            assert len(app._last_citations) == 0
+
+            # Now handle a STATE_DELTA event that adds citations
+            delta_event = StateDeltaEvent(
+                type=EventType.STATE_DELTA,
+                delta=[
+                    {
+                        "op": "replace",
+                        "path": f"/{AGUI_STATE_KEY}/citations",
+                        "value": [
+                            {
+                                "index": 1,
+                                "document_id": "doc1",
+                                "chunk_id": "chunk1",
+                                "document_uri": "test.pdf",
+                                "document_title": "Test Doc",
+                                "content": "Test content from delta",
+                            }
+                        ],
+                    },
+                    {
+                        "op": "add",
+                        "path": f"/{AGUI_STATE_KEY}/citation_registry/chunk1",
+                        "value": 1,
+                    },
+                ],
+            )
+            tool_return2 = ToolReturnPart(
+                tool_name="ask",
+                content="Answer with citations",
+                tool_call_id="test-call-2",
+                metadata=[delta_event],
+            )
+            event2 = FunctionToolResultEvent(result=tool_return2)
+
+            # Handle the delta event
+            await app._handle_stream_event(event2)
+
+            # Citations should be extracted from the delta
+            assert len(app._last_citations) == 1
+            assert app._last_citations[0].chunk_id == "chunk1"
+            assert app._last_citations[0].content == "Test content from delta"
+
+
+@pytest.mark.asyncio
+async def test_handle_stream_event_delta_with_preinitialized_state(
+    temp_db_path: Path,
+):
+    """Test delta handling when _agui_state_snapshot is pre-initialized.
+
+    This is the actual TUI scenario: session_state exists from the start,
+    so the agent emits deltas (not snapshots) even on the first tool call.
+    The TUI pre-initializes _agui_state_snapshot from session_state.
+    """
+    from ag_ui.core import EventType, StateDeltaEvent
+    from pydantic_ai import FunctionToolResultEvent
+    from pydantic_ai.messages import ToolReturnPart
+
+    from haiku.rag.agents.chat.state import AGUI_STATE_KEY
+    from haiku.rag.chat.app import ChatApp
+
+    mock_client = AsyncMock()
+    mock_client.__aenter__ = AsyncMock(return_value=mock_client)
+    mock_client.__aexit__ = AsyncMock(return_value=None)
+
+    with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client):
+        app = ChatApp(temp_db_path, read_only=True)
+
+        async with app.run_test():
+            # Pre-initialize _agui_state_snapshot (simulating what _run_agent does)
+            app._agui_state_snapshot = {
+                AGUI_STATE_KEY: {
+                    "session_id": "test",
+                    "citations": [],
+                    "qa_history": [],
+                    "citation_registry": {},
+                    "document_filter": [],
+                    "initial_context": None,
+                    "session_context": None,
+                }
+            }
+
+            # Now handle a STATE_DELTA event directly (no prior snapshot event)
+            delta_event = StateDeltaEvent(
+                type=EventType.STATE_DELTA,
+                delta=[
+                    {
+                        "op": "replace",
+                        "path": f"/{AGUI_STATE_KEY}/citations",
+                        "value": [
+                            {
+                                "index": 1,
+                                "document_id": "doc1",
+                                "chunk_id": "chunk1",
+                                "document_uri": "test.pdf",
+                                "document_title": "Test Doc",
+                                "content": "Content from first delta",
+                            }
+                        ],
+                    },
+                    {
+                        "op": "add",
+                        "path": f"/{AGUI_STATE_KEY}/citation_registry/chunk1",
+                        "value": 1,
+                    },
+                ],
+            )
+            tool_return = ToolReturnPart(
+                tool_name="ask",
+                content="Answer with citations",
+                tool_call_id="test-call-1",
+                metadata=[delta_event],
+            )
+            event = FunctionToolResultEvent(result=tool_return)
+
+            # Handle the delta event
+            await app._handle_stream_event(event)
+
+            # Citations should be extracted from the delta
+            assert len(app._last_citations) == 1
+            assert app._last_citations[0].chunk_id == "chunk1"
+            assert app._last_citations[0].content == "Content from first delta"
+
+
 @pytest.mark.asyncio
 async def test_citation_expand_collapse_with_enter(temp_db_path: Path):
     """Test that pressing Enter on a focused citation toggles expand/collapse."""