Merge pull request #264 from ggozad/fix/tui

Fix TUI session handling

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

@@ -23,7 +23,6 @@ from haiku.rag.agents.chat.state import (
     ChatDeps,
     ChatSessionState,
 )
-from haiku.rag.agents.research.models import Citation
 from haiku.rag.client import HaikuRAG
 from haiku.rag.config import get_config
 
@@ -103,7 +102,6 @@ class ChatApp(App):
         self.session_state = ChatSessionState()
         self._is_processing = False
         self._tool_call_widgets: dict[str, Any] = {}
-        self._last_citations: list[Citation] = []
         self._current_worker: Worker[None] | None = None
         self._message_history: list[ModelMessage] = []
         self._document_filter: list[str] = []
@@ -170,6 +168,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)
@@ -199,8 +201,7 @@ class ChatApp(App):
                         snapshot = getattr(meta_event, "snapshot", {})
                         self._agui_state_snapshot = snapshot
                         chat_state = snapshot.get(AGUI_STATE_KEY, snapshot)
-                        citations = chat_state.get("citations", [])
+                        self._sync_session_state(chat_state)
-                        self._last_citations = [Citation(**c) for c in citations]
 
                     elif meta_event.type == EventType.STATE_DELTA:
                         delta = getattr(meta_event, "delta", [])
@@ -212,8 +213,7 @@ class ChatApp(App):
                         chat_state = self._agui_state_snapshot.get(
                             AGUI_STATE_KEY, self._agui_state_snapshot
                         )
-                        citations = chat_state.get("citations", [])
+                        self._sync_session_state(chat_state)
-                        self._last_citations = [Citation(**c) for c in citations]
 
     async def _event_stream_handler(
         self,
@@ -247,7 +247,7 @@ class ChatApp(App):
 
         # Clear for new query
         self._tool_call_widgets.clear()
-        self._last_citations.clear()
+        self.session_state.citations.clear()
 
         # Run agent in a worker to keep UI responsive
         self._is_processing = True
@@ -303,8 +303,8 @@ class ChatApp(App):
                 self._message_history = stream.all_messages()
 
                 # Add citations captured from tool metadata
-                if self._last_citations:
+                if self.session_state.citations:
-                    await chat_history.add_citations(self._last_citations)
+                    await chat_history.add_citations(self.session_state.citations)
 
         except asyncio.CancelledError:
             chat_history.hide_thinking()
@@ -323,7 +323,6 @@ class ChatApp(App):
         """Clear the chat history and reset session."""
         chat_history = self.query_one(ChatHistory)
         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)

tests/chat/test_chat_app.py

@@ -337,9 +337,10 @@ async def test_handle_stream_event_extracts_citations_from_state_snapshot(
             # Handle the event
             await app._handle_stream_event(event)
 
-            # Citations should be extracted
+            # Session state should be synced
-            assert len(app._last_citations) == 1
+            assert len(app.session_state.citations) == 1
-            assert app._last_citations[0].chunk_id == "chunk1"
+            assert app.session_state.citations[0].chunk_id == "chunk1"
+            assert app.session_state.citation_registry == {"chunk1": 1}
 
 
 @pytest.mark.asyncio
@@ -387,7 +388,7 @@ async def test_handle_stream_event_extracts_citations_from_state_delta(
             )
             event1 = FunctionToolResultEvent(result=tool_return1)
             await app._handle_stream_event(event1)
-            assert len(app._last_citations) == 0
+            assert len(app.session_state.citations) == 0
 
             # Now handle a STATE_DELTA event that adds citations
             delta_event = StateDeltaEvent(
@@ -425,10 +426,10 @@ async def test_handle_stream_event_extracts_citations_from_state_delta(
             # Handle the delta event
             await app._handle_stream_event(event2)
 
-            # Citations should be extracted from the delta
+            # Session state should be synced
-            assert len(app._last_citations) == 1
+            assert len(app.session_state.citations) == 1
-            assert app._last_citations[0].chunk_id == "chunk1"
+            assert app.session_state.citations[0].chunk_id == "chunk1"
-            assert app._last_citations[0].content == "Test content from delta"
+            assert app.session_state.citations[0].content == "Test content from delta"
 
 
 @pytest.mark.asyncio
@@ -505,10 +506,76 @@ async def test_handle_stream_event_delta_with_preinitialized_state(
             # Handle the delta event
             await app._handle_stream_event(event)
 
-            # Citations should be extracted from the delta
+            # Session state should be synced
-            assert len(app._last_citations) == 1
+            assert len(app.session_state.citations) == 1
-            assert app._last_citations[0].chunk_id == "chunk1"
+            assert app.session_state.citations[0].chunk_id == "chunk1"
-            assert app._last_citations[0].content == "Content from first delta"
+            assert app.session_state.citations[0].content == "Content from first delta"
+
+
+@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