Clena citation handling now that we handle the total state

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] = []
@@ -203,8 +201,6 @@ 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:
@@ -217,8 +213,6 @@ 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(
@@ -253,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
@@ -309,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()
@@ -329,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,10 +337,6 @@ 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"
@@ -392,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(
@@ -430,14 +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
@@ -514,14 +506,10 @@ 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