From 2e0ab87505c4fa8ba7be3efea3b4ddf6b6df78de Mon Sep 17 00:00:00 2001 From: juan <2930882+juacker@users.noreply.github.com> Date: Sat, 11 Jul 2026 22:54:51 +0200 Subject: [PATCH] Stop re-rendering the Workspace shell on every streaming delta MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Streaming text deltas arrive many times per second, and the accumulator lived inside SessionState — so each appendDelta gave sessions[sessionId] a new identity (immer structural sharing), re-rendering every whole-session subscriber per token chunk. The worst offender was the Workspace page (2,200+ lines of component body) which subscribed to the entire session object. Two changes, per external review feedback (validated against the code): 1. Move the accumulator to a top-level streamingText map in the SAME store (sessionId -> messageId -> text). Deltas now invalidate only streamingText subscribers, while completeMessage still swaps accumulator-for-final-message atomically in one set() — no flicker frame. appendDelta guards its isStreaming write so repeat deltas leave the session draft untouched (pinned by a new identity-stability test). 2. Narrow selectors in the Workspace shell: each rendered value gets its own subscription; the per-delta streamingText subscription moves down into ChatFirstLayout so token chunks re-render only the chat area. WorkspaceTaskTranscriptPanel and useAssistantSession get the same treatment. Stable EMPTY_* fallbacks keep memoized children from seeing fresh []/{} identities. A separate store was deliberately NOT used for the accumulator: it would split completeMessage's clear+install across two updates and open a one-frame flicker window. --- src/assistant/sessionStore.test.ts | 73 +++++++++++++++++-- src/assistant/sessionStore.ts | 49 +++++++++---- src/assistant/useAssistantEvents.test.ts | 2 +- src/assistant/useAssistantSession.ts | 9 ++- .../AskUserPanel/AskUserPanel.test.tsx | 1 - .../WorkspaceTaskTranscriptPanel.tsx | 4 +- src/pages/Workspace.tsx | 68 +++++++++++++---- 7 files changed, 167 insertions(+), 39 deletions(-) diff --git a/src/assistant/sessionStore.test.ts b/src/assistant/sessionStore.test.ts index 721cfbad..f0586f37 100644 --- a/src/assistant/sessionStore.test.ts +++ b/src/assistant/sessionStore.test.ts @@ -26,7 +26,12 @@ const ASK_REQUEST = { beforeEach(() => { // Zustand exposes setState on the hook itself; reset the slice that // every test below mutates. Avoids cross-test bleed. - useAssistantStore.setState({ sessions: {}, activeSessionByTab: {}, recoverablePrompts: {} }); + useAssistantStore.setState({ + sessions: {}, + streamingText: {}, + activeSessionByTab: {}, + recoverablePrompts: {}, + }); }); describe('initSession', () => { @@ -90,7 +95,7 @@ describe('loadSessionData — snapshot refresh preserves in-flight FE state', () // Regression: ask_user panel was being unmounted within ~5s because // Workspace.jsx polls workspace_get_snapshot every 5s and the wholesale // replacement in loadSessionData was wiping pendingAskUser. Same race - // existed (and is also fixed) for streamingTextByMessageId/isStreaming. + // existed (and is also fixed) for streaming text/isStreaming. it('preserves pendingAskUser across a snapshot refresh', () => { const store = useAssistantStore.getState(); @@ -102,15 +107,15 @@ describe('loadSessionData — snapshot refresh preserves in-flight FE state', () ); }); - it('preserves streamingTextByMessageId across a snapshot refresh', () => { + it('preserves streaming text across a snapshot refresh', () => { const store = useAssistantStore.getState(); store.initSession(SESSION); store.appendDelta(SESSION.id, 'msg-1', 'Hello '); store.appendDelta(SESSION.id, 'msg-1', 'world'); store.loadSessionData(SESSION.id, SESSION, [], [], []); - expect( - useAssistantStore.getState().sessions[SESSION.id]!.streamingTextByMessageId['msg-1'], - ).toBe('Hello world'); + expect(useAssistantStore.getState().streamingText[SESSION.id]!['msg-1']).toBe( + 'Hello world', + ); expect(useAssistantStore.getState().sessions[SESSION.id]!.isStreaming).toBe(true); }); @@ -226,7 +231,7 @@ describe('setRunStatus', () => { store.setRunStatus(SESSION.id, run('run-1', 'completed')); const s = useAssistantStore.getState().sessions[SESSION.id]!; expect(s.isStreaming).toBe(false); - expect(s.streamingTextByMessageId).toEqual({}); + expect(useAssistantStore.getState().streamingText[SESSION.id]).toBeUndefined(); }); it('sets streaming on queued/running/waiting_for_tool', () => { @@ -277,3 +282,57 @@ describe('removeMessage — recoverable prompt capture', () => { expect(useAssistantStore.getState().recoverablePrompts[SESSION.id]).toBeUndefined(); }); }); + +describe('appendDelta — session identity stability (streaming perf contract)', () => { + // The accumulator lives OUTSIDE SessionState precisely so that per-token + // deltas do not hand `sessions[sessionId]` a new identity, which would + // re-render every whole-session subscriber (the Workspace page shell) + // many times per second. These tests pin that contract. + + it('accumulates text in the top-level streamingText map', () => { + const store = useAssistantStore.getState(); + store.initSession(SESSION); + store.appendDelta(SESSION.id, 'msg-1', 'Hello '); + store.appendDelta(SESSION.id, 'msg-1', 'world'); + expect(useAssistantStore.getState().streamingText[SESSION.id]!['msg-1']).toBe( + 'Hello world', + ); + }); + + it('does not change the session object identity after the first delta', () => { + const store = useAssistantStore.getState(); + store.initSession(SESSION); + // First delta flips isStreaming (one legitimate session write). + store.appendDelta(SESSION.id, 'msg-1', 'a'); + const before = useAssistantStore.getState().sessions[SESSION.id]; + store.appendDelta(SESSION.id, 'msg-1', 'b'); + store.appendDelta(SESSION.id, 'msg-2', 'c'); + const after = useAssistantStore.getState().sessions[SESSION.id]; + expect(after).toBe(before); + }); + + it('is a no-op for an unknown session (no orphan accumulator)', () => { + useAssistantStore.getState().appendDelta('ghost', 'msg-1', 'x'); + expect(useAssistantStore.getState().streamingText['ghost']).toBeUndefined(); + }); + + it('completeMessage atomically clears the accumulator for that message', () => { + const store = useAssistantStore.getState(); + store.initSession(SESSION); + store.addMessage(SESSION.id, msg('msg-1')); + store.appendDelta(SESSION.id, 'msg-1', 'partial'); + store.appendDelta(SESSION.id, 'msg-2', 'other'); + store.completeMessage(SESSION.id, msg('msg-1')); + const streaming = useAssistantStore.getState().streamingText[SESSION.id]!; + expect(streaming['msg-1']).toBeUndefined(); + expect(streaming['msg-2']).toBe('other'); + }); + + it('removeSession drops the session accumulator', () => { + const store = useAssistantStore.getState(); + store.initSession(SESSION); + store.appendDelta(SESSION.id, 'msg-1', 'x'); + store.removeSession(SESSION.id); + expect(useAssistantStore.getState().streamingText[SESSION.id]).toBeUndefined(); + }); +}); diff --git a/src/assistant/sessionStore.ts b/src/assistant/sessionStore.ts index e2eca319..71d03a33 100644 --- a/src/assistant/sessionStore.ts +++ b/src/assistant/sessionStore.ts @@ -35,8 +35,6 @@ export interface SessionState { messages: AssistantMessage[]; runs: AssistantRun[]; toolCalls: ToolInvocation[]; - /** Per-message accumulator for streaming text deltas. Keyed by message id. */ - streamingTextByMessageId: Record; /** True while a run is queued/running/waiting_for_tool. Drives the chat activity indicator. */ isStreaming: boolean; /** Epoch ms (client clock) when the current run started running, for the @@ -61,6 +59,22 @@ export interface SessionState { export interface AssistantStoreState { sessions: Record; + /** + * Per-message accumulator for streaming text deltas, keyed by session id + * then message id. + * + * Kept OUTSIDE `SessionState` on purpose. Deltas arrive many times per + * second while an agent streams, and writing them inside the session draft + * gives `sessions[sessionId]` a new identity on every delta (immer + * structural sharing) — re-rendering every whole-session subscriber (the + * Workspace page shell) once per token chunk. As a top-level sibling map, + * a delta only invalidates `streamingText` subscribers, while actions that + * must touch both (e.g. `completeMessage` swapping the accumulator for the + * final message) still update them atomically inside a single `set()`, so + * there is no flicker frame between "streamed text gone" and "final + * message present". + */ + streamingText: Record>; activeSessionByTab: Record; setActiveSessionForTab: (tabId: string, sessionId: string) => void; @@ -112,7 +126,6 @@ const createInitialSessionState = (session: AssistantSession): SessionState => ( messages: [], runs: [], toolCalls: [], - streamingTextByMessageId: {}, isStreaming: false, runStartedAt: null, pendingAskUser: null, @@ -130,6 +143,7 @@ const useAssistantStore = create()( devtools( immer((set, get) => ({ sessions: {}, + streamingText: {}, activeSessionByTab: {}, recoverablePrompts: {}, @@ -195,7 +209,8 @@ const useAssistantStore = create()( s.totalMessageCount = Math.max(0, s.totalMessageCount - 1); } s.queuedMessageIds = s.queuedMessageIds.filter((id) => id !== messageId); - delete s.streamingTextByMessageId[messageId]; + const acc = state.streamingText[sessionId]; + if (acc) delete acc[messageId]; }), clearRecoverablePrompt: (sessionId) => @@ -257,11 +272,14 @@ const useAssistantStore = create()( appendDelta: (sessionId, messageId, text) => set((state) => { const s = state.sessions[sessionId]; - if (s) { - s.streamingTextByMessageId[messageId] = - (s.streamingTextByMessageId[messageId] || '') + text; - s.isStreaming = true; - } + if (!s) return; + const acc = (state.streamingText[sessionId] ??= {}); + acc[messageId] = (acc[messageId] || '') + text; + // Guarded write: assigning `true` unconditionally would mark the + // session draft modified on every delta, changing the identity of + // `sessions[sessionId]` per token chunk — exactly what keeping the + // accumulator outside the session object is meant to prevent. + if (!s.isStreaming) s.isStreaming = true; }), completeMessage: (sessionId, message) => @@ -272,7 +290,8 @@ const useAssistantStore = create()( if (idx >= 0) { s.messages[idx] = message; } - delete s.streamingTextByMessageId[message.id]; + const acc = state.streamingText[sessionId]; + if (acc) delete acc[message.id]; // Intentionally do NOT clear `isStreaming` here. A run typically // alternates between assistant text turns and tool-execution // phases; clearing on every message completion makes the activity @@ -299,7 +318,8 @@ const useAssistantStore = create()( if (idx >= 0) { s.messages[idx] = message; } - delete s.streamingTextByMessageId[message.id]; + const acc = state.streamingText[sessionId]; + if (acc) delete acc[message.id]; }), setRunStatus: (sessionId, run) => @@ -314,7 +334,7 @@ const useAssistantStore = create()( } if ((TERMINAL_STATUSES as readonly string[]).includes(run.status)) { s.isStreaming = false; - s.streamingTextByMessageId = {}; + delete state.streamingText[sessionId]; s.runStartedAt = null; s.pendingAskUser = null; } else if ((ACTIVE_STATUSES as readonly string[]).includes(run.status)) { @@ -403,7 +423,9 @@ const useAssistantStore = create()( // The DB only persists assistant text at end-of-run, so a poll // tick that lands mid-stream would otherwise wipe the deltas the // user is watching arrive, making text flicker on and off. - streamingTextByMessageId: existing?.streamingTextByMessageId || {}, + // (The streaming-text accumulator itself lives in the top-level + // `streamingText` map, so a session replacement here cannot wipe + // it by construction.) isStreaming: existing?.isStreaming || false, // Run start is FE-only live state the BE snapshot doesn't carry; // preserve it so a poll tick mid-run doesn't reset the elapsed timer. @@ -420,6 +442,7 @@ const useAssistantStore = create()( removeSession: (sessionId) => set((state) => { delete state.sessions[sessionId]; + delete state.streamingText[sessionId]; for (const [tabId, sid] of Object.entries(state.activeSessionByTab)) { if (sid === sessionId) { delete state.activeSessionByTab[tabId]; diff --git a/src/assistant/useAssistantEvents.test.ts b/src/assistant/useAssistantEvents.test.ts index 025dafe0..ae57512c 100644 --- a/src/assistant/useAssistantEvents.test.ts +++ b/src/assistant/useAssistantEvents.test.ts @@ -123,7 +123,7 @@ describe('useAssistantEvents — run lifecycle', () => { payload: { message_id: 'msg-1', text: 'world' }, }); expect( - useAssistantStore.getState().sessions[SESSION.id]!.streamingTextByMessageId['msg-1'], + useAssistantStore.getState().streamingText[SESSION.id]!['msg-1'], ).toBe('Hello world'); }); }); diff --git a/src/assistant/useAssistantSession.ts b/src/assistant/useAssistantSession.ts index 61b2e345..d86baa7d 100644 --- a/src/assistant/useAssistantSession.ts +++ b/src/assistant/useAssistantSession.ts @@ -17,6 +17,8 @@ interface EnsureSessionContext { const normalizeIdList = (ids: string[] | undefined | null): string[] => [...(ids || [])].sort(); +const EMPTY_STREAMING: Record = {}; + export function useAssistantSession(tabId: string) { const sessionId = useAssistantStore( (state) => state.activeSessionByTab[tabId] @@ -24,6 +26,11 @@ export function useAssistantSession(tabId: string) { const sessionState = useAssistantStore((state) => sessionId ? state.sessions[sessionId] : null ); + // Separate subscription: the accumulator lives outside SessionState so + // per-delta updates don't invalidate the whole-session subscription above. + const streamingText = useAssistantStore( + (state) => (sessionId && state.streamingText[sessionId]) || EMPTY_STREAMING + ); // Mirror sessionId into a ref so stable callbacks (sendMessage in particular) // always see the latest value. Writing `ref.current` directly during render @@ -140,7 +147,7 @@ export function useAssistantSession(tabId: string) { session: sessionState?.session || null, messages: sessionState?.messages || [], runs: sessionState?.runs || [], - streamingText: sessionState?.streamingTextByMessageId || {}, + streamingText, isStreaming: sessionState?.isStreaming || false, ensureSession, sendMessage, diff --git a/src/components/AskUserPanel/AskUserPanel.test.tsx b/src/components/AskUserPanel/AskUserPanel.test.tsx index d7854664..d526ba11 100644 --- a/src/components/AskUserPanel/AskUserPanel.test.tsx +++ b/src/components/AskUserPanel/AskUserPanel.test.tsx @@ -53,7 +53,6 @@ const mountWithPending = (pending: PendingAskUser | null) => { messages: [], runs: [], toolCalls: [], - streamingTextByMessageId: {}, isStreaming: false, runStartedAt: null, pendingAskUser: pending, diff --git a/src/components/WorkspaceTaskTranscriptPanel.tsx b/src/components/WorkspaceTaskTranscriptPanel.tsx index 41dd403f..d0f316b9 100644 --- a/src/components/WorkspaceTaskTranscriptPanel.tsx +++ b/src/components/WorkspaceTaskTranscriptPanel.tsx @@ -105,7 +105,9 @@ export default function WorkspaceTaskTranscriptPanel({ const messages = sessionState?.messages || EMPTY_MESSAGES; const toolCalls = sessionState?.toolCalls || EMPTY_TOOL_CALLS; - const streamingText = sessionState?.streamingTextByMessageId || EMPTY_STREAMING; + const streamingText = useAssistantStore( + (state) => (sessionId && state.streamingText[sessionId]) || EMPTY_STREAMING + ); const isStreaming = !!sessionState?.isStreaming; const hasOlderMessages = !!sessionState?.hasOlderMessages; const isLoadingOlderMessages = !!sessionState?.isLoadingOlderMessages; diff --git a/src/pages/Workspace.tsx b/src/pages/Workspace.tsx index f929dcb5..e1fd8d38 100644 --- a/src/pages/Workspace.tsx +++ b/src/pages/Workspace.tsx @@ -67,6 +67,13 @@ type WorkspaceUiState = { previewEntry: PreviewEntry | null; viewingTask: WorkspaceTaskResponse | null; }; +// Stable fallbacks for store-derived values, so re-renders without session +// data don't hand new `[]`/`{}` identities to memoized children each time. +const EMPTY_MESSAGES: AssistantMessage[] = []; +const EMPTY_TOOL_CALLS: ToolInvocation[] = []; +const EMPTY_QUEUED_IDS: string[] = []; +const EMPTY_STREAMING: Record = {}; + const EMPTY_WORKSPACE_UI: WorkspaceUiState = { activePanel: null, previewEntry: null, @@ -1271,7 +1278,6 @@ interface ChatFirstLayoutProps { // conversation is empty before its history finishes loading). isHydrating: boolean; toolCalls: ToolInvocation[]; - streamingText: Record; isStreaming: boolean; runError: string | null; runErrorIsLimit: boolean; @@ -1290,7 +1296,6 @@ const ChatFirstLayout = ({ messages, isHydrating, toolCalls, - streamingText, isStreaming, runError, runErrorIsLimit, @@ -1302,6 +1307,12 @@ const ChatFirstLayout = ({ isLoadingOlderMessages, onLoadOlderMessages, }: ChatFirstLayoutProps) => { + // Streaming deltas are the highest-frequency store updates (many per + // second). Subscribing here — instead of in the Workspace page shell — + // scopes each delta's re-render to the chat area only. + const streamingText = useAssistantStore((state) => + (sessionId && state.streamingText[sessionId]) || EMPTY_STREAMING + ); const cardRef = useRef(null); // Publish the conversation card's viewport geometry as document-level @@ -1533,8 +1544,40 @@ const Workspace = () => { const snapshotReady = snapshot?.workspaceId === workspaceId; const activeSnapshot = snapshotReady ? snapshot : null; const sessionId = activeSnapshot?.session?.id || null; - const sessionState = useAssistantStore((state) => - sessionId ? state.sessions[sessionId] || null : null + // Narrow store subscriptions: each value the page shell renders gets its + // own selector, so this (very large) component body only re-runs when one + // of these actually changes. Subscribing to the whole `sessions[sessionId]` + // object re-rendered the entire page on every session mutation. Streaming + // text deltas are deliberately NOT subscribed here — ChatFirstLayout owns + // that subscription so per-token updates re-render only the chat area. + // Selectors return raw store references (stable across unrelated + // mutations); snapshot fallbacks are applied below. + const storeMessages = useAssistantStore((state) => + sessionId ? state.sessions[sessionId]?.messages : undefined + ); + const storeTotalMessageCount = useAssistantStore((state) => + sessionId ? state.sessions[sessionId]?.totalMessageCount : undefined + ); + const storeToolCalls = useAssistantStore((state) => + sessionId ? state.sessions[sessionId]?.toolCalls : undefined + ); + const isStreaming = useAssistantStore((state) => + sessionId ? state.sessions[sessionId]?.isStreaming || false : false + ); + const runStartedAt = useAssistantStore((state) => + sessionId ? state.sessions[sessionId]?.runStartedAt ?? null : null + ); + const storeQueuedMessageIds = useAssistantStore((state) => + sessionId ? state.sessions[sessionId]?.queuedMessageIds : undefined + ); + const hasOlderMessages = useAssistantStore((state) => + sessionId ? !!state.sessions[sessionId]?.hasOlderMessages : false + ); + const isLoadingOlderMessages = useAssistantStore((state) => + sessionId ? !!state.sessions[sessionId]?.isLoadingOlderMessages : false + ); + const storeRuns = useAssistantStore((state) => + sessionId ? state.sessions[sessionId]?.runs : undefined ); const lastLoadedSessionUpdatedAtRef = useRef(null); @@ -1789,7 +1832,7 @@ const Workspace = () => { }, [artifacts, memories, patchWorkspaceUi] ); - const messages = sessionState?.messages || activeSnapshot?.messages || []; + const messages = storeMessages || activeSnapshot?.messages || EMPTY_MESSAGES; // While the initial entry load is in flight and no messages have arrived // yet, the conversation is unknown — not provably empty. Suppress the // "Start a conversation" empty state until loading settles so an existing @@ -1798,16 +1841,12 @@ const Workspace = () => { // Conversation total from the backend page responses (kept live by the // store as messages stream in); before the first page load reports it, // the loaded window is the best available answer. - const totalMessageCount = sessionState?.totalMessageCount ?? messages.length; - const toolCalls = sessionState?.toolCalls || activeSnapshot?.toolCalls || []; - const streamingText = sessionState?.streamingTextByMessageId || {}; - const isStreaming = sessionState?.isStreaming || false; - const runStartedAt = sessionState?.runStartedAt ?? null; + const totalMessageCount = storeTotalMessageCount ?? messages.length; + const toolCalls = storeToolCalls || activeSnapshot?.toolCalls || EMPTY_TOOL_CALLS; // Store is the live source once the session is hydrated; the snapshot // covers the first render before hydration. - const queuedMessageIds = sessionState?.queuedMessageIds ?? activeSnapshot?.queuedMessageIds ?? []; - const hasOlderMessages = !!sessionState?.hasOlderMessages; - const isLoadingOlderMessages = !!sessionState?.isLoadingOlderMessages; + const queuedMessageIds = + storeQueuedMessageIds ?? activeSnapshot?.queuedMessageIds ?? EMPTY_QUEUED_IDS; const handleDeleteQueuedMessage = useCallback( (messageId: string) => { if (!sessionId) return; @@ -1954,7 +1993,7 @@ const Workspace = () => { // Surface the most recent run's failure in the chat. Derived from the // newest run, so it clears automatically when the next run starts. Without // this, a failed turn (e.g. a provider usage/token limit) shows nothing. - const lastRun = getLastRunInfo(sessionState?.runs || activeSnapshot?.runs); + const lastRun = getLastRunInfo(storeRuns || activeSnapshot?.runs); const runError = lastRun?.status === 'failed' ? lastRun.error?.trim() || 'The run failed.' : null; const runErrorIsLimit = runError ? isUsageLimitError(runError) : false; // Tell the backend this workspace is being viewed so the rail clears its @@ -2053,7 +2092,6 @@ const Workspace = () => { messages={messages} isHydrating={isHydrating} toolCalls={toolCalls} - streamingText={streamingText} isStreaming={isStreaming} runError={runError} runErrorIsLimit={runErrorIsLimit}