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:
authored and GitHub committed 2026-09-28 16:07:20 +08:00
1 parent 93a032c90d
commit 7994191e49
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
}