⚡️ perf(ai-chat): 隔离流式消息重绘并复用工具结果索引

将消息行按引用和相关工具结果精确 memo，避免流式输出重绘全部历史消息。

复用单次工具结果索引，并为最新流式消息增加 O(1) 定位热路径与 500 条消息回归测试。

frontend/src/components/ai/AIChatPanelConversationView.performance.test.tsx

@@ -0,0 +1,175 @@
+import React, { createRef } from 'react';
+import { act, create, type ReactTestRenderer } from 'react-test-renderer';
+import { describe, expect, it, vi } from 'vitest';
+
+import type { AIChatMessage } from '../../types';
+import { buildOverlayWorkbenchTheme } from '../../utils/overlayWorkbenchTheme';
+import AIChatPanelConversationView from './AIChatPanelConversationView';
+
+const bubbleRenderCounts = vi.hoisted(() => new Map<string, number>());
+const toolIndexBuilds = vi.hoisted(() => ({ count: 0 }));
+
+vi.mock('./aiToolResultIndex', async (importOriginal) => {
+  const actual = await importOriginal<typeof import('./aiToolResultIndex')>();
+  return {
+    ...actual,
+    buildAIToolResultIndex: (messages: readonly AIChatMessage[]) => {
+      toolIndexBuilds.count += 1;
+      return actual.buildAIToolResultIndex(messages);
+    },
+  };
+});
+
+vi.mock('./AIMessageBubble', async () => {
+  const ReactModule = await import('react');
+  return {
+    AIMessageBubble: ReactModule.memo(({ msg }: { msg: AIChatMessage }) => {
+      bubbleRenderCounts.set(msg.id, (bubbleRenderCounts.get(msg.id) || 0) + 1);
+      return <div data-message-id={msg.id}>{msg.content}</div>;
+    }),
+  };
+});
+
+vi.mock('./AIMessageRenderBoundary', () => ({
+  default: ({ children }: { children: React.ReactNode }) => <>{children}</>,
+}));
+
+vi.mock('./AIChatPanelModeContent', () => ({
+  default: () => null,
+}));
+
+const overlayTheme = buildOverlayWorkbenchTheme(false);
+const noop = () => {};
+
+const renderConversation = (messages: AIChatMessage[]) => (
+  <AIChatPanelConversationView
+    mode="chat"
+    messages={messages}
+    darkMode={false}
+    overlayTheme={overlayTheme}
+    textColor="#0f172a"
+    mutedColor="#64748b"
+    quickActionBg="rgba(255,255,255,0.8)"
+    quickActionBorder="1px solid rgba(0,0,0,0.06)"
+    showScrollBottom={false}
+    contextTableNames={[]}
+    isV2Ui
+    insights={[]}
+    sessions={[]}
+    activeSessionId="session-performance"
+    messagesEndRef={createRef<HTMLDivElement>()}
+    onScrollMessages={noop}
+    onQuickAction={noop}
+    onSelectSession={noop}
+    onEditMessage={noop}
+    onRetryMessage={noop}
+    onDeleteMessage={noop}
+    onMessageRenderError={noop}
+    onScrollBottom={noop}
+  />
+);
+
+describe('AIChatPanelConversationView streaming render performance', () => {
+  it('does not rerender the previous 499 bubbles when only the newest message streams', () => {
+    bubbleRenderCounts.clear();
+    toolIndexBuilds.count = 0;
+    const messages = Array.from({ length: 500 }, (_, index): AIChatMessage => ({
+      id: `message-${index}`,
+      role: index % 2 === 0 ? 'user' : 'assistant',
+      content: `content-${index}`,
+      timestamp: index,
+    }));
+
+    let renderer: ReactTestRenderer;
+    act(() => {
+      renderer = create(renderConversation(messages));
+    });
+    expect(bubbleRenderCounts.get('message-0')).toBe(1);
+    expect(bubbleRenderCounts.get('message-499')).toBe(1);
+
+    const nextMessages = [...messages];
+    nextMessages[499] = { ...messages[499], content: 'content-499-next-token' };
+    act(() => {
+      renderer.update(renderConversation(nextMessages));
+    });
+
+    for (let index = 0; index < 499; index += 1) {
+      expect(bubbleRenderCounts.get(`message-${index}`)).toBe(1);
+    }
+    expect(bubbleRenderCounts.get('message-499')).toBe(2);
+  });
+
+  it('builds one shared tool-result index instead of scanning history in every tool block', () => {
+    bubbleRenderCounts.clear();
+    toolIndexBuilds.count = 0;
+    const messages: AIChatMessage[] = [];
+    for (let index = 0; index < 40; index += 1) {
+      messages.push({
+        id: `assistant-${index}`,
+        role: 'assistant',
+        content: '',
+        timestamp: index * 2,
+        tool_calls: [{
+          id: `call-${index}`,
+          type: 'function',
+          function: { name: 'inspect_ai_runtime', arguments: '{}' },
+        }],
+      });
+      messages.push({
+        id: `tool-${index}`,
+        role: 'tool',
+        content: `result-${index}`,
+        timestamp: index * 2 + 1,
+        tool_call_id: `call-${index}`,
+        tool_name: 'inspect_ai_runtime',
+      });
+    }
+
+    act(() => {
+      create(renderConversation(messages));
+    });
+
+    expect(toolIndexBuilds.count).toBe(1);
+  });
+
+  it('rerenders only the assistant whose relevant tool result changes', () => {
+    bubbleRenderCounts.clear();
+    toolIndexBuilds.count = 0;
+    const toolCallMessage: AIChatMessage = {
+      id: 'assistant-with-tool',
+      role: 'assistant',
+      content: '',
+      timestamp: 1,
+      tool_calls: [{
+        id: 'call-1',
+        type: 'function',
+        function: { name: 'inspect_ai_runtime', arguments: '{}' },
+      }],
+    };
+    const unrelatedMessage: AIChatMessage = {
+      id: 'assistant-unrelated',
+      role: 'assistant',
+      content: 'stable',
+      timestamp: 2,
+    };
+    let renderer: ReactTestRenderer;
+    act(() => {
+      renderer = create(renderConversation([toolCallMessage, unrelatedMessage]));
+    });
+
+    const resultMessage: AIChatMessage = {
+      id: 'tool-result-1',
+      role: 'tool',
+      content: 'result',
+      timestamp: 3,
+      tool_call_id: 'call-1',
+      tool_name: 'inspect_ai_runtime',
+    };
+    act(() => {
+      renderer.update(renderConversation([toolCallMessage, unrelatedMessage, resultMessage]));
+    });
+
+    expect(bubbleRenderCounts.get('assistant-with-tool')).toBe(2);
+    expect(bubbleRenderCounts.get('assistant-unrelated')).toBe(1);
+  });
+});

