From d0135231fb46b64c1def04ab9f3d8408adfd4e93 Mon Sep 17 00:00:00 2001 From: Saxon Fletcher Date: Fri, 25 Sep 2026 18:09:34 +1000 Subject: [PATCH] chore(studio): improve assistant feed performance (#50888) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ## Problem Assistant conversations with multiple query and Edge Function blocks repeatedly render expensive content while streaming. Scrolling past the feed boundary can also move the surrounding layout. ## Solution Memoize unchanged messages, blocks, and code highlighting; batch streaming UI updates; and skip off-screen query layout while keeping block state mounted. Preserve streamed status updates and contain scrolling in the message viewport. The changes are shared by Next and TanStack. ## Review instructions 1. Compare the base branch and this branch using the same saved conversation containing 10–20 query, result/chart, and Edge Function blocks. Keep the browser, viewport, and conversation identical. 2. In Chrome DevTools, record Performance with 4× CPU throttling while streaming a follow-up, typing in the composer, and scrolling through the feed. Compare scripting/layout time and long tasks. React DevTools Profiler should show unchanged completed blocks avoiding renders during subsequent text updates. 3. Scroll away from query blocks and return. Confirm results, display settings, selections, and controls retain their state. Run a read-only query such as `select 1` and check its results still update. 4. Confirm “Thinking…” finishes, Stop retains the latest streamed text, and approval/skip, copy, edit, and branch actions still work. Repeated scrolling at the feed boundary must leave the outer layout/composer stationary; jump-to-latest and following new messages should still work. 5. Repeat in both runtimes: `STUDIO_FRAMEWORK=next pnpm dev:studio` and `STUDIO_FRAMEWORK=tanstack pnpm dev:studio`. Also check the assistant sidebar, which shares the feed. ## Validation - 170 assistant/Explorer tests and one shared CodeBlock test passed; formatting and Studio lint passed (two existing warnings). - Browser checks covered both route entry points, viewport state/geometry, and scroll behavior. - Review fixes: 23 focused tests, lint, formatting, and full Studio typechecking passed. Full production builds were not verified. ## Checklist - [x] I have read [CONTRIBUTING.md](https://github.com/supabase/supabase/blob/master/CONTRIBUTING.md) - [x] No docs content changed; docs authoring skills are not applicable. ## Summary by CodeRabbit * **Performance** * Improved responsiveness in the AI assistant by reducing unnecessary updates while messages stream and conversation history is displayed. * Optimized query previews, message rendering, and code blocks to keep the interface smoother during use. * **Bug Fixes** * Improved handling of message edits and deletions during generation, and preserved the latest response when generation is stopped. * Improved conversation scrolling behavior while keeping conversation content and scroll areas working as expected. --- .../ui/AIAssistantPanel/AssistantChat.tsx | 42 ++-- .../AIAssistantPanel/AssistantQueryCell.tsx | 26 +- .../ui/AIAssistantPanel/Message.Display.tsx | 14 +- .../ui/AIAssistantPanel/Message.Parts.tsx | 110 ++++----- .../Message.Parts.utils.test.ts | 82 +++++++ .../AIAssistantPanel/Message.Parts.utils.ts | 27 +++ .../Message.performance.test.tsx | 229 ++++++++++++++++++ .../ui/AIAssistantPanel/Message.tsx | 71 ++++-- .../ui/AIAssistantPanel/MessageMarkdown.tsx | 4 +- .../elements/Conversation.test.tsx | 42 ++++ .../elements/Conversation.tsx | 28 ++- .../src/CodeBlock/CodeBlock.test.tsx | 26 ++ .../ui-patterns/src/CodeBlock/CodeBlock.tsx | 49 ++-- 13 files changed, 609 insertions(+), 141 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/Message.performance.test.tsx create mode 100644 apps/studio/components/ui/AIAssistantPanel/elements/Conversation.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..8eb9895e9e9 100644 --- a/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx +++ b/apps/studio/components/ui/AIAssistantPanel/Message.Parts.tsx @@ -1,7 +1,7 @@ import { UIMessage as VercelMessage } from '@ai-sdk/react' import { type DynamicToolUIPart, type ReasoningUIPart, type TextUIPart, type ToolUIPart } from 'ai' 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' @@ -10,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, @@ -333,63 +334,62 @@ 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) => 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/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.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 c2b1279eb20..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 /** @@ -20,7 +22,7 @@ const FADE_GUTTER = 'inset-x-7' export const Conversation = ({ className, children, ...props }: ConversationProps) => ( ) -export const ConversationContent = ({ className, ...props }: ConversationContentProps) => ( - -) +export const ConversationContent = ({ + className, + scrollClassName, + 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 = ({ )} ) -} +})