TUI now syncs full session state from AG-UI events

CHANGELOG.md

@@ -22,6 +22,9 @@
   - First request still sends full snapshot; subsequent requests send only changes
   - Backend logging shows incoming/outgoing state events for debugging
 
+### Fixed
+
+- **Chat TUI Session State Sync**: TUI now syncs full session state from AG-UI events
 
 ## [0.27.1] - 2026-01-27
 

haiku_rag_slim/haiku/rag/chat/app.py

@@ -170,6 +170,10 @@ class ChatApp(App):
         if self.client:
             await self.client.__aexit__(None, None, None)
 
+    def _sync_session_state(self, chat_state: dict[str, Any]) -> None:
+        """Sync session_state from AG-UI state."""
+        self.session_state = ChatSessionState.model_validate(chat_state)
+
     async def _handle_stream_event(self, event: AgentStreamEvent) -> None:
         """Handle streaming events from the agent."""
         chat_history = self.query_one(ChatHistory)
@@ -201,6 +205,7 @@ class ChatApp(App):
                         chat_state = snapshot.get(AGUI_STATE_KEY, snapshot)
                         citations = chat_state.get("citations", [])
                         self._last_citations = [Citation(**c) for c in citations]
+                        self._sync_session_state(chat_state)
 
                     elif meta_event.type == EventType.STATE_DELTA:
                         delta = getattr(meta_event, "delta", [])
@@ -214,6 +219,7 @@ class ChatApp(App):
                         )
                         citations = chat_state.get("citations", [])
                         self._last_citations = [Citation(**c) for c in citations]
+                        self._sync_session_state(chat_state)
 
     async def _event_stream_handler(
         self,

tests/chat/test_chat_app.py

@@ -341,6 +341,11 @@ async def test_handle_stream_event_extracts_citations_from_state_snapshot(
             assert len(app._last_citations) == 1
             assert app._last_citations[0].chunk_id == "chunk1"
 
+            # Session state should be synced
+            assert len(app.session_state.citations) == 1
+            assert app.session_state.citations[0].chunk_id == "chunk1"
+            assert app.session_state.citation_registry == {"chunk1": 1}
+
 
 @pytest.mark.asyncio
 async def test_handle_stream_event_extracts_citations_from_state_delta(
@@ -430,6 +435,10 @@ async def test_handle_stream_event_extracts_citations_from_state_delta(
             assert app._last_citations[0].chunk_id == "chunk1"
             assert app._last_citations[0].content == "Test content from delta"
 
+            # Session state should be synced
+            assert len(app.session_state.citations) == 1
+            assert app.session_state.citations[0].chunk_id == "chunk1"
+
 
 @pytest.mark.asyncio
 async def test_handle_stream_event_delta_with_preinitialized_state(
@@ -510,6 +519,76 @@ async def test_handle_stream_event_delta_with_preinitialized_state(
             assert app._last_citations[0].chunk_id == "chunk1"
             assert app._last_citations[0].content == "Content from first delta"
 
+            # Session state should be synced
+            assert len(app.session_state.citations) == 1
+            assert app.session_state.citations[0].chunk_id == "chunk1"
+
+
+@pytest.mark.asyncio
+async def test_handle_stream_event_syncs_session_context(temp_db_path: Path):
+    """Test that _handle_stream_event syncs session_context to 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 state
+            app._agui_state_snapshot = {
+                AGUI_STATE_KEY: {
+                    "session_id": "test",
+                    "citations": [],
+                    "qa_history": [],
+                    "citation_registry": {},
+                    "document_filter": [],
+                    "initial_context": None,
+                    "session_context": None,
+                }
+            }
+
+            # Verify session_context starts as None
+            assert app.session_state.session_context is None
+
+            # Handle a delta that adds session_context
+            delta_event = StateDeltaEvent(
+                type=EventType.STATE_DELTA,
+                delta=[
+                    {
+                        "op": "replace",
+                        "path": f"/{AGUI_STATE_KEY}/session_context",
+                        "value": {
+                            "summary": "User asked about Python async patterns.",
+                            "last_updated": "2025-01-15T10:30:00",
+                        },
+                    },
+                ],
+            )
+            tool_return = ToolReturnPart(
+                tool_name="ask",
+                content="Answer about async",
+                tool_call_id="test-call-1",
+                metadata=[delta_event],
+            )
+            event = FunctionToolResultEvent(result=tool_return)
+
+            await app._handle_stream_event(event)
+
+            # Session context should be synced to session_state
+            assert app.session_state.session_context is not None
+            assert (
+                app.session_state.session_context.summary
+                == "User asked about Python async patterns."
+            )
+
 
 @pytest.mark.asyncio
 async def test_citation_expand_collapse_with_enter(temp_db_path: Path):