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..3c68a8ed 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/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._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", []) @@ -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._last_citations = [Citation(**c) for c in citations] + self._sync_session_state(chat_state) 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: - await chat_history.add_citations(self._last_citations) + if self.session_state.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) diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 0f148860..4313fba1 100644 --- a/tests/chat/test_chat_app.py +++ b/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 - 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 @@ -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 - assert len(app._last_citations) == 1 - 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" + 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 - assert len(app._last_citations) == 1 - 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" + 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