Stop re-rendering the Workspace shell on every streaming delta

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();
+  });
+});

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<string, string>;
   /** 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<string, SessionState>;
+  /**
+   * 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<string, Record<string, string>>;
   activeSessionByTab: Record<string, string>;
 
   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<AssistantStoreState>()(
   devtools(
     immer((set, get) => ({
       sessions: {},
+      streamingText: {},
       activeSessionByTab: {},
       recoverablePrompts: {},
 
@@ -195,7 +209,8 @@ const useAssistantStore = create<AssistantStoreState>()(
             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<AssistantStoreState>()(
       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<AssistantStoreState>()(
           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<AssistantStoreState>()(
           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<AssistantStoreState>()(
           }
           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<AssistantStoreState>()(
             // 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<AssistantStoreState>()(
       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];

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');
   });
 });

src/assistant/useAssistantSession.ts

@@ -17,13 +17,20 @@ interface EnsureSessionContext {
 
 const normalizeIdList = (ids: string[] | undefined | null): string[] => [...(ids || [])].sort();
 
+const EMPTY_STREAMING: Record<string, string> = {};
+
 export function useAssistantSession(tabId: string) {
   const sessionId = useAssistantStore(
     (state) => state.activeSessionByTab[tabId]
   );
   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,

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,

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;