Fix state delta computation by avoiding mutation of original state

app/backend/main.py

@@ -80,30 +80,30 @@ async def stream_chat(request: Request) -> Response:
     run_input = AGUIAdapter.build_run_input(body)
 
     # Restore session state from incoming AG-UI state (look under namespaced key)
-    initial_qa_history: list[QAResponse] = []
-    session_id: str | None = None
-    document_filter: list[str] = []
-    initial_context: str | None = None
+    session_state: ChatSessionState | None = None
     state = getattr(run_input, "state", None)
-    if state:
-        chat_state = state.get(AGUI_STATE_KEY, state)
-        if "qa_history" in chat_state:
-            initial_qa_history = [
-                QAResponse(**qa) for qa in chat_state.get("qa_history", [])
-            ]
-        session_id = chat_state.get("session_id")
-        document_filter = chat_state.get("document_filter", [])
-        initial_context = chat_state.get("initial_context")
+    if state and AGUI_STATE_KEY in state:
+        chat_state = state[AGUI_STATE_KEY]
+        if chat_state and chat_state.get("session_id"):
+            # Only restore state if client has a session_id (not first request)
+            # This ensures first request gets a full snapshot with generated UUID
+            # NOTE: We intentionally do NOT restore session_context from the client.
+            # The server maintains session_context via background summarization tasks,
+            # and the agent fetches it from the server-side cache (get_cached_session_context).
+            session_state = ChatSessionState(
+                session_id=chat_state["session_id"],
+                qa_history=[
+                    QAResponse(**qa) for qa in chat_state.get("qa_history", [])
+                ],
+                document_filter=chat_state.get("document_filter", []),
+                initial_context=chat_state.get("initial_context"),
+                citation_registry=chat_state.get("citation_registry", {}),
+            )
 
     deps = ChatDeps(
         client=get_client(db_path),
         config=Config,
-        session_state=ChatSessionState(
-            qa_history=initial_qa_history,
-            document_filter=document_filter,
-            initial_context=initial_context,
-            **({"session_id": session_id} if session_id else {}),
-        ),
+        session_state=session_state,
         state_key=AGUI_STATE_KEY,
     )
 
@@ -111,25 +111,7 @@ async def stream_chat(request: Request) -> Response:
     adapter = AGUIAdapter(agent=chat_agent, run_input=run_input, accept=accept)
     event_stream = adapter.run_stream(deps=deps)
 
-    # Wrap to log state events
-    async def logged_event_stream():
-        async for event in event_stream:
-            event_type = getattr(event, "type", None)
-            if event_type and "state" in str(event_type).lower():
-                delta: list[dict[str, str]] | None = getattr(event, "delta", None)
-                snapshot: dict[str, object] | None = getattr(event, "snapshot", None)
-                if delta is not None:
-                    logger.info(f"StateDeltaEvent: {len(delta)} ops")
-                    for op in delta[:3]:  # Log first 3 ops
-                        logger.info(f"  {op['op']} {op['path']}")
-                    if len(delta) > 3:
-                        logger.info(f"  ... and {len(delta) - 3} more ops")
-                elif snapshot is not None:
-                    snapshot_keys = list(snapshot.keys())
-                    logger.info(f"StateSnapshotEvent: keys={snapshot_keys}")
-            yield event
-
-    sse_event_stream = adapter.encode_stream(logged_event_stream())
+    sse_event_stream = adapter.encode_stream(event_stream)
 
     return StreamingResponse(
         sse_event_stream,

app/frontend/components/Chat.tsx

@@ -47,6 +47,7 @@ interface ChatSessionState {
 	qa_history: QAResponse[];
 	session_context: SessionContext | null;
 	document_filter: string[];
+	citation_registry: Record<string, number>;
 }
 
 // AG-UI state is namespaced under AGUI_STATE_KEY
@@ -404,6 +405,7 @@ function ChatContentInner() {
 					qa_history: [],
 					session_context: null,
 					document_filter: [],
+					citation_registry: {},
 				},
 			},
 		},
@@ -429,6 +431,8 @@ function ChatContentInner() {
 				qa_history: agentState?.[AGUI_STATE_KEY]?.qa_history ?? [],
 				session_context: agentState?.[AGUI_STATE_KEY]?.session_context ?? null,
 				document_filter: selected,
+				citation_registry:
+					agentState?.[AGUI_STATE_KEY]?.citation_registry ?? {},
 			},
 		});
 	};
@@ -445,6 +449,8 @@ function ChatContentInner() {
 				qa_history: agentState?.[AGUI_STATE_KEY]?.qa_history ?? [],
 				session_context: agentState?.[AGUI_STATE_KEY]?.session_context ?? null,
 				document_filter: agentState?.[AGUI_STATE_KEY]?.document_filter ?? [],
+				citation_registry:
+					agentState?.[AGUI_STATE_KEY]?.citation_registry ?? {},
 			},
 		});
 	};

haiku_rag_slim/haiku/rag/agents/chat/agent.py

@@ -116,12 +116,19 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
         if not results:
             return ToolReturn(return_value="No results found.")
 
-        # Build citation infos using stable registry indices
+        # Copy session state to work with (avoids mutating original for delta computation)
+        new_state = (
+            ctx.deps.session_state.model_copy(deep=True)
+            if ctx.deps.session_state
+            else ChatSessionState()
+        )
+
+        # Build citation infos using the copy's registry
         citation_infos = []
         for r in results:
             chunk_id = r.chunk_id or ""