frontend/src/components/ai/AIChatPanelConversationView.tsx

@@ -12,6 +12,11 @@ import AIChatPanelModeContent, {
   type AIChatInsightItem,
   type AIChatPanelMode,
 } from './AIChatPanelModeContent';
+import {
+  buildAIToolResultIndex,
+  haveSameRelevantToolResults,
+  type AIToolResultIndex,
+} from './aiToolResultIndex';
 
 interface AIChatPanelConversationViewProps {
   mode: AIChatPanelMode;
@@ -43,6 +48,110 @@ interface AIChatPanelConversationViewProps {
   onScrollBottom: () => void;
 }
 
+interface AIChatMessageRowProps {
+  message: AIChatMessage;
+  toolResultsById: AIToolResultIndex;
+  darkMode: boolean;
+  overlayTheme: OverlayWorkbenchTheme;
+  textColor: string;
+  activeConnectionId?: string;
+  activeConnectionConfig?: RpcConnectionConfig;
+  activeDbName?: string;
+  onEditMessage: (message: AIChatMessage) => void;
+  onRetryMessage: (message: AIChatMessage) => void;
+  onDeleteMessage: (id: string) => void;
+  onMessageRenderError: (error: Error, errorInfo: React.ErrorInfo, message: AIChatMessage) => void;
+}
+
+const areAIChatMessageRowPropsEqual = (
+  previous: AIChatMessageRowProps,
+  next: AIChatMessageRowProps,
+): boolean => (
+  previous.message === next.message
+  && previous.darkMode === next.darkMode
+  && previous.overlayTheme === next.overlayTheme
+  && previous.textColor === next.textColor
+  && previous.activeConnectionId === next.activeConnectionId
+  && previous.activeConnectionConfig === next.activeConnectionConfig
+  && previous.activeDbName === next.activeDbName
+  && previous.onEditMessage === next.onEditMessage
+  && previous.onRetryMessage === next.onRetryMessage
+  && previous.onDeleteMessage === next.onDeleteMessage
+  && previous.onMessageRenderError === next.onMessageRenderError
+  && haveSameRelevantToolResults(
+    next.message.tool_calls,
+    previous.toolResultsById,
+    next.toolResultsById,
+  )
+);
+
+const AIChatMessageRow: React.FC<AIChatMessageRowProps> = React.memo(({
+  message,
+  toolResultsById,
+  darkMode,
+  overlayTheme,
+  textColor,
+  activeConnectionId,
+  activeConnectionConfig,
+  activeDbName,
+  onEditMessage,
+  onRetryMessage,
+  onDeleteMessage,
+  onMessageRenderError,
+}) => (
+  <AIMessageRenderBoundary
+    msg={message}
+    darkMode={darkMode}
+    overlayTheme={overlayTheme}
+    onDeleteMessage={onDeleteMessage}
+    onError={onMessageRenderError}
+  >
+    <AIMessageBubble
+      msg={message}
+      darkMode={darkMode}
+      overlayTheme={overlayTheme}
+      textColor={textColor}
+      onEdit={onEditMessage}
+      onRetry={onRetryMessage}
+      onDelete={onDeleteMessage}
+      activeConnectionId={activeConnectionId}
+      activeConnectionConfig={activeConnectionConfig}
+      activeDbName={activeDbName}
+      toolResultsById={toolResultsById}
+    />
+  </AIMessageRenderBoundary>
+), areAIChatMessageRowPropsEqual);
+
+interface AIChatMessageListProps extends Omit<
+  AIChatMessageRowProps,
+  'message' | 'toolResultsById'
+> {
+  messages: AIChatMessage[];
+}
+
+const AIChatMessageList: React.FC<AIChatMessageListProps> = ({
+  messages,
+  ...rowProps
+}) => {
+  const toolResultsById = React.useMemo(
+    () => buildAIToolResultIndex(messages),
+    [messages],
+  );
+
+  return (
+    <>
+      {messages.map((message) => (
+        <AIChatMessageRow
+          key={message.id}
+          {...rowProps}
+          message={message}
+          toolResultsById={toolResultsById}
+        />
+      ))}
+    </>
+  );
+};
+
 const AIChatPanelConversationView: React.FC<AIChatPanelConversationViewProps> = ({
   mode,
   messages,
@@ -87,30 +196,19 @@ const AIChatPanelConversationView: React.FC<AIChatPanelConversationViewProps> =
             isV2Ui={isV2Ui}
           />
         ) : (
-          messages.map((message) => (
-            <AIMessageRenderBoundary
-              key={message.id}
-              msg={message}
-              darkMode={darkMode}
-              overlayTheme={overlayTheme}
-              onDeleteMessage={onDeleteMessage}
-              onError={onMessageRenderError}
-            >
-              <AIMessageBubble
-                msg={message}
-                darkMode={darkMode}
-                overlayTheme={overlayTheme}
-                textColor={textColor}
-                onEdit={onEditMessage}
-                onRetry={onRetryMessage}
-                onDelete={onDeleteMessage}
-                activeConnectionId={activeConnectionId}
-                activeConnectionConfig={activeConnectionConfig}
-                activeDbName={activeDbName}
-                allMessages={messages}
-              />
-            </AIMessageRenderBoundary>
-          ))
+          <AIChatMessageList
+            messages={messages}
+            darkMode={darkMode}
+            overlayTheme={overlayTheme}
+            textColor={textColor}
+            activeConnectionId={activeConnectionId}
+            activeConnectionConfig={activeConnectionConfig}
+            activeDbName={activeDbName}
+            onEditMessage={onEditMessage}
+            onRetryMessage={onRetryMessage}
+            onDeleteMessage={onDeleteMessage}
+            onMessageRenderError={onMessageRenderError}
+          />
         )
       )}
 

