From 84a8e78f6bfb040083a278a16eaa74c601853bd5 Mon Sep 17 00:00:00 2001 From: Saxon Fletcher Date: Fri, 25 Sep 2026 11:02:05 +1000 Subject: [PATCH 1/2] chore(studio): improve assistant feed performance --- .../ui/AIAssistantPanel/AssistantChat.tsx | 42 ++-- .../AIAssistantPanel/AssistantQueryCell.tsx | 26 +- .../ui/AIAssistantPanel/Message.Display.tsx | 14 +- .../ui/AIAssistantPanel/Message.Parts.tsx | 114 ++++----- .../Message.performance.test.tsx | 229 ++++++++++++++++++ .../ui/AIAssistantPanel/Message.tsx | 71 ++++-- .../ui/AIAssistantPanel/MessageMarkdown.tsx | 4 +- .../elements/Conversation.tsx | 20 +- .../src/CodeBlock/CodeBlock.test.tsx | 26 ++ .../ui-patterns/src/CodeBlock/CodeBlock.tsx | 49 ++-- 10 files changed, 455 insertions(+), 140 deletions(-) create mode 100644 apps/studio/components/ui/AIAssistantPanel/Message.performance.test.tsx create mode 100644 packages/ui-patterns/src/CodeBlock/CodeBlock.test.tsx diff --git a/apps/studio/components/ui/AIAssistantPanel/AssistantChat.tsx b/apps/studio/components/ui/AIAssistantPanel/AssistantChat.tsx index 24c7add23cc..96a7546c00e 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AssistantChat.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AssistantChat.tsx @@ -33,6 +33,7 @@ import { Markdown } from '@/components/interfaces/Markdown' import { useCheckOpenAIKeyQuery } from '@/data/ai/check-api-key-query' import { useRateMessageMutation } from '@/data/ai/rate-message-mutation' import { useTablesQuery } from '@/data/tables/tables-query' +import { useLatest } from '@/hooks/misc/useLatest' import { useLocalStorageQuery } from '@/hooks/misc/useLocalStorage' import { useOrgAiOptInLevel } from '@/hooks/misc/useOrgOptedIntoAi' import { useSelectedOrganizationQuery } from '@/hooks/misc/useSelectedOrganization' @@ -172,6 +173,8 @@ export const AssistantChat = ({ regenerate, } = useChat({ id: chatId, + // Batch token updates without throttling the SDK's tool execution or approval state. + throttle: 50, ...(chatInstance ? { chat: chatInstance } : {}), sendAutomaticallyWhen: lastAssistantMessageIsCompleteWithApprovalResponses, onError: onErrorChat, @@ -186,32 +189,36 @@ export const AssistantChat = ({ const isChatInputDisabled = !isApiKeySet || disablePrompts || isLoadingOrganization || isSupportChatClosed + const messagesRef = useLatest(chatMessages) + const isChatLoadingRef = useLatest(isChatLoading) + const branchedFrom = currentChat?.branchedFrom const branchedConversation = branchedFrom ? snap.chats[branchedFrom.chatId] : undefined const deleteMessageFromHere = useCallback( (messageId: string) => { - // Find the message index in current chatMessages - const messageIndex = chatMessages.findIndex((msg) => msg.id === messageId) + const messages = messagesRef.current + const messageIndex = messages.findIndex((msg) => msg.id === messageId) if (messageIndex === -1) return - if (isChatLoading) stop() + if (isChatLoadingRef.current) stop() - snap.deleteMessagesAfter(messageId, { includeSelf: true, chatId }) + state.deleteMessagesAfter(messageId, { includeSelf: true, chatId }) - const updatedMessages = chatMessages.slice(0, messageIndex) + const updatedMessages = messages.slice(0, messageIndex) setMessages(updatedMessages) }, - [snap, setMessages, chatMessages, isChatLoading, stop, chatId] + [state, setMessages, messagesRef, isChatLoadingRef, stop, chatId] ) const editMessage = useCallback( (messageId: string) => { - const messageIndex = chatMessages.findIndex((msg) => msg.id === messageId) + const messages = messagesRef.current + const messageIndex = messages.findIndex((msg) => msg.id === messageId) if (messageIndex === -1) return // Target message - const messageToEdit = chatMessages[messageIndex] + const messageToEdit = messages[messageIndex] // Activate editing mode setEditingMessageId(messageId) @@ -233,7 +240,7 @@ export const AssistantChat = ({ } }, 100) }, - [chatMessages, setValue] + [messagesRef, setValue] ) const cancelEdit = useCallback(() => { @@ -251,7 +258,7 @@ export const AssistantChat = ({ try { const result = await rateMessage({ rating, - messages: chatMessages, + messages: messagesRef.current, messageId, projectRef: project.ref, orgSlug: selectedOrganization.slug, @@ -274,7 +281,7 @@ export const AssistantChat = ({ }) } }, - [chatMessages, project?.ref, selectedOrganization?.slug, rateMessage, track, state, chatId] + [messagesRef, project?.ref, selectedOrganization?.slug, rateMessage, track, state, chatId] ) const isContextExceededError = @@ -282,13 +289,15 @@ export const AssistantChat = ({ (error.message?.includes('context_length_exceeded') || error.message?.includes('exceeds the context window')) + const editedMessageIndex = editingMessageId + ? chatMessages.findIndex((message) => message.id === editingMessageId) + : -1 + const renderedMessages = useMemo( () => chatMessages.map((message, index) => { const isBeingEdited = editingMessageId === message.id - const isAfterEditedMessage = editingMessageId - ? chatMessages.findIndex((m) => m.id === editingMessageId) < index - : false + const isAfterEditedMessage = !!editingMessageId && editedMessageIndex < index const isLastMessage = index === chatMessages.length - 1 return ( @@ -334,6 +343,7 @@ export const AssistantChat = ({ editMessage, cancelEdit, editingMessageId, + editedMessageIndex, chatStatus, addToolApprovalResponse, handleRateMessage, @@ -516,7 +526,9 @@ export const AssistantChat = ({ onStop={() => { stop() // to save partial responses from the AI - const lastMessage = chatMessages[chatMessages.length - 1] + // Read the live SDK state: the rendered snapshot may trail the stream by 50ms. + const messages = chatInstance?.messages ?? chatMessages + const lastMessage = messages[messages.length - 1] if (lastMessage && lastMessage.role === 'assistant') { state.updateMessage(lastMessage, chatId) } diff --git a/apps/studio/components/ui/AIAssistantPanel/AssistantQueryCell.tsx b/apps/studio/components/ui/AIAssistantPanel/AssistantQueryCell.tsx index 1cb22012512..120468bc62e 100644 --- a/apps/studio/components/ui/AIAssistantPanel/AssistantQueryCell.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/AssistantQueryCell.tsx @@ -1,4 +1,4 @@ -import { useRef, useState } from 'react' +import { useMemo, useRef, useState } from 'react' import { identifyQueryType } from './AIAssistant.utils' import { @@ -106,15 +106,18 @@ export const AssistantQueryCell = ({ } const result = resultOverride === undefined ? initialResult : (resultOverride ?? undefined) - const display = - localDisplay ?? - getAssistantQueryDisplay({ - view, - xAxis, - yAxis, - sql: query.uncheckedSql, - rows: result?.rows, - }) + const inferredDisplay = useMemo( + () => + getAssistantQueryDisplay({ + view, + xAxis, + yAxis, + sql: query.uncheckedSql, + rows: result?.rows, + }), + [view, xAxis, yAxis, query.uncheckedSql, result?.rows] + ) + const display = localDisplay ?? inferredDisplay const handleTitleChange = (value: string) => { const nextTitle = value.trim() @@ -165,11 +168,12 @@ export const AssistantQueryCell = ({ onCancel={onDeny} onConfirm={onApprove} > + {/* Keep editor state mounted; the fixed height preserves scroll geometry when skipped. */} {children} } -function MessageDisplayContent({ message }: { message: VercelMessage }) { +const MessageDisplayContent = memo(function MessageDisplayContent({ + message, +}: { + message: VercelMessage +}) { const { id, isLoading, readOnly } = useMessageInfoContext() const messageParts = message.parts @@ -53,7 +57,9 @@ function MessageDisplayContent({ message }: { message: VercelMessage }) {
{messageParts?.length > 0 ? messageParts.map((part: NonNullable, idx) => { - return + // The SDK exposes its mutable object on the first write, then publishes clones. + // Capture state/text now so later mutations cannot change memo's previous props. + return }) : content && (
@@ -64,7 +70,7 @@ function MessageDisplayContent({ message }: { message: VercelMessage }) { )}
) -} +}) function MessageDisplayTextMessage({ id, diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx index 8c3d255bab6..3ecf20f7477 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx @@ -1,7 +1,8 @@ import { UIMessage as VercelMessage } from '@ai-sdk/react' import { type DynamicToolUIPart, type ReasoningUIPart, type TextUIPart, type ToolUIPart } from 'ai' +import isEqual from 'lodash/isEqual' import { BrainIcon, CheckIcon, Loader2 } from 'lucide-react' -import { type ReactNode } from 'react' +import { memo, type ReactNode } from 'react' import { cn } from 'ui' import { AssistantQueryCell } from './AssistantQueryCell' @@ -333,63 +334,66 @@ const isCompactToolPart = (part: NonNullable[number]) => part.type === 'tool-get_active_incidents' || part.type === 'tool-load_knowledge' -export function MessagePartSwitcher({ - part, -}: { - part: NonNullable[number] -}) { - const content = (() => { - switch (part.type) { - case 'dynamic-tool': { - if (part.toolName === 'query_logs') { +export const MessagePartSwitcher = memo( + function MessagePartSwitcher({ part }: { part: NonNullable[number] }) { + const content = (() => { + switch (part.type) { + case 'dynamic-tool': { + if (part.toolName === 'query_logs') { + return + } + return + } + case 'tool-list_policies': + case 'tool-search_docs': + case 'tool-get_active_incidents': + case 'tool-load_knowledge': { + return + } + case 'reasoning': + return + case 'text': + return + + case 'tool-execute_sql': { + return + } + case 'tool-query_logs': { return } - return - } - case 'tool-list_policies': - case 'tool-search_docs': - case 'tool-get_active_incidents': - case 'tool-load_knowledge': { - return - } - case 'reasoning': - return - case 'text': - return + case 'tool-deploy_edge_function': { + return + } + case 'tool-create_notebook': { + return + } + case 'tool-update_notebook': { + return + } + case 'tool-delete_notebook': { + return + } + case 'tool-run_notebook': { + return + } - case 'tool-execute_sql': { - return - } - case 'tool-query_logs': { - return - } - case 'tool-deploy_edge_function': { - return - } - case 'tool-create_notebook': { - return - } - case 'tool-update_notebook': { - return - } - case 'tool-delete_notebook': { - return - } - case 'tool-run_notebook': { - return + case 'source-url': + case 'source-document': + case 'file': + default: + return null } + })() - case 'source-url': - case 'source-document': - case 'file': - default: - return null - } - })() + if (content === null) return null + // Tool rows depend on being direct siblings to share their compact spacing and dividers. + if (isCompactToolPart(part)) return content - if (content === null) return null - // Tool rows depend on being direct siblings to share their compact spacing and dividers. - if (isCompactToolPart(part)) return content - - return {content} -} + return {content} + }, + (previous, next) => { + // The AI SDK clones every part of the active message on each stream update. + // Compare the part, not its identity, so completed tools retain their rendered subtree. + return isEqual(previous.part, next.part) + } +) diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.performance.test.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.performance.test.tsx new file mode 100644 index 00000000000..bdd8b19a2d2 --- /dev/null +++ b/apps/studio/components/ui/AIAssistantPanel/Message.performance.test.tsx @@ -0,0 +1,229 @@ +import { fireEvent, render, screen } from '@testing-library/react' +import type { UIMessage } from 'ai' +import { useState } from 'react' +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import { Message } from './Message' + +const { renderQuery, renderFunction, renderMarkdown } = vi.hoisted(() => ({ + renderQuery: vi.fn(), + renderFunction: vi.fn(), + renderMarkdown: vi.fn(), +})) + +// Keep the message/context/part pipeline real, substituting stateful probes for the +// expensive leaves so we can detect both extra renders and lost local state. +vi.mock('./AssistantQueryCell', () => ({ + AssistantQueryCell: (props: { initialResult?: { rows: unknown[] }; onApprove?: () => void }) => { + renderQuery(props) + const [runs, setRuns] = useState(0) + return ( +
+ + + {JSON.stringify(props.initialResult?.rows)} +
+ ) + }, +})) + +vi.mock('./EdgeFunctionRenderer', () => ({ + EdgeFunctionRenderer: (props: { code: string }) => { + renderFunction(props) + return 
{props.code}
+ }, +})) + +vi.mock('./MessageMarkdown', () => ({ + MessageMarkdown: ({ children }: { children: string }) => { + renderMarkdown(children) + return

{children}

+ }, +})) + +vi.mock('./Message.Actions', () => ({ MessageActions: () => null })) + +const initialMessage: UIMessage = { + id: 'assistant-1', + role: 'assistant', + parts: [ + { + type: 'tool-execute_sql', + toolCallId: 'query-1', + state: 'approval-requested', + approval: { id: 'approval-1' }, + input: { sql: 'select 1', label: 'Query', view: 'table' }, + }, + { + type: 'tool-deploy_edge_function', + toolCallId: 'function-1', + state: 'input-available', + input: { + code: 'Deno.serve(() => new Response("hello"))', + label: 'Function', + functionName: 'hello', + }, + }, + { type: 'text', text: 'Working' }, + ], +} + +const callbacks = { + onDelete: vi.fn(), + onEdit: vi.fn(), + onBranch: vi.fn(), + onCancelEdit: vi.fn(), + addToolApprovalResponse: vi.fn(), +} + +function FeedMessage({ + message = initialMessage, + isLoading = true, + addToolApprovalResponse = callbacks.addToolApprovalResponse, +}: { + message?: UIMessage + isLoading?: boolean + addToolApprovalResponse?: typeof callbacks.addToolApprovalResponse +}) { + return ( + + ) +} + +describe('assistant feed rendering', () => { + beforeEach(() => vi.clearAllMocks()) + + it('does not rerender unchanged history when the parent renders', () => { + const { rerender } = render() + const counts = [ + renderQuery.mock.calls.length, + renderFunction.mock.calls.length, + renderMarkdown.mock.calls.length, + ] + + rerender() + + expect([ + renderQuery.mock.calls.length, + renderFunction.mock.calls.length, + renderMarkdown.mock.calls.length, + ]).toEqual(counts) + }) + + it('updates streamed text without rerendering cloned tools or losing their local state', () => { + const { rerender } = render() + fireEvent.click(screen.getByRole('button', { name: 'Run count: 0' })) + const queryRenders = renderQuery.mock.calls.length + const functionRenders = renderFunction.mock.calls.length + + const updated = structuredClone(initialMessage) + updated.parts[2] = { type: 'text', text: 'Working on the next step' } + rerender() + + expect(screen.getByText('Working on the next step')).toBeInTheDocument() + expect(screen.getByRole('button', { name: 'Run count: 1' })).toBeInTheDocument() + expect(renderQuery).toHaveBeenCalledTimes(queryRenders) + expect(renderFunction).toHaveBeenCalledTimes(functionRenders) + }) + + it('updates tool output and code when their content changes', () => { + const { rerender } = render() + const updated = structuredClone(initialMessage) + updated.parts[0] = { + type: 'tool-execute_sql', + toolCallId: 'query-1', + state: 'output-available', + input: { sql: 'select 1', label: 'Query', view: 'table' }, + output: [{ value: 1 }], + } + updated.parts[1] = { + type: 'tool-deploy_edge_function', + toolCallId: 'function-1', + state: 'input-available', + input: { code: 'updated code', label: 'Function', functionName: 'hello' }, + } + rerender() + + expect(screen.getByText('[{"value":1}]')).toBeInTheDocument() + expect(screen.getByText('updated code')).toBeInTheDocument() + }) + + it('uses the latest approval callback even when the tool part is unchanged', () => { + const { rerender } = render() + const approve = vi.fn() + rerender() + fireEvent.click(screen.getByRole('button', { name: 'Approve query' })) + + expect(approve).toHaveBeenCalledWith({ id: 'approval-1', approved: true }) + expect(callbacks.addToolApprovalResponse).not.toHaveBeenCalled() + }) + + it('retains query state when streaming completes', () => { + const { rerender } = render() + fireEvent.click(screen.getByRole('button', { name: 'Run count: 0' })) + rerender() + expect(screen.getByRole('button', { name: 'Run count: 1' })).toBeInTheDocument() + }) + + it('finishes reasoning when the SDK mutates the first streamed part before publishing a snapshot', () => { + const reasoning = { + type: 'reasoning' as const, + text: '', + state: 'streaming' as 'streaming' | 'done', + } + const message: UIMessage = { id: 'reasoning-1', role: 'assistant', parts: [reasoning] } + const { rerender } = render() + expect(screen.getByText('Thinking...')).toBeInTheDocument() + + // Chat.pushMessage exposes the initial object; subsequent replaceMessage calls clone it. + reasoning.state = 'done' + rerender() + + expect(screen.queryByText('Thinking...')).not.toBeInTheDocument() + expect(screen.getByText('Reasoned')).toBeInTheDocument() + }) + + it('updates text when the SDK mutates the first streamed part', () => { + const text = { type: 'text' as const, text: 'First token', state: 'streaming' as const } + const message: UIMessage = { id: 'text-1', role: 'assistant', parts: [text] } + const { rerender } = render() + + text.text = 'First token and the rest of the response' + rerender() + + expect(screen.getByText(text.text)).toBeInTheDocument() + }) + + it('updates a tool when its initial input-streaming part is mutated to a completed result', () => { + const tool = { + type: 'tool-execute_sql' as const, + toolCallId: 'query-1', + state: 'input-streaming' as const, + } + const message: UIMessage = { id: 'tool-1', role: 'assistant', parts: [tool] } + const { rerender } = render() + expect(screen.getByText('Writing SQL...')).toBeInTheDocument() + + Object.assign(tool, { + state: 'output-available', + input: { sql: 'select 1' }, + output: [{ value: 1 }], + }) + rerender() + + expect(screen.queryByText('Writing SQL...')).not.toBeInTheDocument() + expect(screen.getByText('[{"value":1}]')).toBeInTheDocument() + }) +}) diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.tsx index b719c2cab0f..60185f9a821 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/Message.tsx @@ -1,5 +1,5 @@ import { UIMessage as VercelMessage } from '@ai-sdk/react' -import { useState } from 'react' +import { memo, useMemo, useState } from 'react' import { toast } from 'sonner' import { cn, copyToClipboard } from 'ui' @@ -114,38 +114,59 @@ interface MessageProps { rating?: 'positive' | 'negative' | null } -export function Message(props: MessageProps) { +export const Message = memo(function Message(props: MessageProps) { const message = props.message const { role } = message const isUserMessage = role === 'user' + let messageState: MessageInfo['state'] = 'idle' + if (props.isBeingEdited) messageState = 'editing' + else if (props.isAfterEditedMessage) messageState = 'predecessor-editing' - const messageInfo = { - id: props.id, - isLoading: props.isLoading, - readOnly: props.readOnly, - variant: props.variant, - isUserMessage, - state: props.isBeingEdited - ? 'editing' - : props.isAfterEditedMessage - ? 'predecessor-editing' - : 'idle', - isLastMessage: props.isLastMessage, - rating: props.rating, - } satisfies MessageInfo + const messageInfo = useMemo( + () => ({ + id: props.id, + isLoading: props.isLoading, + readOnly: props.readOnly, + variant: props.variant, + isUserMessage, + state: messageState, + isLastMessage: props.isLastMessage, + rating: props.rating, + }), + [ + props.id, + props.isLoading, + props.readOnly, + props.variant, + isUserMessage, + messageState, + props.isLastMessage, + props.rating, + ] + ) - const messageActions = { - addToolApprovalResponse: props.addToolApprovalResponse, - onDelete: props.onDelete, - onEdit: props.onEdit, - onBranch: props.onBranch, - onCancelEdit: props.onCancelEdit, - onRate: props.onRate, - } + const messageActions = useMemo( + () => ({ + addToolApprovalResponse: props.addToolApprovalResponse, + onDelete: props.onDelete, + onEdit: props.onEdit, + onBranch: props.onBranch, + onCancelEdit: props.onCancelEdit, + onRate: props.onRate, + }), + [ + props.addToolApprovalResponse, + props.onDelete, + props.onEdit, + props.onBranch, + props.onCancelEdit, + props.onRate, + ] + ) return ( {isUserMessage ? : } ) -} +}) diff --git a/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx b/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx index 70f452668a2..d31d37696d2 100644 --- a/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/MessageMarkdown.tsx @@ -128,7 +128,7 @@ const baseMarkdownComponents = { ), } -export function MessageMarkdown({ +export const MessageMarkdown = memo(function MessageMarkdown({ id, isLoading, readOnly, @@ -171,7 +171,7 @@ export function MessageMarkdown({ {markdownSource} ) -} +}) export const MarkdownPre = ({ children, diff --git a/apps/studio/components/ui/AIAssistantPanel/elements/Conversation.tsx b/apps/studio/components/ui/AIAssistantPanel/elements/Conversation.tsx index c2b1279eb20..8c67f8735a2 100644 --- a/apps/studio/components/ui/AIAssistantPanel/elements/Conversation.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/elements/Conversation.tsx @@ -20,7 +20,7 @@ const FADE_GUTTER = 'inset-x-7' export const Conversation = ({ className, children, ...props }: ConversationProps) => ( ) -export const ConversationContent = ({ className, ...props }: ConversationContentProps) => ( - -) +export const ConversationContent = ({ + className, + children, + ...props +}: ConversationContentProps) => { + const context = useStickToBottomContext() + + return ( +
+
+ {typeof children === 'function' ? children(context) : children} +
+
+ ) +} export const ConversationScrollButton = ({ className, diff --git a/packages/ui-patterns/src/CodeBlock/CodeBlock.test.tsx b/packages/ui-patterns/src/CodeBlock/CodeBlock.test.tsx new file mode 100644 index 00000000000..efcc97036cb --- /dev/null +++ b/packages/ui-patterns/src/CodeBlock/CodeBlock.test.tsx @@ -0,0 +1,26 @@ +import { fireEvent, render, screen } from '@testing-library/react' +import { describe, expect, it, vi } from 'vitest' + +import { CodeBlock } from './CodeBlock' + +describe('CodeBlock', () => { + it('highlights updated code and copies it with the latest callback', () => { + const copy = vi.fn() + const nextCopy = vi.fn() + const { container, rerender } = render( + + ) + + expect(container.querySelector('code')?.textContent).toContain('select 1') + expect(screen.getByText('select', { selector: 'span' })).toHaveAttribute('style') + fireEvent.click(screen.getByRole('button', { name: 'Copy' })) + expect(copy).toHaveBeenCalledWith('select 1') + + rerender() + + expect(container.querySelector('code')?.textContent).toContain('select 2') + fireEvent.click(screen.getByRole('button', { name: 'Copied' })) + expect(nextCopy).toHaveBeenCalledWith('select 2') + expect(copy).toHaveBeenCalledTimes(1) + }) +}) diff --git a/packages/ui-patterns/src/CodeBlock/CodeBlock.tsx b/packages/ui-patterns/src/CodeBlock/CodeBlock.tsx index 536f4b65fa9..8ae54154398 100644 --- a/packages/ui-patterns/src/CodeBlock/CodeBlock.tsx +++ b/packages/ui-patterns/src/CodeBlock/CodeBlock.tsx @@ -5,7 +5,7 @@ import curl from 'highlightjs-curl' import { noop } from 'lodash' import { Check, Copy } from 'lucide-react' import { useTheme } from 'next-themes' -import { Children, ReactNode, useState } from 'react' +import { Children, memo, ReactNode, useState } from 'react' import { Light as SyntaxHighlighter, SyntaxHighlighterProps } from 'react-syntax-highlighter' import bash from 'react-syntax-highlighter/dist/cjs/languages/hljs/bash' import csharp from 'react-syntax-highlighter/dist/cjs/languages/hljs/csharp' @@ -32,6 +32,27 @@ import { Button, cn, copyToClipboard, FloatingPlate } from 'ui' import { monokaiCustomTheme } from './CodeBlock.utils' +SyntaxHighlighter.registerLanguage('js', js) +SyntaxHighlighter.registerLanguage('ts', ts) +SyntaxHighlighter.registerLanguage('py', py) +SyntaxHighlighter.registerLanguage('sql', sql) +SyntaxHighlighter.registerLanguage('bash', bash) +SyntaxHighlighter.registerLanguage('dart', dart) +SyntaxHighlighter.registerLanguage('csharp', csharp) +SyntaxHighlighter.registerLanguage('json', json) +SyntaxHighlighter.registerLanguage('kotlin', kotlin) +SyntaxHighlighter.registerLanguage('curl', curl) +SyntaxHighlighter.registerLanguage('http', http) +SyntaxHighlighter.registerLanguage('php', php) +SyntaxHighlighter.registerLanguage('python', python) +SyntaxHighlighter.registerLanguage('go', go) +SyntaxHighlighter.registerLanguage('pgsql', pgsql) +SyntaxHighlighter.registerLanguage('swift', swift) +SyntaxHighlighter.registerLanguage('html', xml) +SyntaxHighlighter.registerLanguage('toml', ini) +SyntaxHighlighter.registerLanguage('yaml', yaml) +SyntaxHighlighter.registerLanguage('markdown', markdown) + const codeBlockLangs = [ 'js', 'jsx', @@ -106,7 +127,7 @@ export interface CodeBlockProps { * @param {boolean} [props.focusable=true] - Whether the code block is focusable. When true, users can focus the code block to select text or use ⌘A (Cmd+A) to select all. This is so we don't need to load Monaco Editor. * @param {function} [props.handleCopy] - Optional override behaviour for copying value. For e.g if the code block contains obfuscated values, but the copy behaviour should reveal those values instead. */ -export const CodeBlock = ({ +export const CodeBlock = memo(function CodeBlock({ title, language, linesToHighlight = [], @@ -125,7 +146,7 @@ export const CodeBlock = ({ focusable = true, onCopyCallback = noop, handleCopy, -}: CodeBlockProps) => { +}: CodeBlockProps) { const { resolvedTheme } = useTheme() const isDarkTheme = resolvedTheme?.includes('dark')! const monokaiTheme = theme ?? monokaiCustomTheme(isDarkTheme) @@ -161,26 +182,6 @@ export const CodeBlock = ({ let lang = language ? language : className ? className.replace('language-', '') : 'js' // force jsx to be js highlighted if (lang === 'jsx') lang = 'js' - SyntaxHighlighter.registerLanguage('js', js) - SyntaxHighlighter.registerLanguage('ts', ts) - SyntaxHighlighter.registerLanguage('py', py) - SyntaxHighlighter.registerLanguage('sql', sql) - SyntaxHighlighter.registerLanguage('bash', bash) - SyntaxHighlighter.registerLanguage('dart', dart) - SyntaxHighlighter.registerLanguage('csharp', csharp) - SyntaxHighlighter.registerLanguage('json', json) - SyntaxHighlighter.registerLanguage('kotlin', kotlin) - SyntaxHighlighter.registerLanguage('curl', curl) - SyntaxHighlighter.registerLanguage('http', http) - SyntaxHighlighter.registerLanguage('php', php) - SyntaxHighlighter.registerLanguage('python', python) - SyntaxHighlighter.registerLanguage('go', go) - SyntaxHighlighter.registerLanguage('pgsql', pgsql) - SyntaxHighlighter.registerLanguage('swift', swift) - SyntaxHighlighter.registerLanguage('html', xml) - SyntaxHighlighter.registerLanguage('toml', ini) - SyntaxHighlighter.registerLanguage('yaml', yaml) - SyntaxHighlighter.registerLanguage('markdown', markdown) const large = false // don't show line numbers if bash == lang @@ -289,4 +290,4 @@ export const CodeBlock = ({ )} ) -} +}) From 46209b17036f8994c61dbe274ff1741b21e90e7a Mon Sep 17 00:00:00 2001 From: Saxon Fletcher Date: Fri, 25 Sep 2026 11:45:57 +1000 Subject: [PATCH 2/2] fix(studio): address assistant performance review --- .../ui/AIAssistantPanel/Message.Parts.tsx | 8 +- .../Message.Parts.utils.test.ts | 82 +++++++++++++++++++ .../AIAssistantPanel/Message.Parts.utils.ts | 27 ++++++ .../elements/Conversation.test.tsx | 42 ++++++++++ .../elements/Conversation.tsx | 10 ++- 5 files changed, 161 insertions(+), 8 deletions(-) create mode 100644 apps/studio/components/ui/AIAssistantPanel/Message.Parts.utils.test.ts create mode 100644 apps/studio/components/ui/AIAssistantPanel/Message.Parts.utils.ts create mode 100644 apps/studio/components/ui/AIAssistantPanel/elements/Conversation.test.tsx diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx index 3ecf20f7477..8eb9895e9e9 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx @@ -1,6 +1,5 @@ import { UIMessage as VercelMessage } from '@ai-sdk/react' import { type DynamicToolUIPart, type ReasoningUIPart, type TextUIPart, type ToolUIPart } from 'ai' -import isEqual from 'lodash/isEqual' import { BrainIcon, CheckIcon, Loader2 } from 'lucide-react' import { memo, type ReactNode } from 'react' import { cn } from 'ui' @@ -11,6 +10,7 @@ import { getManualToolApprovalHandlers } from './Confirm.utils' import { EdgeFunctionRenderer } from './EdgeFunctionRenderer' import { Tool } from './elements/Tool' import { useMessageActionsContext, useMessageInfoContext } from './Message.Context' +import { areMessagePartsEqual } from './Message.Parts.utils' import { deployEdgeFunctionInputSchema, deployEdgeFunctionOutputSchema, @@ -391,9 +391,5 @@ export const MessagePartSwitcher = memo( return {content} }, - (previous, next) => { - // The AI SDK clones every part of the active message on each stream update. - // Compare the part, not its identity, so completed tools retain their rendered subtree. - return isEqual(previous.part, next.part) - } + (previous, next) => areMessagePartsEqual(previous.part, next.part) ) diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.Parts.utils.test.ts b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.utils.test.ts new file mode 100644 index 00000000000..d6d9ab015aa --- /dev/null +++ b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.utils.test.ts @@ -0,0 +1,82 @@ +import type { DynamicToolUIPart, ToolUIPart } from 'ai' +import { describe, expect, it, vi } from 'vitest' + +import { areMessagePartsEqual } from './Message.Parts.utils' + +const completedTool = { + type: 'tool-execute_sql', + toolCallId: 'query-1', + state: 'output-available', + input: { sql: 'select 1' }, + output: [{ value: 1 }], +} satisfies ToolUIPart + +describe('areMessagePartsEqual', () => { + it.each(['static', 'dynamic'])('does not traverse finalized %s tool output', (kind) => { + const readRows = vi.fn(() => [{ value: 1 }]) + const output = () => ({ + get rows() { + return readRows() + }, + }) + const tool = + kind === 'static' + ? completedTool + : { ...completedTool, type: 'dynamic-tool' as const, toolName: 'query_logs' } + + expect(areMessagePartsEqual({ ...tool, output: output() }, { ...tool, output: output() })).toBe( + true + ) + expect(readRows).not.toHaveBeenCalled() + }) + + const changedTools: Array<[string, ToolUIPart | DynamicToolUIPart]> = [ + ['tool identity', { ...completedTool, toolCallId: 'query-2' }], + ['tool type', { ...completedTool, type: 'tool-query_logs' }], + ['input', { ...completedTool, input: { sql: 'select 2' } }], + ['approval', { ...completedTool, approval: { id: 'approval-1', approved: true } }], + ['metadata', { ...completedTool, toolMetadata: { title: 'Updated title' } }], + ['preliminary flag', { ...completedTool, preliminary: true }], + [ + 'error state', + { + type: 'tool-execute_sql', + toolCallId: 'query-1', + state: 'output-error', + input: completedTool.input, + errorText: 'Query failed', + }, + ], + ] + + it.each(changedTools)('rerenders when %s changes', (_label, next) => { + expect(areMessagePartsEqual(completedTool, next)).toBe(false) + }) + + it('updates preliminary results before the tool state changes', () => { + const previous = { ...completedTool, preliminary: true } + expect(areMessagePartsEqual(previous, { ...previous, output: [{ value: 2 }] })).toBe(false) + }) + + it('renders the final result after preliminary output', () => { + expect(areMessagePartsEqual({ ...completedTool, preliminary: true }, completedTool)).toBe(false) + }) + + it('checks live text and reasoning state', () => { + expect( + areMessagePartsEqual({ type: 'text', text: 'Hello' }, { type: 'text', text: 'Hello again' }) + ).toBe(false) + expect( + areMessagePartsEqual( + { type: 'reasoning', text: '', state: 'streaming' }, + { type: 'reasoning', text: '', state: 'done' } + ) + ).toBe(false) + }) + + it('keeps identical and cloned unchanged parts memoized', () => { + expect(areMessagePartsEqual(completedTool, completedTool)).toBe(true) + const text = { type: 'text' as const, text: 'Hello' } + expect(areMessagePartsEqual(text, { ...text })).toBe(true) + }) +}) diff --git a/apps/studio/components/ui/AIAssistantPanel/Message.Parts.utils.ts b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.utils.ts new file mode 100644 index 00000000000..3c3f5eb9cea --- /dev/null +++ b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.utils.ts @@ -0,0 +1,27 @@ +import { isToolUIPart, type UIMessage } from 'ai' +import isEqual from 'lodash/isEqual' + +type MessagePart = UIMessage['parts'][number] + +export function areMessagePartsEqual(previous: MessagePart, next: MessagePart): boolean { + if (previous === next) return true + if (previous.type !== next.type) return false + + if ( + isToolUIPart(previous) && + isToolUIPart(next) && + previous.state === 'output-available' && + next.state === 'output-available' && + !previous.preliminary && + !next.preliminary + ) { + // Final output is fixed for a tool call, but the SDK clones it on every text update. + // Keep checking identity, input, approval and metadata without walking result rows. + const { output: _previousOutput, ...previousFields } = previous + const { output: _nextOutput, ...nextFields } = next + return isEqual(previousFields, nextFields) + } + + // Preliminary output and live text/reasoning can still change without a state transition. + return isEqual(previous, next) +} diff --git a/apps/studio/components/ui/AIAssistantPanel/elements/Conversation.test.tsx b/apps/studio/components/ui/AIAssistantPanel/elements/Conversation.test.tsx new file mode 100644 index 00000000000..d3ba8734312 --- /dev/null +++ b/apps/studio/components/ui/AIAssistantPanel/elements/Conversation.test.tsx @@ -0,0 +1,42 @@ +import { render, screen } from '@testing-library/react' +import { createRef } from 'react' +import type { StickToBottomContext } from 'use-stick-to-bottom' +import { describe, expect, it } from 'vitest' + +import { Conversation, ConversationContent } from './Conversation' + +describe('ConversationContent', () => { + it('keeps scroll viewport classes separate from content classes and DOM attributes', () => { + const context = createRef() + const { rerender } = render( + + + Message + + + ) + + const viewport = context.current?.scrollRef.current + const content = context.current?.contentRef.current + expect(viewport).toHaveClass('scroll-pt-4', 'overscroll-y-contain') + expect(viewport).not.toHaveClass('space-y-4') + expect(content).toHaveClass('space-y-4') + expect(content).not.toHaveClass('scroll-pt-4') + expect(content).toHaveAttribute('id', 'messages') + expect(content).not.toHaveAttribute('scrollClassName') + + rerender( + + + {() => 'Updated message'} + + + ) + + expect(context.current?.scrollRef.current).toBe(viewport) + expect(context.current?.contentRef.current).toBe(content) + expect(viewport).toHaveClass('scroll-pt-8') + expect(viewport).not.toHaveClass('scroll-pt-4') + expect(screen.getByText('Updated message')).toBeInTheDocument() + }) +}) diff --git a/apps/studio/components/ui/AIAssistantPanel/elements/Conversation.tsx b/apps/studio/components/ui/AIAssistantPanel/elements/Conversation.tsx index 8c67f8735a2..3e59882dcc3 100644 --- a/apps/studio/components/ui/AIAssistantPanel/elements/Conversation.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/elements/Conversation.tsx @@ -7,7 +7,9 @@ import { StickToBottom, useStickToBottomContext } from 'use-stick-to-bottom' type ConversationProps = Omit, 'children'> & { children?: ReactNode } -type ConversationContentProps = ComponentProps +type ConversationContentProps = ComponentProps & { + scrollClassName?: string +} type ConversationScrollButtonProps = ComponentProps /** @@ -46,13 +48,17 @@ export const Conversation = ({ className, children, ...props }: ConversationProp export const ConversationContent = ({ className, + scrollClassName, children, ...props }: ConversationContentProps) => { const context = useStickToBottomContext() return ( -
+
{typeof children === 'function' ? children(context) : children}