-            if ctx.deps.session_state is not None and chunk_id:
-                index = ctx.deps.session_state.get_or_assign_index(chunk_id)
+            if chunk_id:
+                index = new_state.get_or_assign_index(chunk_id)
             else:
                 index = len(citation_infos) + 1
             citation_infos.append(
@@ -137,26 +144,10 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
                 )
             )
 
-        # Build new state with citations and registry
-        session_id = ctx.deps.session_state.session_id if ctx.deps.session_state else ""
-        new_state = ChatSessionState(
-            session_id=session_id,
-            citations=citation_infos,
-            qa_history=(
-                ctx.deps.session_state.qa_history if ctx.deps.session_state else []
-            ),
-            session_context=get_cached_session_context(session_id)
-            if session_id
-            else None,
-            document_filter=(
-                ctx.deps.session_state.document_filter if ctx.deps.session_state else []
-            ),
-            citation_registry=(
-                ctx.deps.session_state.citation_registry
-                if ctx.deps.session_state
-                else {}
-            ),
-        )
+        # Update new_state with citations and fresh session_context
+        new_state.citations = citation_infos
+        if new_state.session_id:
+            new_state.session_context = get_cached_session_context(new_state.session_id)
 
         # Return detailed results for the agent to present
         result_lines = []
@@ -276,14 +267,17 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
 
         result = await graph.run(state=state, deps=deps)
 
-        # Build citation infos using stable registry indices
+        # Copy session state to work with (avoids mutating original for delta computation)
+        new_state = (
+            ctx.deps.session_state.model_copy(deep=True)
+            if ctx.deps.session_state
+            else ChatSessionState()
+        )
+
+        # Build citation infos using the copy's registry
         citation_infos = []
         for c in result.citations:
-            # Use registry for stable indices across calls
-            if ctx.deps.session_state is not None:
-                index = ctx.deps.session_state.get_or_assign_index(c.chunk_id)
-            else:
-                index = len(citation_infos) + 1
+            index = new_state.get_or_assign_index(c.chunk_id)
             citation_infos.append(
                 Citation(
                     index=index,
@@ -297,54 +291,37 @@ def create_chat_agent(config: AppConfig) -> Agent[ChatDeps, str]:
                 )
             )
 
-        # Accumulate Q&A in session state with full citation metadata
-        if ctx.deps.session_state is not None:
-            qa_response = QAResponse(
-                question=question,
-                answer=result.answer,
-                confidence=result.confidence,
-                citations=citation_infos,
-            )
-            ctx.deps.session_state.qa_history.append(qa_response)
-            # Enforce FIFO limit
-            if len(ctx.deps.session_state.qa_history) > MAX_QA_HISTORY:
-                ctx.deps.session_state.qa_history = ctx.deps.session_state.qa_history[
-                    -MAX_QA_HISTORY:
-                ]
-
-            # Spawn background task to update session context
-            # Cancel any previous summarization for this session
-            if session_id in _summarization_tasks:
-                _summarization_tasks[session_id].cancel()
-
-            task = asyncio.create_task(
-                _update_context_background(
-                    qa_history=list(ctx.deps.session_state.qa_history),
-                    config=ctx.deps.config,
-                    session_state=ctx.deps.session_state,
-                )
-            )
-            _summarization_tasks[session_id] = task
-            task.add_done_callback(lambda t: _summarization_tasks.pop(session_id, None))
-
-        # Build new state with citations, qa_history, and registry
-        new_state = ChatSessionState(
-            session_id=session_id,
+        # Add Q&A to the copy's history
+        qa_response = QAResponse(
+            question=question,
+            answer=result.answer,
+            confidence=result.confidence,
             citations=citation_infos,
-            qa_history=(
-                ctx.deps.session_state.qa_history if ctx.deps.session_state else []
-            ),
-            session_context=get_cached_session_context(session_id)
-            if session_id
-            else None,
-            document_filter=(
-                ctx.deps.session_state.document_filter if ctx.deps.session_state else []
-            ),
-            citation_registry=(
-                ctx.deps.session_state.citation_registry
-                if ctx.deps.session_state
-                else {}
-            ),
+        )
+        new_state.qa_history.append(qa_response)
+        # Enforce FIFO limit
+        if len(new_state.qa_history) > MAX_QA_HISTORY:
+            new_state.qa_history = new_state.qa_history[-MAX_QA_HISTORY:]
+
+        # Update citations and session_context
+        new_state.citations = citation_infos
+        if new_state.session_id:
+            new_state.session_context = get_cached_session_context(new_state.session_id)
+
+        # Spawn background task to update session context
+        if new_state.session_id in _summarization_tasks:
+            _summarization_tasks[new_state.session_id].cancel()
+
+        task = asyncio.create_task(
+            _update_context_background(
+                qa_history=list(new_state.qa_history),
+                config=ctx.deps.config,
+                session_state=new_state,
+            )
+        )
+        _summarization_tasks[new_state.session_id] = task
+        task.add_done_callback(
+            lambda t, sid=new_state.session_id: _summarization_tasks.pop(sid, None)
         )
 
         # Format answer with citation references using stable indices