diff --git a/CHANGELOG.md b/CHANGELOG.md index 81e13659..e41b53eb 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 637f123f..98062865 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/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, diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 0f148860..db12f161 100644 --- a/tests/chat/test_chat_app.py +++ b/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):