diff --git a/e2e-chatbot-app-next/client/src/components/message.tsx b/e2e-chatbot-app-next/client/src/components/message.tsx index 2a296d7f..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'; @@ -105,6 +111,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 +169,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 +249,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 +331,208 @@ 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 TOOL_GROUP_COLLAPSE_THRESHOLD = 5; + +const MessageToolGroup = ({ + tools, + isLoading, + submitApproval, + isSubmitting, + pendingApprovalId, +}: { + tools: ToolPart[]; + isLoading: boolean; + submitApproval: ReturnType['submitApproval']; + isSubmitting: boolean; + 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 ( +
+ {visibleTools.map(renderTool)} + {shouldCollapse && ( + + {hiddenTools.map(renderTool)} + + + + +{hiddenTools.length} more tool use(s) + + + Show less + + + + )} +
+ ); +}; + +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';