Add citations_history to SessionState for unified citation rendering between search and ask

app/frontend/components/Chat.tsx

@@ -272,7 +272,7 @@ function MessageViewWithCitations({
 	return (
 		<CopilotChatMessageView messages={messages} isRunning={isRunning}>
 			{({ messageElements }) => {
-				if (!chatState?.qa_history?.length) {
+				if (!chatState?.citations_history?.length) {
 					return (
 						<>
 							{messageElements}
@@ -285,13 +285,11 @@ function MessageViewWithCitations({
 				// message (tool messages produce nothing). We correlate elements with
 				// messages to inject CitationBlocks after the right assistant responses.
 				//
-				// Tool call objects on messages only carry `id` (no `name`), so we
-				// can't identify which tool was called from the message alone. Instead
-				// we rely on the fact that qa_history only grows when the `ask` tool
-				// runs: after each assistant text response that followed tool calls,
-				// we inject any new qa_history citations.
+				// Both search and ask tools append to citations_history in order,
+				// so after each assistant text response that followed tool calls,
+				// we inject the next citations_history entry.
 				const result: React.ReactNode[] = [];
-				let qaIdx = 0;
+				let citIdx = 0;
 				let seenToolCalls = false;
 				let elemIdx = 0;
 
@@ -320,19 +318,19 @@ function MessageViewWithCitations({
 					}
 
 					// After an assistant text response that followed tool calls,
-					// inject the next qa_history entry's citations (one per turn)
+					// inject the next citations_history entry (one per turn)
 					if (msg.role === "assistant" && msg.content && seenToolCalls) {
-						if (qaIdx < chatState.qa_history.length) {
-							const qa = chatState.qa_history[qaIdx];
-							if (qa.citations?.length) {
+						if (citIdx < chatState.citations_history.length) {
+							const citations = chatState.citations_history[citIdx];
+							if (citations?.length) {
 								result.push(
 									<CitationBlock
-										key={`citations-${qaIdx}`}
-										citations={qa.citations}
+										key={`citations-${citIdx}`}
+										citations={citations}
 									/>,
 								);
 							}
-							qaIdx++;
+							citIdx++;
 						}
 						seenToolCalls = false;
 					}

app/frontend/lib/sessionStorage.ts

@@ -24,6 +24,7 @@ export interface SessionContext {
 export interface ChatSessionState {
 	initial_context: string | null;
 	citations: Citation[];
+	citations_history: Citation[][];
 	qa_history: QAResponse[];
 	session_context: SessionContext | null;
 	document_filter: string[];
@@ -53,6 +54,7 @@ export function normalizeChatState(state?: ChatSessionState): ChatSessionState {
 	return {
 		initial_context: state?.initial_context ?? null,
 		citations: state?.citations ?? [],
+		citations_history: state?.citations_history ?? [],
 		qa_history: state?.qa_history ?? [],
 		session_context: state?.session_context ?? null,
 		document_filter: state?.document_filter ?? [],

haiku_rag_slim/haiku/rag/tools/qa.py

@@ -185,6 +185,7 @@ async def run_qa_core(
 
     if session_state is not None:
         session_state.citations = citations
+        session_state.citations_history.append(citations)
 
     if qa_session_state is not None:
         qa_session_state.qa_history.append(

haiku_rag_slim/haiku/rag/tools/search.py

@@ -112,6 +112,7 @@ def create_search_toolset(
                     )
                 )
             session_state.citations = citations
+            session_state.citations_history.append(citations)
 
             result_lines = []
             for c in citations:

haiku_rag_slim/haiku/rag/tools/session.py

@@ -29,6 +29,7 @@ class SessionState(BaseModel):
     document_filter: list[str] = []
     citation_registry: dict[str, int] = {}
     citations: list[Citation] = []
+    citations_history: list[list[Citation]] = []
 
     def get_or_assign_index(self, chunk_id: str) -> int:
         """Get or assign a stable citation index for a chunk_id.

tests/tools/test_qa.py

@@ -76,6 +76,26 @@ class TestRunQACore:
         if result.citations:
             assert len(session_state.citation_registry) > 0
 
+    @pytest.mark.asyncio
+    async def test_run_qa_core_populates_citations_history(
+        self, allow_model_requests, qa_client, qa_config
+    ):
+        """run_qa_core appends to SessionState.citations_history."""
+        context = ToolContext()
+        prepare_context(context, features=["qa"])
+
+        await run_qa_core(
+            client=qa_client,
+            config=qa_config,
+            question="What is Python?",
+            context=context,
+        )
+
+        session_state = context.get(SESSION_NAMESPACE, SessionState)
+        assert session_state is not None
+        assert len(session_state.citations_history) == 1
+        assert session_state.citations_history[0] == session_state.citations
+
     @pytest.mark.asyncio
     async def test_run_qa_core_without_context(
         self, allow_model_requests, qa_client, qa_config

tests/tools/test_search.py

@@ -301,6 +301,44 @@ class TestSearchWithSessionState:
         # New chunks should get higher indices
         assert len(session_state.citation_registry) >= first_count
 
+    @pytest.mark.asyncio
+    async def test_search_populates_citations_history(
+        self, search_client, search_config
+    ):
+        """Search appends to SessionState.citations_history."""
+        context = ToolContext()
+        prepare_context(context, features=["search"])
+        toolset = create_search_toolset(search_config)
+
+        search_tool = toolset.tools["search"]
+        ctx = make_ctx(search_client, context)
+        await search_tool.function(ctx, "Python")
+
+        session_state = context.get(SESSION_NAMESPACE, SessionState)
+        assert session_state is not None
+        assert len(session_state.citations_history) == 1
+        assert session_state.citations_history[0] == session_state.citations
+
+    @pytest.mark.asyncio
+    async def test_search_multiple_appends_separate_entries(
+        self, search_client, search_config
+    ):
+        """Multiple searches append separate entries to citations_history."""
+        context = ToolContext()
+        prepare_context(context, features=["search"])
+        toolset = create_search_toolset(search_config)
+
+        search_tool = toolset.tools["search"]
+        ctx = make_ctx(search_client, context)
+        await search_tool.function(ctx, "Python")
+        await search_tool.function(ctx, "JavaScript")
+
+        session_state = context.get(SESSION_NAMESPACE, SessionState)
+        assert session_state is not None
+        assert len(session_state.citations_history) == 2
+        # Latest citations should match the last entry
+        assert session_state.citations_history[1] == session_state.citations
+
 
 @pytest.fixture
 def search_config():

tests/tools/test_session.py

@@ -1,5 +1,6 @@
 from ag_ui.core import EventType, StateDeltaEvent
 
+from haiku.rag.agents.research.models import Citation
 from haiku.rag.tools.session import (
     SessionState,
     compute_combined_state_delta,
@@ -7,6 +8,34 @@ from haiku.rag.tools.session import (
 )
 
 
+class TestSessionState:
+    """Tests for SessionState model."""
+
+    def test_citations_history_defaults_to_empty(self):
+        """SessionState.citations_history defaults to empty list."""
+        state = SessionState()
+        assert state.citations_history == []
+
+    def test_citations_history_serialization_roundtrip(self):
+        """citations_history survives serialize/deserialize."""
+        citation = Citation(
+            index=1,
+            document_id="d1",
+            chunk_id="c1",
+            document_uri="test://doc",
+            document_title="Doc",
+            page_numbers=[],
+            headings=None,
+            content="some content",
+        )
+        state = SessionState(citations_history=[[citation]])
+        data = state.model_dump(mode="json")
+        restored = SessionState.model_validate(data)
+        assert len(restored.citations_history) == 1
+        assert len(restored.citations_history[0]) == 1
+        assert restored.citations_history[0][0].chunk_id == "c1"
+
+
 class TestComputeStateDelta:
     """Tests for compute_state_delta."""