diff --git a/frontend/src/components/ai/AIChatPanelConversationView.performance.test.tsx b/frontend/src/components/ai/AIChatPanelConversationView.performance.test.tsx new file mode 100644 index 00000000..8a1c386a --- /dev/null +++ b/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()); +const toolIndexBuilds = vi.hoisted(() => ({ count: 0 })); + +vi.mock('./aiToolResultIndex', async (importOriginal) => { + const actual = await importOriginal(); + 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
{msg.content}
; + }), + }; +}); + +vi.mock('./AIMessageRenderBoundary', () => ({ + default: ({ children }: { children: React.ReactNode }) => <>{children}, +})); + +vi.mock('./AIChatPanelModeContent', () => ({ + default: () => null, +})); + +const overlayTheme = buildOverlayWorkbenchTheme(false); +const noop = () => {}; + +const renderConversation = (messages: AIChatMessage[]) => ( + ()} + 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); + }); +}); diff --git a/frontend/src/components/ai/AIChatPanelConversationView.tsx b/frontend/src/components/ai/AIChatPanelConversationView.tsx index 2ffe68b7..c765c5cb 100644 --- a/frontend/src/components/ai/AIChatPanelConversationView.tsx +++ b/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 = React.memo(({ + message, + toolResultsById, + darkMode, + overlayTheme, + textColor, + activeConnectionId, + activeConnectionConfig, + activeDbName, + onEditMessage, + onRetryMessage, + onDeleteMessage, + onMessageRenderError, +}) => ( + + + +), areAIChatMessageRowPropsEqual); + +interface AIChatMessageListProps extends Omit< + AIChatMessageRowProps, + 'message' | 'toolResultsById' +> { + messages: AIChatMessage[]; +} + +const AIChatMessageList: React.FC = ({ + messages, + ...rowProps +}) => { + const toolResultsById = React.useMemo( + () => buildAIToolResultIndex(messages), + [messages], + ); + + return ( + <> + {messages.map((message) => ( + + ))} + + ); +}; + const AIChatPanelConversationView: React.FC = ({ mode, messages, @@ -87,30 +196,19 @@ const AIChatPanelConversationView: React.FC = isV2Ui={isV2Ui} /> ) : ( - messages.map((message) => ( - - - - )) + ) )} diff --git a/frontend/src/components/ai/AIMessageBubble.test.tsx b/frontend/src/components/ai/AIMessageBubble.test.tsx index b7b2f638..837625b2 100644 --- a/frontend/src/components/ai/AIMessageBubble.test.tsx +++ b/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', - }, - ]} + }], + ])} />, ); diff --git a/frontend/src/components/ai/AIMessageBubble.tsx b/frontend/src/components/ai/AIMessageBubble.tsx index 12b1a766..7193a951 100644 --- a/frontend/src/components/ai/AIMessageBubble.tsx +++ b/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 = React.memo(({ activeConnectionId, activeConnectionConfig, activeDbName, - allMessages, + toolResultsById, }) => { const [isCopied, setIsCopied] = useState(false); const i18n = useOptionalI18n(); @@ -205,7 +206,6 @@ export const AIMessageBubble: React.FC = 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 = React.memo(({ = React.memo(({ ; + +export const buildAIToolResultIndex = ( + messages: readonly AIChatMessage[], +): AIToolResultIndex => { + const toolResultsById = new Map(); + 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) + )); +}; diff --git a/frontend/src/components/ai/messageBubble/AIMessageMarkdown.test.tsx b/frontend/src/components/ai/messageBubble/AIMessageMarkdown.test.tsx index 7f71bd27..49010e92 100644 --- a/frontend/src/components/ai/messageBubble/AIMessageMarkdown.test.tsx +++ b/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} diff --git a/frontend/src/components/ai/messageBubble/AIMessageStatusBlocks.tsx b/frontend/src/components/ai/messageBubble/AIMessageStatusBlocks.tsx index 9858b344..858caadc 100644 --- a/frontend/src/components/ai/messageBubble/AIMessageStatusBlocks.tsx +++ b/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 = ({ export const AIToolCallingBlock: React.FC = ({ 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); diff --git a/frontend/src/store.test.ts b/frontend/src/store.test.ts index b201b139..f5dcad02 100644 --- a/frontend/src/store.test.ts +++ b/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(); @@ -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'); diff --git a/frontend/src/store.ts b/frontend/src/store.ts index 0d832153..d45a7484 100644 --- a/frontend/src/store.ts +++ b/frontend/src/store.ts @@ -5368,7 +5368,12 @@ export const useStore = create()( 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 };