From 491145fa41e3a24b6dc2da4cdaef7a475976614b Mon Sep 17 00:00:00 2001 From: Ann Zhang Date: Thu, 16 Apr 2026 16:59:48 -0700 Subject: [PATCH 1/2] Group consecutive tool calls into a shared render block Co-authored-by: Ann Zhang --- .../client/src/components/message.tsx | 312 +++++++++++------- 1 file changed, 198 insertions(+), 114 deletions(-) diff --git a/e2e-chatbot-app-next/client/src/components/message.tsx b/e2e-chatbot-app-next/client/src/components/message.tsx index 2a296d7f..198fa8ae 100644 --- a/e2e-chatbot-app-next/client/src/components/message.tsx +++ b/e2e-chatbot-app-next/client/src/components/message.tsx @@ -105,6 +105,11 @@ const PurePreviewMessage = ({ [message.parts], ); + const renderBlocks = React.useMemo( + () => groupConsecutiveToolSegments(partSegments), + [partSegments], + ); + // Check if message only contains non-OAuth errors (no other content) const hasOnlyErrors = React.useMemo(() => { const nonErrorParts = message.parts.filter( @@ -158,7 +163,22 @@ const PurePreviewMessage = ({ )} - {partSegments?.map((parts, index) => { + {renderBlocks.map((block) => { + if (block.kind === 'tool-group') { + return ( + + ); + } + + const parts = block.parts; + const index = block.index; const [part] = parts; const { type } = part; const key = `message-${message.id}-part-${index}`; @@ -223,119 +243,7 @@ const PurePreviewMessage = ({ } } - // Render Databricks tool calls and results - if (part.type === `dynamic-tool`) { - const { toolCallId, input, state, errorText, output, toolName } = - part; - - // Check if this is an MCP tool call by looking for approvalRequestId in metadata - // This works across all states (approval-requested, approval-denied, output-available) - const isMcpApproval = - part.callProviderMetadata?.databricks?.approvalRequestId != - null; - const mcpServerName = - part.callProviderMetadata?.databricks?.mcpServerName?.toString(); - - // Extract approval outcome for 'approval-responded' state - // When addToolApprovalResponse is called, AI SDK sets the `approval` property - // on the tool-call part and changes state to 'approval-responded' - const approved: boolean | undefined = - 'approval' in part ? part.approval?.approved : undefined; - - // When approved but only have approval status (not actual output), show as input-available - const effectiveState: ToolState = (() => { - if ( - part.providerExecuted && - !isLoading && - state === 'input-available' - ) { - return 'output-available'; - } - return state; - })(); - - // Render MCP tool calls with special styling - if (isMcpApproval) { - return ( - - - - - {state === 'approval-requested' && ( - - submitApproval({ - approvalRequestId: toolCallId, - approve: true, - }) - } - onDeny={() => - submitApproval({ - approvalRequestId: toolCallId, - approve: false, - }) - } - isSubmitting={ - isSubmitting && pendingApprovalId === toolCallId - } - /> - )} - {state === 'output-available' && output != null && ( - - Error: {errorText} - - ) : ( -
- {typeof output === 'string' - ? output - : JSON.stringify(output, null, 2)} -
- ) - } - errorText={undefined} - /> - )} - - - ); - } - - // Render regular tool calls - return ( - - - - - {state === 'output-available' && ( - - Error: {errorText} - - ) : ( -
- {typeof output === 'string' - ? output - : JSON.stringify(output, null, 2)} -
- ) - } - errorText={undefined} - /> - )} - - - ); - } + // dynamic-tool parts are rendered by MessageToolGroup above. // Support for citations/annotations if (type === 'source-url') { @@ -417,6 +325,182 @@ export const PreviewMessage = memo( }, ); +type ChatPart = ChatMessage['parts'][number]; +type ToolPart = Extract; + +type RenderBlock = + | { kind: 'segment'; parts: ChatPart[]; index: number } + | { kind: 'tool-group'; tools: ToolPart[]; startIndex: number }; + +const groupConsecutiveToolSegments = ( + partSegments: ChatPart[][], +): RenderBlock[] => { + const blocks: RenderBlock[] = []; + let i = 0; + while (i < partSegments.length) { + const segment = partSegments[i]; + const firstPart = segment[0]; + if (firstPart?.type === 'dynamic-tool') { + const startIndex = i; + const tools: ToolPart[] = [firstPart as ToolPart]; + i++; + while ( + i < partSegments.length && + partSegments[i][0]?.type === 'dynamic-tool' + ) { + tools.push(partSegments[i][0] as ToolPart); + i++; + } + blocks.push({ kind: 'tool-group', tools, startIndex }); + } else { + blocks.push({ kind: 'segment', parts: segment, index: i }); + i++; + } + } + return blocks; +}; + +const MessageToolGroup = ({ + tools, + isLoading, + submitApproval, + isSubmitting, + pendingApprovalId, +}: { + tools: ToolPart[]; + isLoading: boolean; + submitApproval: ReturnType['submitApproval']; + isSubmitting: boolean; + pendingApprovalId: string | null; +}) => { + const isMultiple = tools.length > 1; + return ( +
+ {tools.map((tool) => ( + + ))} +
+ ); +}; + +const ToolPartRenderer = ({ + part, + isLoading, + submitApproval, + isSubmitting, + pendingApprovalId, +}: { + part: ToolPart; + isLoading: boolean; + submitApproval: ReturnType['submitApproval']; + isSubmitting: boolean; + pendingApprovalId: string | null; +}) => { + const { toolCallId, input, state, errorText, output, toolName } = part; + + const isMcpApproval = + part.callProviderMetadata?.databricks?.approvalRequestId != null; + const mcpServerName = + part.callProviderMetadata?.databricks?.mcpServerName?.toString(); + + const approved: boolean | undefined = + 'approval' in part ? part.approval?.approved : undefined; + + const effectiveState: ToolState = (() => { + if (part.providerExecuted && !isLoading && state === 'input-available') { + return 'output-available'; + } + return state; + })(); + + if (isMcpApproval) { + return ( + + + + + {state === 'approval-requested' && ( + + submitApproval({ approvalRequestId: toolCallId, approve: true }) + } + onDeny={() => + submitApproval({ + approvalRequestId: toolCallId, + approve: false, + }) + } + isSubmitting={isSubmitting && pendingApprovalId === toolCallId} + /> + )} + {state === 'output-available' && output != null && ( + + Error: {errorText} + + ) : ( +
+ {typeof output === 'string' + ? output + : JSON.stringify(output, null, 2)} +
+ ) + } + errorText={undefined} + /> + )} + + + ); + } + + return ( + + + + + {state === 'output-available' && ( + + Error: {errorText} + + ) : ( +
+ {typeof output === 'string' + ? output + : JSON.stringify(output, null, 2)} +
+ ) + } + errorText={undefined} + /> + )} + + + ); +}; + export const AwaitingResponseMessage = () => { const role = 'assistant'; From 93767cdfd13787427b10ba486e162a98cec23fee Mon Sep 17 00:00:00 2001 From: Ann Zhang Date: Thu, 16 Apr 2026 17:00:08 -0700 Subject: [PATCH 2/2] Collapse tool groups with more than 5 tool calls Co-authored-by: Ann Zhang --- .../client/src/components/message.tsx | 52 +++++++++++++++---- 1 file changed, 42 insertions(+), 10 deletions(-) diff --git a/e2e-chatbot-app-next/client/src/components/message.tsx b/e2e-chatbot-app-next/client/src/components/message.tsx index 198fa8ae..c533eaba 100644 --- a/e2e-chatbot-app-next/client/src/components/message.tsx +++ b/e2e-chatbot-app-next/client/src/components/message.tsx @@ -1,4 +1,10 @@ import React, { memo, useState } from 'react'; +import { + Collapsible, + CollapsibleContent, + CollapsibleTrigger, +} from '@/components/ui/collapsible'; +import { ChevronDownIcon } from 'lucide-react'; import { AnimatedAssistantIcon } from './animation-assistant-icon'; import { Response } from './elements/response'; import { MessageContent } from './elements/message'; @@ -360,6 +366,8 @@ const groupConsecutiveToolSegments = ( return blocks; }; +const TOOL_GROUP_COLLAPSE_THRESHOLD = 5; + const MessageToolGroup = ({ tools, isLoading, @@ -374,6 +382,25 @@ const MessageToolGroup = ({ pendingApprovalId: string | null; }) => { const isMultiple = tools.length > 1; + const shouldCollapse = tools.length > TOOL_GROUP_COLLAPSE_THRESHOLD; + const visibleTools = shouldCollapse + ? tools.slice(0, TOOL_GROUP_COLLAPSE_THRESHOLD) + : tools; + const hiddenTools = shouldCollapse + ? tools.slice(TOOL_GROUP_COLLAPSE_THRESHOLD) + : []; + + const renderTool = (tool: ToolPart) => ( + + ); + return (
- {tools.map((tool) => ( - - ))} + {visibleTools.map(renderTool)} + {shouldCollapse && ( + + {hiddenTools.map(renderTool)} + + + + +{hiddenTools.length} more tool use(s) + + + Show less + + + + )}
); };