frontend/src/components/ai/AIMessageBubble.test.tsx

@@ -56,16 +56,16 @@ describe('AIMessageBubble', () => {
         onEdit={() => {}}
         onRetry={() => {}}
         onDelete={() => {}}
-        allMessages={[
-          {
+        toolResultsById={new Map([
+          ['tool-1', {
             id: 'tool-result-1',
             role: 'tool',
             content: '[{\"fk\":\"orders.customer_id\"}]',
             timestamp: Date.now(),
             tool_call_id: 'tool-1',
             tool_name: 'get_foreign_keys',
-          },
-        ]}
+          }],
+        ])}
       />,
     );
 

frontend/src/components/ai/AIMessageBubble.tsx

@@ -26,6 +26,7 @@ import {
 import { AIMessageMarkdown } from './messageBubble/AIMessageMarkdown';
 import { AIThinkingBlock, AIToolCallingBlock } from './messageBubble/AIMessageStatusBlocks';
 import { formatAIChatAttachmentSize } from './aiChatAttachments';
+import type { AIToolResultIndex } from './aiToolResultIndex';
 
 interface AIMessageBubbleProps {
   msg: AIChatMessage;
@@ -38,7 +39,7 @@ interface AIMessageBubbleProps {
   activeConnectionId?: string;
   activeConnectionConfig?: any;
   activeDbName?: string;
-  allMessages?: AIChatMessage[];
+  toolResultsById: AIToolResultIndex;
 }
 
 interface AIMessageActionBarProps {
@@ -197,7 +198,7 @@ export const AIMessageBubble: React.FC<AIMessageBubbleProps> = React.memo(({
   activeConnectionId,
   activeConnectionConfig,
   activeDbName,
-  allMessages,
+  toolResultsById,
 }) => {
   const [isCopied, setIsCopied] = useState(false);
   const i18n = useOptionalI18n();
@@ -205,7 +206,6 @@ export const AIMessageBubble: React.FC<AIMessageBubbleProps> = React.memo(({
     i18n?.t ?? ((catalogKey, catalogParams) => catalogTranslate('en-US', catalogKey, catalogParams))
   )(key, params);
   const isUser = msg.role === 'user';
-  const toolMessages = allMessages || [];
 
   const { displayContent, parsedThinking } = React.useMemo(() => {
     const content = msg.content || '';
@@ -280,7 +280,7 @@ export const AIMessageBubble: React.FC<AIMessageBubbleProps> = React.memo(({
               <AIToolCallingBlock
                 toolCalls={msg.tool_calls}
                 loading={Boolean(msg.loading)}
-                allMessages={toolMessages}
+                toolResultsById={toolResultsById}
                 darkMode={darkMode}
                 overlayTheme={overlayTheme}
                 hasContent={false}
@@ -450,7 +450,7 @@ export const AIMessageBubble: React.FC<AIMessageBubbleProps> = React.memo(({
             <AIToolCallingBlock
               toolCalls={msg.tool_calls}
               loading={Boolean(msg.loading)}
-              allMessages={toolMessages}
+              toolResultsById={toolResultsById}
               darkMode={darkMode}
               overlayTheme={overlayTheme}
               hasContent={Boolean(msg.content)}

frontend/src/components/ai/aiToolResultIndex.ts

@@ -0,0 +1,28 @@
+import type { AIChatMessage, AIToolCall } from '../../types';
+
+export type AIToolResultIndex = ReadonlyMap<string, AIChatMessage>;
+
+export const buildAIToolResultIndex = (
+  messages: readonly AIChatMessage[],
+): AIToolResultIndex => {
+  const toolResultsById = new Map<string, AIChatMessage>();
+  for (const message of messages) {
+    if (message.role === 'tool' && message.tool_call_id) {
+      toolResultsById.set(message.tool_call_id, message);
+    }
+  }
+  return toolResultsById;
+};
+
+export const haveSameRelevantToolResults = (
+  toolCalls: readonly AIToolCall[] | undefined,
+  previous: AIToolResultIndex,
+  next: AIToolResultIndex,
+): boolean => {
+  if (previous === next || !toolCalls || toolCalls.length === 0) {
+    return true;
+  }
+  return toolCalls.every((toolCall) => (
+    previous.get(toolCall.id) === next.get(toolCall.id)
+  ));
+};

frontend/src/components/ai/messageBubble/AIMessageMarkdown.test.tsx

@@ -134,14 +134,14 @@ describe('AIMessageStatusBlocks', () => {
           function: { name: 'inspect_ai_runtime', arguments: '{}' },
         }]}
         loading={false}
-        allMessages={[{
+        toolResultsById={new Map([['call-1', {
           id: 'tool-1',
           role: 'tool',
           content: 'result payload',
           timestamp: 1,
           tool_call_id: 'call-1',
           tool_name: 'inspect_ai_runtime',
-        }]}
+        }]])}
         darkMode={false}
         overlayTheme={overlayTheme}
         hasContent={false}
@@ -175,14 +175,14 @@ describe('AIMessageStatusBlocks', () => {
             function: { name: 'inspect_ai_runtime', arguments: '{}' },
           }]}
           loading={false}
-          allMessages={[{
+          toolResultsById={new Map([['call-1', {
             id: 'tool-1',
             role: 'tool',
             content: 'result payload',
             timestamp: 1,
             tool_call_id: 'call-1',
             tool_name: 'inspect_ai_runtime',
-          }]}
+          }]])}
           darkMode={false}
           overlayTheme={overlayTheme}
           hasContent={false}

frontend/src/components/ai/messageBubble/AIMessageStatusBlocks.tsx

@@ -1,4 +1,4 @@
-import React, { useEffect, useMemo, useState } from 'react';
+import React, { useEffect, useState } from 'react';
 import { ApiOutlined, CaretDownOutlined, CaretRightOutlined, CheckOutlined } from '@ant-design/icons';
 
 import { t as catalogTranslate } from '../../../i18n/catalog';
@@ -6,6 +6,7 @@ import type { I18nParams } from '../../../i18n/types';
 import { useOptionalI18n } from '../../../i18n/provider';
 import type { AIChatMessage, AIToolCall } from '../../../types';
 import type { OverlayWorkbenchTheme } from '../../../utils/overlayWorkbenchTheme';
+import type { AIToolResultIndex } from '../aiToolResultIndex';
 
 interface AIThinkingBlockProps {
   displayThinking: string;
@@ -19,7 +20,7 @@ interface AIThinkingBlockProps {
 interface AIToolCallingBlockProps {
   toolCalls: AIToolCall[];
   loading: boolean;
-  allMessages: AIChatMessage[];
+  toolResultsById: AIToolResultIndex;
   darkMode: boolean;
   overlayTheme: OverlayWorkbenchTheme;
   hasContent: boolean;
@@ -199,19 +200,12 @@ export const AIThinkingBlock: React.FC<AIThinkingBlockProps> = ({
 export const AIToolCallingBlock: React.FC<AIToolCallingBlockProps> = ({
   toolCalls,
   loading,
-  allMessages,
+  toolResultsById,
   darkMode,
   overlayTheme,
   hasContent,
 }) => {
   const copy = useMessageCopy();
-  const toolResultsById = useMemo(() => {
-    return new Map(
-      allMessages
-        .filter((message) => message.role === 'tool' && message.tool_call_id)
-        .map((message) => [message.tool_call_id as string, message]),
-    );
-  }, [allMessages]);
   const allDone = toolCalls.every((toolCall) => toolResultsById.has(toolCall.id));
   const [expanded, setExpanded] = useState(!allDone && loading);
 

frontend/src/store.test.ts

@@ -1,5 +1,6 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import { SIDEBAR_RESIZE_MAX_WIDTH } from './utils/sidebarLayout';
+import type { AIChatMessage } from './types';
 
 class MemoryStorage implements Storage {
   private data = new Map<string, string>();
@@ -1942,6 +1943,43 @@ describe('store appearance persistence', () => {
     }
   });
 
+  it('finds the newest streaming message without scanning the full session history', async () => {
+    vi.useFakeTimers();
+    try {
+      const { useStore } = await importStore();
+      let messageIdReads = 0;
+      const messages = Array.from({ length: 500 }, (_, index): AIChatMessage => ({
+        id: `message-${index}`,
+        role: index % 2 === 0 ? 'user' : 'assistant',
+        content: `content-${index}`,
+        timestamp: index,
+      })).map((message) => new Proxy(message, {
+        get(target, property, receiver) {
+          if (property === 'id') {
+            messageIdReads += 1;
+          }
+          return Reflect.get(target, property, receiver);
+        },
+      }));
+      useStore.setState({
+        aiChatHistory: { 'session-stream': messages },
+      });
+      messageIdReads = 0;
+
+      useStore.getState().updateAIChatMessage('session-stream', 'message-499', {
+        content: 'content-499-next-token',
+      });
+
+      expect(messageIdReads).toBeLessThanOrEqual(2);
+      expect(useStore.getState().aiChatHistory['session-stream'][499]?.content).toBe(
+        'content-499-next-token',
+      );
+      expect(useStore.getState().aiChatHistory['session-stream'][0]).toBe(messages[0]);
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
   it('keeps store fallback titles out of production source literals', async () => {
     const { readFileSync } = await import('node:fs');
     const source = readFileSync(new URL('./store.ts', import.meta.url), 'utf8');

frontend/src/store.ts

@@ -5368,7 +5368,12 @@ export const useStore = create<AppState>()(
         set((state) => {
           const messages = state.aiChatHistory[sessionId];
           if (!messages) return state;
-          const idx = messages.findIndex((m) => m.id === messageId);
+          // Message IDs are unique within a session and are also used as React keys.
+          // Streaming updates target the newest assistant message, so keep that hot path O(1).
+          const lastIndex = messages.length - 1;
+          const idx = lastIndex >= 0 && messages[lastIndex].id === messageId
+            ? lastIndex
+            : messages.findIndex((m) => m.id === messageId);
           if (idx < 0) return state;
           const newMessages = [...messages];
           newMessages[idx] = { ...newMessages[idx], ...updates };