Skip to content

Collapse tool groups with more than 5 tool calls #191

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
annzhang-db wants to merge 5 commits into
base: main
Choose a base branch
from
Open

Collapse tool groups with more than 5 tool calls #191

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
a28146d
update
annzhang-db Nov 6, 2025
ca676d1
Merge remote-tracking branch 'upstream/main'
annzhang-db Apr 16, 2026
c201bcb
Remove agent/v1/responses as a supported endpoint type
annzhang-db Apr 17, 2026
b0e0251
Group consecutive tool calls into a shared render block
annzhang-db Apr 16, 2026
5a3bd90
Collapse tool groups with more than 5 tool calls
annzhang-db Apr 17, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion dash-chatbot-app/model_serving_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ def _get_endpoint_task_type(endpoint_name: str) -> str:
def is_endpoint_supported(endpoint_name: str) -> bool:
"""Check if the endpoint has a supported task type."""
task_type = _get_endpoint_task_type(endpoint_name)
supported_task_types = ["agent/v1/chat", "agent/v2/chat", "llm/v1/chat"]
supported_task_types = ["agent/v1/chat", "agent/v2/chat", "llm/v1/chat", "agent/v1/responses"]
return task_type in supported_task_types

def _validate_endpoint_task_type(endpoint_name: str) -> None:
Expand Down
344 changes: 230 additions & 114 deletions e2e-chatbot-app-next/client/src/components/message.tsx
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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';
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -158,7 +169,22 @@ const PurePreviewMessage = ({
</div>
)}

{partSegments?.map((parts, index) => {
{renderBlocks.map((block) => {
if (block.kind === 'tool-group') {
return (
<MessageToolGroup
key={`tool-group-${block.startIndex}`}
tools={block.tools}
isLoading={isLoading}
submitApproval={submitApproval}
isSubmitting={isSubmitting}
pendingApprovalId={pendingApprovalId}
/>
);
}

const parts = block.parts;
const index = block.index;
const [part] = parts;
const { type } = part;
const key = `message-${message.id}-part-${index}`;
Expand Down Expand Up @@ -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 (
<McpTool key={toolCallId} defaultOpen={true}>
<McpToolHeader
serverName={mcpServerName}
toolName={toolName}
state={effectiveState}
approved={approved}
/>
<McpToolContent>
<McpToolInput input={input} />
{state === 'approval-requested' && (
<McpApprovalActions
onApprove={() =>
submitApproval({
approvalRequestId: toolCallId,
approve: true,
})
}
onDeny={() =>
submitApproval({
approvalRequestId: toolCallId,
approve: false,
})
}
isSubmitting={
isSubmitting && pendingApprovalId === toolCallId
}
/>
)}
{state === 'output-available' && output != null && (
<ToolOutput
output={
errorText ? (
<div className="rounded border p-2 text-red-500">
Error: {errorText}
</div>
) : (
<div className="whitespace-pre-wrap font-mono text-sm">
{typeof output === 'string'
? output
: JSON.stringify(output, null, 2)}
</div>
)
}
errorText={undefined}
/>
)}
</McpToolContent>
</McpTool>
);
}

// Render regular tool calls
return (
<Tool key={toolCallId} defaultOpen={true}>
<ToolHeader type={toolName} state={effectiveState} />
<ToolContent>
<ToolInput input={input} />
{state === 'output-available' && (
<ToolOutput
output={
errorText ? (
<div className="rounded border p-2 text-red-500">
Error: {errorText}
</div>
) : (
<div className="whitespace-pre-wrap font-mono text-sm">
{typeof output === 'string'
? output
: JSON.stringify(output, null, 2)}
</div>
)
}
errorText={undefined}
/>
)}
</ToolContent>
</Tool>
);
}
// dynamic-tool parts are rendered by MessageToolGroup above.

// Support for citations/annotations
if (type === 'source-url') {
Expand Down Expand Up @@ -417,6 +331,208 @@ export const PreviewMessage = memo(
},
);

type ChatPart = ChatMessage['parts'][number];
type ToolPart = Extract<ChatPart, { type: 'dynamic-tool' }>;

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<typeof useApproval>['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) => (
<ToolPartRenderer
key={tool.toolCallId}
part={tool}
isLoading={isLoading}
submitApproval={submitApproval}
isSubmitting={isSubmitting}
pendingApprovalId={pendingApprovalId}
/>
);

return (
<div
className={cn('flex flex-col gap-2', {
'rounded-md border border-border/60 bg-muted/20 p-2': isMultiple,
})}
data-testid={isMultiple ? 'tool-group' : undefined}
>
{visibleTools.map(renderTool)}
{shouldCollapse && (
<Collapsible className="group">
<CollapsibleContent>{hiddenTools.map(renderTool)}</CollapsibleContent>
<CollapsibleTrigger className="flex items-center gap-1 text-sm text-muted-foreground hover:text-foreground">
<ChevronDownIcon className="size-4 transition-transform group-data-[state=open]:rotate-180" />
<span className="group-data-[state=open]:hidden">
+{hiddenTools.length} more tool use(s)
</span>
<span className="hidden group-data-[state=open]:inline">
Show less
</span>
</CollapsibleTrigger>
</Collapsible>
)}
</div>
);
};

const ToolPartRenderer = ({
part,
isLoading,
submitApproval,
isSubmitting,
pendingApprovalId,
}: {
part: ToolPart;
isLoading: boolean;
submitApproval: ReturnType<typeof useApproval>['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 (
<McpTool defaultOpen={true}>
<McpToolHeader
serverName={mcpServerName}
toolName={toolName}
state={effectiveState}
approved={approved}
/>
<McpToolContent>
<McpToolInput input={input} />
{state === 'approval-requested' && (
<McpApprovalActions
onApprove={() =>
submitApproval({ approvalRequestId: toolCallId, approve: true })
}
onDeny={() =>
submitApproval({
approvalRequestId: toolCallId,
approve: false,
})
}
isSubmitting={isSubmitting && pendingApprovalId === toolCallId}
/>
)}
{state === 'output-available' && output != null && (
<ToolOutput
output={
errorText ? (
<div className="rounded border p-2 text-red-500">
Error: {errorText}
</div>
) : (
<div className="whitespace-pre-wrap font-mono text-sm">
{typeof output === 'string'
? output
: JSON.stringify(output, null, 2)}
</div>
)
}
errorText={undefined}
/>
)}
</McpToolContent>
</McpTool>
);
}

return (
<Tool key={toolCallId} defaultOpen={true}>
<ToolHeader type={toolName} state={effectiveState} />
<ToolContent>
<ToolInput input={input} />
{state === 'output-available' && (
<ToolOutput
output={
errorText ? (
<div className="rounded border p-2 text-red-500">
Error: {errorText}
</div>
) : (
<div className="whitespace-pre-wrap font-mono text-sm">
{typeof output === 'string'
? output
: JSON.stringify(output, null, 2)}
</div>
)
}
errorText={undefined}
/>
)}
</ToolContent>
</Tool>
);
};

export const AwaitingResponseMessage = () => {
const role = 'assistant';

Expand Down
Loading