mirror of
https://github.com/supabase/supabase.git
synced 2026-10-05 09:25:06 +03:00
feat(studio): group consecutive assistant tool calls (#50893)
<img width="913" height="572" alt="image" src="https://github.com/user-attachments/assets/e8112085-508a-49e4-abb1-7550248e611e" /> ## Problem A single Assistant response often produces 10+ reasoning and lookup rows ("Reasoned", "Ran search_docs", …). They push the answer down the chat, use raw tool names, and a fast tool call flashes past before the row goes back to "Thinking...". ## Solution Consecutive reasoning and lookup rows fold into one collapsible group. - **Running:** the header shows a tool only while it executes ("Checking policies in public..."). Between calls it reads "Thinking...", however long that lasts. Each header label stays up for at least 1 second, so quick calls no longer flash. - **Finished:** the header lists what the tools did, e.g. "Searched docs and checked policies", or "…, and 2 more". - **Expanded (any time):** every call is listed under a vertical rule. Rows still in progress shimmer, including several at once for parallel calls. ## How to test 1. Run `pnpm dev:studio` and open the Assistant on a project with a few tables. 2. Ask something that needs several lookups, e.g. "What RLS policies do I have and what do the docs recommend for them?" 3. While it streams, check the collapsed header: - It shows each tool while it runs, then goes back to "Thinking..." between calls. - Labels don't flash. Each stays up for about a second. - Only the shimmer marks progress, with no blinking cursor underneath. 4. Expand the group mid-stream. Rows read like "Checking policies in public..." rather than tool names, and only rows still in progress shimmer. 5. When it finishes, the header lists the actions ("Searched docs and checked policies") and stops shimmering. 6. Press Stop while a group is running. The unfinished row reads "Response interrupted" and stops spinning. 7. Reload the chat. Older groups show their collapsed summaries. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * AI assistant tool activity is grouped into collapsible sections with progress labels while work is underway and summaries when complete. * Expand grouped activity to review reasoning and tool details. Active tools and reasoning are highlighted, while completed reasoning without text is hidden. * Progress labels remain visible briefly during transitions, and active responses display a shimmer effect. * **Bug Fixes** * The loading indicator no longer appears while the assistant is processing a tool group. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Joshen Lim <joshenlimek@gmail.com>
This commit is contained in:
10 files changed
+618
-107
No files matched your search
@@ -29,6 +29,7 @@ import {
|
||||
ConversationScrollButton,
|
||||
} from './elements/Conversation'
|
||||
import { Message } from './Message'
|
||||
import { groupMessageParts } from './Message.Parts.utils'
|
||||
import { Markdown } from '@/components/interfaces/Markdown'
|
||||
import { useCheckOpenAIKeyQuery } from '@/data/ai/check-api-key-query'
|
||||
import { useRateMessageMutation } from '@/data/ai/rate-message-mutation'
|
||||
@@ -293,7 +294,14 @@ export const AssistantChat = ({
|
||||
(error.message?.includes('context_length_exceeded') ||
|
||||
error.message?.includes('exceeds the context window'))
|
||||
|
||||
const isTimedOut = !error && !isChatLoading && isTimedOutMessage(chatMessages.at(-1))
|
||||
const lastMessage = chatMessages.at(-1)
|
||||
// A running tool group shimmers already, so the cursor would be a second loading indicator
|
||||
const isToolGroupRunning =
|
||||
isChatLoading &&
|
||||
lastMessage?.role === 'assistant' &&
|
||||
groupMessageParts(lastMessage.parts).at(-1)?.type === 'tool-group'
|
||||
|
||||
const isTimedOut = !error && !isChatLoading && isTimedOutMessage(lastMessage)
|
||||
let displayError = IS_PLATFORM ? ASSISTANT_ERRORS['default'] : error
|
||||
if (isContextExceededError) displayError = ASSISTANT_ERRORS['context-exceeded']
|
||||
if (isTimedOut) displayError = { message: ASSISTANT_TIMEOUT_MESSAGE }
|
||||
@@ -620,7 +628,7 @@ export const AssistantChat = ({
|
||||
}
|
||||
/>
|
||||
)}
|
||||
{isChatLoading && (
|
||||
{isChatLoading && !isToolGroupRunning && (
|
||||
<motion.span
|
||||
animate={shouldReduceMotion ? { opacity: 1 } : { opacity: [1, 0] }}
|
||||
transition={
|
||||
|
||||
@@ -3,7 +3,8 @@ import { memo, type PropsWithChildren } from 'react'
|
||||
import { cn } from 'ui'
|
||||
|
||||
import { useMessageInfoContext } from './Message.Context'
|
||||
import { MessagePartSwitcher } from './Message.Parts'
|
||||
import { MessagePartSwitcher, MessagePartToolGroup } from './Message.Parts'
|
||||
import { groupMessageParts, isRunningToolCall } from './Message.Parts.utils'
|
||||
import { MessageMarkdown } from './MessageMarkdown'
|
||||
import { ProfileImage as ProfileImageDisplay } from '@/components/ui/ProfileImage'
|
||||
import { useProfileNameAndPicture } from '@/lib/profile'
|
||||
@@ -46,20 +47,35 @@ const MessageDisplayContent = memo(function MessageDisplayContent({
|
||||
}: {
|
||||
message: VercelMessage
|
||||
}) {
|
||||
const { id, isLoading, readOnly } = useMessageInfoContext()
|
||||
const { id, isLoading, isLastMessage, readOnly } = useMessageInfoContext()
|
||||
|
||||
const messageParts = message.parts
|
||||
const content =
|
||||
('content' in message && typeof message.content === 'string' && message.content.trim()) ||
|
||||
undefined
|
||||
|
||||
// The SDK exposes its mutable object on the first write, then publishes clones. Capture
|
||||
// state/text now so later mutations cannot change memoized parts' previous props.
|
||||
const items = groupMessageParts((messageParts ?? []).map((part) => ({ ...part })))
|
||||
const isStreaming = isLoading && !!isLastMessage
|
||||
|
||||
return (
|
||||
<div className="flex-1 min-w-0">
|
||||
{messageParts?.length > 0
|
||||
? messageParts.map((part: NonNullable<VercelMessage['parts'][number]>, idx) => {
|
||||
// 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 <MessagePartSwitcher key={idx} part={{ ...part }} />
|
||||
? items.map((item, idx) => {
|
||||
if (item.type === 'part') {
|
||||
return <MessagePartSwitcher key={item.partIndex} part={item.part} />
|
||||
}
|
||||
return (
|
||||
<MessagePartToolGroup
|
||||
key={`tool-group-${item.groupIndex}`}
|
||||
parts={item.parts}
|
||||
// The trailing group can still grow, and a parallel call can outlast a later block
|
||||
isRunning={
|
||||
isStreaming && (idx === items.length - 1 || item.parts.some(isRunningToolCall))
|
||||
}
|
||||
/>
|
||||
)
|
||||
})
|
||||
: content && (
|
||||
<div className="w-full max-w-3xl mx-auto">
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
import { act, screen } from '@testing-library/react'
|
||||
import userEvent from '@testing-library/user-event'
|
||||
import type { ToolUIPart } from 'ai'
|
||||
import { type PropsWithChildren } from 'react'
|
||||
import { describe, expect, it } from 'vitest'
|
||||
import { afterEach, describe, expect, it, vi } from 'vitest'
|
||||
|
||||
import { MessageProvider } from './Message.Context'
|
||||
import { MessagePartSwitcher } from './Message.Parts'
|
||||
import { MessagePartSwitcher, MessagePartToolGroup } from './Message.Parts'
|
||||
import type { CompactPart } from './Message.Parts.utils'
|
||||
import { customRender } from '@/tests/lib/custom-render'
|
||||
|
||||
type MessagePart = Parameters<typeof MessagePartSwitcher>[0]['part']
|
||||
@@ -117,3 +120,86 @@ describe('MessagePartSwitcher', () => {
|
||||
expect(container.querySelector('.animate-spin')).toBeNull()
|
||||
})
|
||||
})
|
||||
|
||||
describe('MessagePartToolGroup', () => {
|
||||
const reasoning = (text = 'Looking at the schema'): CompactPart => ({
|
||||
type: 'reasoning',
|
||||
state: 'done',
|
||||
text,
|
||||
})
|
||||
const streamingReasoning: CompactPart = { type: 'reasoning', state: 'streaming', text: '' }
|
||||
const tool: CompactPart = {
|
||||
type: 'tool-load_knowledge',
|
||||
toolCallId: 'knowledge-1',
|
||||
state: 'output-available',
|
||||
input: {},
|
||||
output: {},
|
||||
}
|
||||
const runningTool: CompactPart = {
|
||||
type: 'tool-load_knowledge',
|
||||
toolCallId: 'knowledge-1',
|
||||
state: 'input-available',
|
||||
input: {},
|
||||
}
|
||||
|
||||
afterEach(() => vi.useRealTimers())
|
||||
|
||||
// A group only runs while its message streams
|
||||
const toolGroup = (parts: CompactPart[], isRunning: boolean) => (
|
||||
<Provider isLoading={isRunning}>
|
||||
<MessagePartToolGroup parts={parts} isRunning={isRunning} />
|
||||
</Provider>
|
||||
)
|
||||
|
||||
it('folds its rows behind a summary until expanded', async () => {
|
||||
const { container } = customRender(toolGroup([reasoning(), tool], false))
|
||||
expect(container.querySelectorAll('.tool-item')).toHaveLength(0)
|
||||
|
||||
await userEvent.click(screen.getByRole('button', { name: 'Read up' }))
|
||||
|
||||
expect(container.querySelectorAll('.tool-item')).toHaveLength(2)
|
||||
expect(screen.getByText('Reasoned')).toBeInTheDocument()
|
||||
})
|
||||
|
||||
it('hides finished reasoning rows with nothing to expand', async () => {
|
||||
const { container } = customRender(toolGroup([reasoning(''), tool], false))
|
||||
await userEvent.click(screen.getByRole('button', { name: 'Read up' }))
|
||||
expect(container.querySelectorAll('.tool-item')).toHaveLength(1)
|
||||
})
|
||||
|
||||
it('holds each running header long enough to read', () => {
|
||||
vi.useFakeTimers()
|
||||
const { rerender } = customRender(toolGroup([streamingReasoning], true))
|
||||
const trigger = screen.getByRole('button', { name: 'Thinking...' })
|
||||
|
||||
rerender(toolGroup([reasoning(), runningTool], true))
|
||||
expect(trigger).toHaveAccessibleName('Thinking...')
|
||||
act(() => vi.advanceTimersByTime(1000))
|
||||
expect(trigger).toHaveAccessibleName('Reading up...')
|
||||
|
||||
// The call finished right away, but its label stays up before going back to thinking
|
||||
rerender(toolGroup([reasoning(), tool], true))
|
||||
expect(trigger).toHaveAccessibleName('Reading up...')
|
||||
act(() => vi.advanceTimersByTime(1000))
|
||||
expect(trigger).toHaveAccessibleName('Thinking...')
|
||||
|
||||
// The summary replaces the header as soon as the group stops
|
||||
rerender(toolGroup([reasoning(), tool], false))
|
||||
expect(trigger).toHaveAccessibleName('Read up')
|
||||
expect(trigger.querySelector('.shimmer')).toBeNull()
|
||||
})
|
||||
|
||||
it('shimmers every row still in progress', async () => {
|
||||
const { container, rerender } = customRender(
|
||||
toolGroup([reasoning(), runningTool, tool, runningTool], true)
|
||||
)
|
||||
await userEvent.click(screen.getByRole('button', { name: 'Reading up...' }))
|
||||
|
||||
const shimmeringRows = () =>
|
||||
[...container.querySelectorAll('.tool-item')].map((row) => !!row.querySelector('.shimmer'))
|
||||
expect(shimmeringRows()).toEqual([false, true, false, true])
|
||||
|
||||
rerender(toolGroup([reasoning(), tool, tool, tool, streamingReasoning], true))
|
||||
expect(shimmeringRows()).toEqual([false, false, false, false, true])
|
||||
})
|
||||
})
|
||||
@@ -1,22 +1,28 @@
|
||||
import { UIMessage as VercelMessage } from '@ai-sdk/react'
|
||||
import {
|
||||
isToolUIPart,
|
||||
type DynamicToolUIPart,
|
||||
type ReasoningUIPart,
|
||||
type TextUIPart,
|
||||
type ToolUIPart,
|
||||
} from 'ai'
|
||||
import { BrainIcon, CheckIcon, CircleStop, Loader2 } from 'lucide-react'
|
||||
import { isToolUIPart, type TextUIPart, type ToolUIPart } from 'ai'
|
||||
import { BrainIcon, CheckIcon, CircleStop, Loader2, XIcon } from 'lucide-react'
|
||||
import { memo, type ReactNode } from 'react'
|
||||
import { cn } from 'ui'
|
||||
import { Markdown } from 'ui-patterns/Markdown'
|
||||
|
||||
import { AssistantQueryCell } from './AssistantQueryCell'
|
||||
import { toAssistantQueryResult } from './AssistantQueryCell.utils'
|
||||
import { getManualToolApprovalHandlers } from './Confirm.utils'
|
||||
import { EdgeFunctionRenderer } from './EdgeFunctionRenderer'
|
||||
import { Tool } from './elements/Tool'
|
||||
import { ToolGroup } from './elements/ToolGroup'
|
||||
import { useMessageActionsContext, useMessageInfoContext } from './Message.Context'
|
||||
import { areMessagePartsEqual } from './Message.Parts.utils'
|
||||
import {
|
||||
areMessagePartsEqual,
|
||||
getCompactPartLabel,
|
||||
getCompactPartStatus,
|
||||
getToolGroupSummary,
|
||||
INTERRUPTED_LABEL,
|
||||
isCompactToolCall,
|
||||
isRunningToolCall,
|
||||
type CompactPart,
|
||||
type CompactPartStatus,
|
||||
} from './Message.Parts.utils'
|
||||
import {
|
||||
deployEdgeFunctionInputSchema,
|
||||
deployEdgeFunctionOutputSchema,
|
||||
@@ -27,6 +33,7 @@ import { MessagePartQueryLogs } from './MessagePartQueryLogs'
|
||||
import { NotebookProposalRenderer, type NotebookProposalMode } from './NotebookProposalRenderer'
|
||||
import { NotebookRunRenderer } from './NotebookRunRenderer'
|
||||
import { parseSupportRequestMessage, SupportRequestMessage } from './SupportRequestMessage'
|
||||
import { useMinimumDisplayTime } from '@/hooks/misc/useMinimumDisplayTime'
|
||||
|
||||
function MessagePartText({ textPart }: { textPart: TextUIPart }) {
|
||||
const { id, isLoading, readOnly, isUserMessage, state } = useMessageInfoContext()
|
||||
@@ -52,59 +59,33 @@ function MessagePartText({ textPart }: { textPart: TextUIPart }) {
|
||||
)
|
||||
}
|
||||
|
||||
function MessagePartDynamicTool({ toolPart }: { toolPart: DynamicToolUIPart }) {
|
||||
return (
|
||||
<Tool
|
||||
icon={
|
||||
toolPart.state === 'input-streaming' ? (
|
||||
<Loader2 strokeWidth={1.5} size={12} className="animate-spin" />
|
||||
) : (
|
||||
<CheckIcon strokeWidth={1.5} size={12} className="text-foreground-muted" />
|
||||
)
|
||||
}
|
||||
label={
|
||||
<div>
|
||||
{toolPart.state === 'input-streaming' ? 'Running ' : 'Ran '}
|
||||
<span className="text-foreground-lighter">{`${toolPart.toolName}`}</span>
|
||||
</div>
|
||||
}
|
||||
/>
|
||||
)
|
||||
const COMPACT_STATUS_ICONS: Record<CompactPartStatus, ReactNode> = {
|
||||
running: <Loader2 strokeWidth={1.5} size={12} className="animate-spin" />,
|
||||
done: <CheckIcon strokeWidth={1.5} size={12} className="text-foreground-muted" />,
|
||||
failed: <XIcon strokeWidth={1.5} size={12} className="text-destructive" />,
|
||||
}
|
||||
|
||||
function MessagePartTool({ toolPart }: { toolPart: ToolUIPart }) {
|
||||
return (
|
||||
<Tool
|
||||
icon={
|
||||
toolPart.state === 'input-streaming' ? (
|
||||
<Loader2 strokeWidth={1.5} size={12} className="animate-spin" />
|
||||
) : (
|
||||
<CheckIcon strokeWidth={1.5} size={12} className="text-foreground-muted" />
|
||||
)
|
||||
}
|
||||
label={
|
||||
<div>
|
||||
{toolPart.state === 'input-streaming' ? 'Running ' : 'Ran '}
|
||||
<span className="text-foreground-lighter">{`${toolPart.type.replace('tool-', '')}`}</span>
|
||||
</div>
|
||||
}
|
||||
/>
|
||||
)
|
||||
}
|
||||
function MessagePartCompact({ part, isActive }: { part: CompactPart; isActive?: boolean }) {
|
||||
const status = getCompactPartStatus(part)
|
||||
const isReasoning = part.type === 'reasoning'
|
||||
|
||||
function MessagePartReasoning({ reasoningPart }: { reasoningPart: ReasoningUIPart }) {
|
||||
return (
|
||||
<Tool
|
||||
isActive={isActive}
|
||||
icon={
|
||||
reasoningPart.state === 'streaming' ? (
|
||||
<Loader2 strokeWidth={1.5} size={12} className="animate-spin" />
|
||||
) : (
|
||||
isReasoning && status === 'done' ? (
|
||||
<BrainIcon strokeWidth={1.5} size={12} className="text-foreground-muted" />
|
||||
) : (
|
||||
COMPACT_STATUS_ICONS[status]
|
||||
)
|
||||
}
|
||||
label={reasoningPart.state === 'streaming' ? 'Thinking...' : 'Reasoned'}
|
||||
label={getCompactPartLabel(part)}
|
||||
>
|
||||
{reasoningPart.text}
|
||||
{isReasoning ? (
|
||||
<Markdown className="text-xs text-foreground-lighter [&>p]:m-0 flex flex-col gap-y-1">
|
||||
{part.text}
|
||||
</Markdown>
|
||||
) : undefined}
|
||||
</Tool>
|
||||
)
|
||||
}
|
||||
@@ -293,9 +274,7 @@ function MessagePartNotebookRun({ toolPart }: { toolPart: ToolUIPart }) {
|
||||
|
||||
const MessagePart = {
|
||||
Text: MessagePartText,
|
||||
Dynamic: MessagePartDynamicTool,
|
||||
Tool: MessagePartTool,
|
||||
Reasoning: MessagePartReasoning,
|
||||
Compact: MessagePartCompact,
|
||||
ExecuteSql: MessagePartExecuteSql,
|
||||
QueryLogs: MessagePartQueryLogs,
|
||||
DeployEdgeFunction: MessagePartDeployEdgeFunction,
|
||||
@@ -332,16 +311,15 @@ const isWideMessagePart = (part: NonNullable<VercelMessage['parts']>[number]) =>
|
||||
// Unlabelled code fences resolve to SQL in MessageMarkdown, too.
|
||||
(part.type === 'text' && /```(?:sql)?(?:\s|$)/i.test(part.text))
|
||||
|
||||
const isCompactToolPart = (part: NonNullable<VercelMessage['parts']>[number]) =>
|
||||
part.type === 'reasoning' ||
|
||||
(part.type === 'dynamic-tool' && part.toolName !== 'query_logs') ||
|
||||
part.type === 'tool-list_policies' ||
|
||||
part.type === 'tool-search_docs' ||
|
||||
part.type === 'tool-get_active_incidents' ||
|
||||
part.type === 'tool-load_knowledge'
|
||||
|
||||
export const MessagePartSwitcher = memo(
|
||||
function MessagePartSwitcher({ part }: { part: NonNullable<VercelMessage['parts']>[number] }) {
|
||||
function MessagePartSwitcher({
|
||||
part,
|
||||
isActive,
|
||||
}: {
|
||||
part: NonNullable<VercelMessage['parts']>[number]
|
||||
/** Marks the in-progress call within a running tool group. */
|
||||
isActive?: boolean
|
||||
}) {
|
||||
const { isLoading, isLastMessage } = useMessageInfoContext()
|
||||
const isActiveMessage = isLoading && isLastMessage
|
||||
// Compact rows and query_logs run on the server, so `input-available` means the tool never
|
||||
@@ -349,7 +327,7 @@ export const MessagePartSwitcher = memo(
|
||||
const isServerToolAwaitingOutput =
|
||||
isToolUIPart(part) &&
|
||||
part.state === 'input-available' &&
|
||||
(isCompactToolPart(part) ||
|
||||
(isCompactToolCall(part) ||
|
||||
part.type === 'tool-query_logs' ||
|
||||
(part.type === 'dynamic-tool' && part.toolName === 'query_logs'))
|
||||
const isIncompletePart =
|
||||
@@ -361,29 +339,23 @@ export const MessagePartSwitcher = memo(
|
||||
return (
|
||||
<Tool
|
||||
icon={<CircleStop strokeWidth={1.5} size={12} className="text-foreground-muted" />}
|
||||
label="Response interrupted"
|
||||
label={INTERRUPTED_LABEL}
|
||||
>
|
||||
{part.type === 'reasoning' ? part.text : undefined}
|
||||
</Tool>
|
||||
)
|
||||
}
|
||||
|
||||
// Tool rows depend on being direct siblings to share their compact spacing and dividers.
|
||||
if (part.type === 'reasoning' || (isToolUIPart(part) && isCompactToolCall(part))) {
|
||||
return <MessagePart.Compact part={part} isActive={isActive} />
|
||||
}
|
||||
|
||||
const content = (() => {
|
||||
switch (part.type) {
|
||||
case 'dynamic-tool': {
|
||||
if (part.toolName === 'query_logs') {
|
||||
return <MessagePart.QueryLogs toolPart={part} />
|
||||
}
|
||||
return <MessagePart.Dynamic toolPart={part} />
|
||||
return <MessagePart.QueryLogs toolPart={part} />
|
||||
}
|
||||
case 'tool-list_policies':
|
||||
case 'tool-search_docs':
|
||||
case 'tool-get_active_incidents':
|
||||
case 'tool-load_knowledge': {
|
||||
return <MessagePart.Tool toolPart={part} />
|
||||
}
|
||||
case 'reasoning':
|
||||
return <MessagePart.Reasoning reasoningPart={part} />
|
||||
case 'text':
|
||||
return <MessagePart.Text textPart={part} />
|
||||
|
||||
@@ -418,10 +390,42 @@ export const MessagePartSwitcher = memo(
|
||||
})()
|
||||
|
||||
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 <MessagePartContainer isWide={isWideMessagePart(part)}>{content}</MessagePartContainer>
|
||||
},
|
||||
(previous, next) => areMessagePartsEqual(previous.part, next.part)
|
||||
(previous, next) =>
|
||||
previous.isActive === next.isActive && areMessagePartsEqual(previous.part, next.part)
|
||||
)
|
||||
|
||||
// Long enough to read a short label before the next one replaces it
|
||||
const MIN_HEADER_DISPLAY_MS = 1000
|
||||
|
||||
export function MessagePartToolGroup({
|
||||
parts,
|
||||
isRunning,
|
||||
}: {
|
||||
parts: CompactPart[]
|
||||
isRunning: boolean
|
||||
}) {
|
||||
// A tool call leads the header only while it executes. The rest of the time the model is thinking.
|
||||
const runningIndex = useMinimumDisplayTime(
|
||||
parts.findLastIndex(isRunningToolCall),
|
||||
MIN_HEADER_DISPLAY_MS
|
||||
)
|
||||
const runningToolCall: CompactPart | undefined = parts[runningIndex]
|
||||
|
||||
let header = runningToolCall ? getCompactPartLabel(runningToolCall, 'running') : 'Thinking...'
|
||||
if (!isRunning) header = getToolGroupSummary(parts)
|
||||
|
||||
return (
|
||||
<ToolGroup label={header} isActive={isRunning}>
|
||||
{parts.map((part, idx) => {
|
||||
// Some models don't share their reasoning, leaving finished rows with nothing to expand
|
||||
if (part.type === 'reasoning' && part.state === 'done' && !part.text.trim()) return null
|
||||
|
||||
// Parallel calls can leave several rows in progress at once
|
||||
const isActive = isRunning && getCompactPartStatus(part) === 'running'
|
||||
return <MessagePartSwitcher key={idx} part={part} isActive={isActive} />
|
||||
})}
|
||||
</ToolGroup>
|
||||
)
|
||||
}
|
||||
@@ -1,7 +1,13 @@
|
||||
import type { DynamicToolUIPart, ToolUIPart } from 'ai'
|
||||
import type { DynamicToolUIPart, ToolUIPart, UIMessage } from 'ai'
|
||||
import { describe, expect, it, vi } from 'vitest'
|
||||
|
||||
import { areMessagePartsEqual } from './Message.Parts.utils'
|
||||
import {
|
||||
areMessagePartsEqual,
|
||||
getCompactPartLabel,
|
||||
getToolGroupSummary,
|
||||
groupMessageParts,
|
||||
type CompactPart,
|
||||
} from './Message.Parts.utils'
|
||||
|
||||
const completedTool = {
|
||||
type: 'tool-execute_sql',
|
||||
@@ -80,3 +86,132 @@ describe('areMessagePartsEqual', () => {
|
||||
expect(areMessagePartsEqual(text, { ...text })).toBe(true)
|
||||
})
|
||||
})
|
||||
|
||||
const reasoning = (state: 'streaming' | 'done' = 'done'): CompactPart => ({
|
||||
type: 'reasoning',
|
||||
state,
|
||||
text: 'Thinking about it',
|
||||
})
|
||||
const tool = (name: string, input: unknown = {}): CompactPart => ({
|
||||
type: `tool-${name}`,
|
||||
toolCallId: name,
|
||||
state: 'output-available',
|
||||
input,
|
||||
output: {},
|
||||
})
|
||||
const runningTool = (name: string, input: unknown = {}): CompactPart => ({
|
||||
type: `tool-${name}`,
|
||||
toolCallId: name,
|
||||
state: 'input-available',
|
||||
input,
|
||||
})
|
||||
const failedTool = (name: string): CompactPart => ({
|
||||
type: `tool-${name}`,
|
||||
toolCallId: name,
|
||||
state: 'output-error',
|
||||
input: {},
|
||||
errorText: 'Boom',
|
||||
})
|
||||
const mcpTool = (toolName: string): CompactPart => ({
|
||||
type: 'dynamic-tool',
|
||||
toolName,
|
||||
toolCallId: toolName,
|
||||
state: 'output-available',
|
||||
input: {},
|
||||
output: {},
|
||||
})
|
||||
const text = (value: string): UIMessage['parts'][number] => ({ type: 'text', text: value })
|
||||
|
||||
describe('groupMessageParts', () => {
|
||||
it('folds consecutive compact parts and starts a new group after each block', () => {
|
||||
const parts = [
|
||||
reasoning(),
|
||||
tool('search_docs'),
|
||||
text('Let me check'),
|
||||
mcpTool('list_tables'),
|
||||
mcpTool('query_logs'),
|
||||
tool('execute_sql'),
|
||||
reasoning(),
|
||||
]
|
||||
|
||||
expect(groupMessageParts(parts)).toEqual([
|
||||
{ type: 'tool-group', parts: [parts[0], parts[1]], groupIndex: 0 },
|
||||
{ type: 'part', part: parts[2], partIndex: 2 },
|
||||
{ type: 'tool-group', parts: [parts[3]], groupIndex: 1 },
|
||||
{ type: 'part', part: parts[4], partIndex: 4 },
|
||||
{ type: 'part', part: parts[5], partIndex: 5 },
|
||||
{ type: 'tool-group', parts: [parts[6]], groupIndex: 2 },
|
||||
])
|
||||
})
|
||||
|
||||
it('drops parts that render nothing so they do not split a group', () => {
|
||||
const [first, second, third] = [reasoning(), tool('search_docs'), reasoning()]
|
||||
const parts = [
|
||||
first,
|
||||
{ type: 'step-start' as const },
|
||||
second,
|
||||
tool('rename_chat'),
|
||||
text(' '),
|
||||
third,
|
||||
]
|
||||
|
||||
expect(groupMessageParts(parts)).toEqual([
|
||||
{ type: 'tool-group', parts: [first, second, third], groupIndex: 0 },
|
||||
])
|
||||
})
|
||||
})
|
||||
|
||||
describe('getCompactPartLabel', () => {
|
||||
it.each([
|
||||
{ part: reasoning('streaming'), label: 'Thinking...' },
|
||||
{ part: reasoning(), label: 'Reasoned' },
|
||||
{ part: runningTool('list_policies'), label: 'Checking policies...' },
|
||||
{ part: tool('list_policies'), label: 'Checked policies' },
|
||||
{ part: failedTool('list_policies'), label: 'Checking policies failed' },
|
||||
{
|
||||
part: runningTool('list_policies', { schemas: ['public'] }),
|
||||
label: 'Checking policies in public...',
|
||||
},
|
||||
{
|
||||
part: tool('list_policies', { schemas: ['public', 'auth'] }),
|
||||
label: 'Checked policies in public, auth',
|
||||
},
|
||||
{
|
||||
part: tool('search_docs', {
|
||||
graphql_query: '{ searchDocs(query: "row level security") { nodes { title } } }',
|
||||
}),
|
||||
label: 'Searched docs for "row level security"',
|
||||
},
|
||||
{ part: tool('search_docs', { graphql_query: '{ schema }' }), label: 'Searched docs' },
|
||||
{ part: tool('load_knowledge', { name: 'rls' }), label: 'Read up on RLS' },
|
||||
{ part: mcpTool('get_project_url'), label: 'Ran get project url' },
|
||||
])('labels "$label"', ({ part, label }) => {
|
||||
expect(getCompactPartLabel(part)).toBe(label)
|
||||
})
|
||||
})
|
||||
|
||||
describe('getToolGroupSummary', () => {
|
||||
it.each([
|
||||
{ parts: [reasoning(), tool('search_docs')], summary: 'Searched docs' },
|
||||
{
|
||||
parts: [tool('search_docs'), reasoning(), tool('search_docs'), tool('list_policies')],
|
||||
summary: 'Searched docs and checked policies',
|
||||
},
|
||||
{
|
||||
parts: [
|
||||
tool('search_docs'),
|
||||
tool('list_policies'),
|
||||
tool('list_reports'),
|
||||
mcpTool('list_tables'),
|
||||
],
|
||||
summary: 'Searched docs, checked policies, and 2 more',
|
||||
},
|
||||
{ parts: [tool('search_docs'), failedTool('list_policies')], summary: 'Searched docs' },
|
||||
// Nothing finished
|
||||
{ parts: [reasoning()], summary: 'Reasoned' },
|
||||
{ parts: [failedTool('list_policies')], summary: 'Checking policies failed' },
|
||||
{ parts: [reasoning(), runningTool('search_docs')], summary: 'Response interrupted' },
|
||||
])('summarizes "$summary"', ({ parts, summary }) => {
|
||||
expect(getToolGroupSummary(parts)).toBe(summary)
|
||||
})
|
||||
})
|
||||
@@ -1,7 +1,17 @@
|
||||
import { isToolUIPart, type UIMessage } from 'ai'
|
||||
import {
|
||||
getToolName,
|
||||
isToolUIPart,
|
||||
type DynamicToolUIPart,
|
||||
type ReasoningUIPart,
|
||||
type ToolUIPart,
|
||||
type UIMessage,
|
||||
} from 'ai'
|
||||
import isEqual from 'lodash/isEqual'
|
||||
|
||||
type MessagePart = UIMessage['parts'][number]
|
||||
type ToolPart = ToolUIPart | DynamicToolUIPart
|
||||
/** A one-line row that folds into a tool group: reasoning, or a lookup tool call. */
|
||||
export type CompactPart = ReasoningUIPart | ToolPart
|
||||
|
||||
export function areMessagePartsEqual(previous: MessagePart, next: MessagePart): boolean {
|
||||
if (previous === next) return true
|
||||
@@ -25,3 +35,184 @@ export function areMessagePartsEqual(previous: MessagePart, next: MessagePart):
|
||||
// Preliminary output and live text/reasoning can still change without a state transition.
|
||||
return isEqual(previous, next)
|
||||
}
|
||||
|
||||
type ToolLabels = {
|
||||
running: string
|
||||
done: string
|
||||
/** What the call looks at, from its input, e.g. "in public" */
|
||||
detail?: (input: unknown) => string | undefined
|
||||
}
|
||||
|
||||
const getField = (input: unknown, key: string): unknown =>
|
||||
typeof input === 'object' && input !== null ? Reflect.get(input, key) : undefined
|
||||
|
||||
function inSchemas(input: unknown) {
|
||||
const schemas = getField(input, 'schemas')
|
||||
return Array.isArray(schemas) && schemas.length > 0 ? `in ${schemas.join(', ')}` : undefined
|
||||
}
|
||||
|
||||
function forDocsQuery(input: unknown) {
|
||||
// e.g. { searchDocs(query: "row level security", limit: 5) { ... } }
|
||||
const query = String(getField(input, 'graphql_query')).match(
|
||||
/searchDocs\s*\(\s*query\s*:\s*"((?:[^"\\]|\\.)+)"/
|
||||
)?.[1]
|
||||
return query ? `for "${query}"` : undefined
|
||||
}
|
||||
|
||||
const KNOWLEDGE_TOPICS: Record<string, string> = {
|
||||
pg_best_practices: 'Postgres best practices',
|
||||
rls: 'RLS',
|
||||
storage: 'Storage',
|
||||
edge_functions: 'Edge Functions',
|
||||
realtime: 'Realtime',
|
||||
logs: 'logs',
|
||||
}
|
||||
|
||||
function onKnowledgeTopic(input: unknown) {
|
||||
const name = String(getField(input, 'name'))
|
||||
return Object.hasOwn(KNOWLEDGE_TOPICS, name) ? `on ${KNOWLEDGE_TOPICS[name]}` : undefined
|
||||
}
|
||||
|
||||
function forAdvisorType(input: unknown) {
|
||||
const type = getField(input, 'type')
|
||||
return type === 'security' || type === 'performance' ? `for ${type} issues` : undefined
|
||||
}
|
||||
|
||||
const TOOL_LABELS: Record<string, ToolLabels> = {
|
||||
search_docs: { running: 'Searching docs', done: 'Searched docs', detail: forDocsQuery },
|
||||
load_knowledge: { running: 'Reading up', done: 'Read up', detail: onKnowledgeTopic },
|
||||
get_active_incidents: { running: 'Checking Supabase status', done: 'Checked Supabase status' },
|
||||
list_policies: { running: 'Checking policies', done: 'Checked policies', detail: inSchemas },
|
||||
list_tables: { running: 'Listing tables', done: 'Listed tables', detail: inSchemas },
|
||||
list_extensions: { running: 'Listing extensions', done: 'Listed extensions' },
|
||||
list_edge_functions: { running: 'Listing Edge Functions', done: 'Listed Edge Functions' },
|
||||
list_branches: { running: 'Listing branches', done: 'Listed branches' },
|
||||
get_advisors: { running: 'Checking advisors', done: 'Checked advisors', detail: forAdvisorType },
|
||||
list_reports: { running: 'Listing reports', done: 'Listed reports' },
|
||||
get_report: { running: 'Reading report', done: 'Read report' },
|
||||
list_databases: { running: 'Listing databases', done: 'Listed databases' },
|
||||
list_notebooks: { running: 'Listing notebooks', done: 'Listed notebooks' },
|
||||
get_notebook: { running: 'Reading notebook', done: 'Read notebook' },
|
||||
// Self-hosted fallbacks
|
||||
getSchemaTables: { running: 'Listing tables', done: 'Listed tables', detail: inSchemas },
|
||||
getRlsKnowledge: { running: 'Checking policies', done: 'Checked policies', detail: inSchemas },
|
||||
getFunctions: {
|
||||
running: 'Listing database functions',
|
||||
done: 'Listed database functions',
|
||||
detail: inSchemas,
|
||||
},
|
||||
getEdgeFunctionKnowledge: {
|
||||
running: 'Reading up',
|
||||
done: 'Read up',
|
||||
detail: () => 'on Edge Functions',
|
||||
},
|
||||
}
|
||||
|
||||
function getToolLabels(toolName: string): ToolLabels {
|
||||
if (Object.hasOwn(TOOL_LABELS, toolName)) return TOOL_LABELS[toolName]
|
||||
const name = toolName.replaceAll('_', ' ')
|
||||
return { running: `Running ${name}`, done: `Ran ${name}` }
|
||||
}
|
||||
|
||||
const BLOCK_TOOLS = new Set([
|
||||
'execute_sql',
|
||||
'query_logs',
|
||||
'deploy_edge_function',
|
||||
'create_notebook',
|
||||
'update_notebook',
|
||||
'delete_notebook',
|
||||
'run_notebook',
|
||||
])
|
||||
|
||||
/** Lookups get a compact row. Studio's own tools need a label; any MCP tool gets a generic one. */
|
||||
export function isCompactToolCall(part: ToolPart): boolean {
|
||||
const toolName = getToolName(part)
|
||||
return part.type === 'dynamic-tool'
|
||||
? toolName !== 'query_logs'
|
||||
: Object.hasOwn(TOOL_LABELS, toolName)
|
||||
}
|
||||
|
||||
const isBlockPart = (part: MessagePart) =>
|
||||
part.type === 'text'
|
||||
? part.text.trim().length > 0
|
||||
: isToolUIPart(part) && BLOCK_TOOLS.has(getToolName(part))
|
||||
|
||||
export type MessagePartItem =
|
||||
| { type: 'part'; part: MessagePart; partIndex: number }
|
||||
| { type: 'tool-group'; parts: CompactPart[]; groupIndex: number }
|
||||
|
||||
/**
|
||||
* Folds each run of consecutive compact parts into a tool group. Parts that render nothing
|
||||
* (step markers, empty text, tools without UI) are dropped so they don't split a run.
|
||||
*/
|
||||
export function groupMessageParts(parts: MessagePart[]): MessagePartItem[] {
|
||||
const items: MessagePartItem[] = []
|
||||
let groupCount = 0
|
||||
|
||||
parts.forEach((part, partIndex) => {
|
||||
if (part.type === 'reasoning' || (isToolUIPart(part) && isCompactToolCall(part))) {
|
||||
const lastItem = items.at(-1)
|
||||
if (lastItem?.type === 'tool-group') lastItem.parts.push(part)
|
||||
else items.push({ type: 'tool-group', parts: [part], groupIndex: groupCount++ })
|
||||
} else if (isBlockPart(part)) {
|
||||
items.push({ type: 'part', part, partIndex })
|
||||
}
|
||||
})
|
||||
|
||||
return items
|
||||
}
|
||||
|
||||
export const INTERRUPTED_LABEL = 'Response interrupted'
|
||||
|
||||
/** `running` covers any unfinished part. Once its message stops, it renders as interrupted. */
|
||||
export type CompactPartStatus = 'running' | 'done' | 'failed'
|
||||
|
||||
export function getCompactPartStatus(part: CompactPart): CompactPartStatus {
|
||||
if (part.type === 'reasoning') return part.state === 'streaming' ? 'running' : 'done'
|
||||
if (part.state === 'output-available') return 'done'
|
||||
if (part.state === 'output-error' || part.state === 'output-denied') return 'failed'
|
||||
return 'running'
|
||||
}
|
||||
|
||||
export const isRunningToolCall = (part: CompactPart) =>
|
||||
part.type !== 'reasoning' && getCompactPartStatus(part) === 'running'
|
||||
|
||||
export function getCompactPartLabel(
|
||||
part: CompactPart,
|
||||
status = getCompactPartStatus(part)
|
||||
): string {
|
||||
if (part.type === 'reasoning') return status === 'running' ? 'Thinking...' : 'Reasoned'
|
||||
|
||||
const labels = getToolLabels(getToolName(part))
|
||||
const detail = labels.detail?.(part.input)
|
||||
const target = detail ? ` ${detail}` : ''
|
||||
if (status === 'running') return `${labels.running}${target}...`
|
||||
if (status === 'failed') return `${labels.running}${target} failed`
|
||||
return `${labels.done}${target}`
|
||||
}
|
||||
|
||||
const lowerFirst = (text: string) => text.charAt(0).toLowerCase() + text.slice(1)
|
||||
|
||||
/** A finished group's header: what its tool calls did, e.g. "Searched docs and checked policies". */
|
||||
export function getToolGroupSummary(parts: CompactPart[]): string {
|
||||
const actions = new Set(
|
||||
parts.flatMap((part) =>
|
||||
part.type !== 'reasoning' && getCompactPartStatus(part) === 'done'
|
||||
? [getToolLabels(getToolName(part)).done]
|
||||
: []
|
||||
)
|
||||
)
|
||||
const [first, second, ...rest] = actions
|
||||
|
||||
if (!first) {
|
||||
// Nothing finished: the group only reasoned, failed or was cut off
|
||||
const lastPart = parts.at(-1)
|
||||
if (!lastPart) return 'Reasoned'
|
||||
return getCompactPartStatus(lastPart) === 'running'
|
||||
? INTERRUPTED_LABEL
|
||||
: getCompactPartLabel(lastPart)
|
||||
}
|
||||
if (!second) return first
|
||||
if (rest.length === 0) return `${first} and ${lowerFirst(second)}`
|
||||
return `${first}, ${lowerFirst(second)}, and ${rest.length} more`
|
||||
}
|
||||
@@ -181,19 +181,36 @@ describe('assistant feed rendering', () => {
|
||||
it('finishes reasoning when the SDK mutates the first streamed part before publishing a snapshot', () => {
|
||||
const reasoning = {
|
||||
type: 'reasoning' as const,
|
||||
text: '',
|
||||
text: 'Looking at the schema',
|
||||
state: 'streaming' as 'streaming' | 'done',
|
||||
}
|
||||
const message: UIMessage = { id: 'reasoning-1', role: 'assistant', parts: [reasoning] }
|
||||
const { rerender } = render(<FeedMessage message={message} />)
|
||||
expect(screen.getByText('Thinking...')).toBeInTheDocument()
|
||||
// Expand the tool group so the reasoning row itself is rendered
|
||||
fireEvent.click(screen.getByRole('button', { name: 'Thinking...' }))
|
||||
// The group header and the reasoning row
|
||||
expect(screen.getAllByText('Thinking...')).toHaveLength(2)
|
||||
|
||||
// Chat.pushMessage exposes the initial object; subsequent replaceMessage calls clone it.
|
||||
reasoning.state = 'done'
|
||||
rerender(<FeedMessage message={structuredClone(message)} isLoading={false} />)
|
||||
|
||||
expect(screen.queryByText('Thinking...')).not.toBeInTheDocument()
|
||||
expect(screen.getByText('Reasoned')).toBeInTheDocument()
|
||||
expect(screen.getAllByText('Reasoned')).toHaveLength(2)
|
||||
})
|
||||
|
||||
it('keeps a tool group running while a call in it outlasts a later block', () => {
|
||||
const message: UIMessage = {
|
||||
id: 'parallel-1',
|
||||
role: 'assistant',
|
||||
parts: [
|
||||
{ type: 'tool-search_docs', toolCallId: 'docs-1', state: 'input-available', input: {} },
|
||||
{ type: 'text', text: 'Checking your schema' },
|
||||
],
|
||||
}
|
||||
render(<FeedMessage message={message} />)
|
||||
|
||||
expect(screen.getByRole('button', { name: 'Searching docs...' })).toBeInTheDocument()
|
||||
})
|
||||
|
||||
it('updates text when the SDK mutates the first streamed part', () => {
|
||||
|
||||
@@ -5,9 +5,11 @@ type ToolProps = PropsWithChildren<{
|
||||
className?: string
|
||||
label: ReactNode
|
||||
icon?: ReactNode
|
||||
/** Shimmers the label to mark the tool call that is currently in progress. */
|
||||
isActive?: boolean
|
||||
}>
|
||||
|
||||
export function Tool({ className, label, icon, children }: ToolProps) {
|
||||
export function Tool({ className, label, icon, isActive = false, children }: ToolProps) {
|
||||
const isCollapsible = !!children
|
||||
|
||||
return (
|
||||
@@ -20,22 +22,20 @@ export function Tool({ className, label, icon, children }: ToolProps) {
|
||||
className
|
||||
)}
|
||||
>
|
||||
<Collapsible>
|
||||
<Collapsible className="min-w-0 flex-1">
|
||||
<CollapsibleTrigger
|
||||
className={cn('flex items-center gap-2 w-full text-left')}
|
||||
className="flex items-center gap-2 w-full min-w-0 text-left [&>svg]:shrink-0"
|
||||
disabled={!children}
|
||||
>
|
||||
{icon}
|
||||
{typeof label === 'string' ? (
|
||||
<span className="text-foreground-lighter">{label}</span>
|
||||
) : (
|
||||
label
|
||||
)}
|
||||
<span className={cn('min-w-0 truncate text-foreground-lighter', isActive && 'shimmer')}>
|
||||
{label}
|
||||
</span>
|
||||
</CollapsibleTrigger>
|
||||
|
||||
{isCollapsible && (
|
||||
<CollapsibleContent
|
||||
className={cn('pl-6 py-2 text-xs leading-normal', 'max-h-64 overflow-y-auto')}
|
||||
className={cn('pl-5 py-2 text-xs leading-normal', 'max-h-64 overflow-y-auto')}
|
||||
>
|
||||
{children}
|
||||
</CollapsibleContent>
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
import { ChevronRight } from 'lucide-react'
|
||||
import type { PropsWithChildren, ReactNode } from 'react'
|
||||
import { cn, Collapsible, CollapsibleContent, CollapsibleTrigger } from 'ui'
|
||||
|
||||
type ToolGroupProps = PropsWithChildren<{
|
||||
label: ReactNode
|
||||
/** Shimmers the label to show work is in progress. */
|
||||
isActive?: boolean
|
||||
}>
|
||||
|
||||
/** Folds a run of `Tool` rows behind a single summary row. */
|
||||
export function ToolGroup({ label, isActive = false, children }: ToolGroupProps) {
|
||||
return (
|
||||
<Collapsible className="w-full max-w-3xl mx-auto my-4 first:mt-0 last:mb-0 text-foreground-lighter">
|
||||
<CollapsibleTrigger className="group/tool-group flex items-center gap-2 w-full py-2 text-left">
|
||||
<ChevronRight
|
||||
strokeWidth={1.5}
|
||||
size={12}
|
||||
className="shrink-0 transition-transform duration-200 ease-out group-data-[state=open]/tool-group:rotate-90"
|
||||
/>
|
||||
<span className={cn('min-w-0 truncate', isActive && 'shimmer')}>{label}</span>
|
||||
</CollapsibleTrigger>
|
||||
<CollapsibleContent className="overflow-hidden data-open:animate-collapsible-down data-closed:animate-collapsible-up">
|
||||
<div className="ml-1.5 pl-4 border-l border-muted">{children}</div>
|
||||
</CollapsibleContent>
|
||||
</Collapsible>
|
||||
)
|
||||
}
|
||||
@@ -0,0 +1,26 @@
|
||||
import { useEffect, useRef, useState } from 'react'
|
||||
|
||||
/**
|
||||
* Returns `value`, keeping each one for at least `minimumMs` so quick changes don't flash past.
|
||||
* Values that arrive in the meantime are skipped in favor of the latest.
|
||||
*/
|
||||
export function useMinimumDisplayTime<T>(value: T, minimumMs: number): T {
|
||||
const [displayed, setDisplayed] = useState(value)
|
||||
const shownAtRef = useRef(Date.now())
|
||||
|
||||
useEffect(() => {
|
||||
if (Object.is(displayed, value)) return
|
||||
|
||||
const remainingMs = minimumMs - (Date.now() - shownAtRef.current)
|
||||
const timeout = setTimeout(
|
||||
() => {
|
||||
shownAtRef.current = Date.now()
|
||||
setDisplayed(value)
|
||||
},
|
||||
Math.max(0, remainingMs)
|
||||
)
|
||||
return () => clearTimeout(timeout)
|
||||
}, [value, displayed, minimumMs])
|
||||
|
||||
return displayed
|
||||
}
|
||||
Reference in new issue
Block a user