Skip to content

Group consecutive tool calls into a shared render block #190

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 4 commits into
base: main
Choose a base branch
from
Open

Group consecutive tool calls into a shared render block #190

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 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
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
312 changes: 198 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
Expand Up @@ -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(
Expand Down Expand Up @@ -158,7 +163,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 +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 (
<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 +325,182 @@ 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 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;
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}
>
{tools.map((tool) => (
<ToolPartRenderer
key={tool.toolCallId}
part={tool}
isLoading={isLoading}
submitApproval={submitApproval}
isSubmitting={isSubmitting}
pendingApprovalId={pendingApprovalId}
/>
))}
</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
6 changes: 2 additions & 4 deletions e2e-chatbot-app-next/packages/ai-sdk-providers/src/request-context.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
*
* Context is injected when:
* 1. Using API_PROXY environment variable, OR
* 2. Endpoint task type is 'agent/v2/chat' or 'agent/v1/responses'
* 2. Endpoint task type is 'agent/v2/chat'
*
* @param endpointTask - The task type of the serving endpoint (optional)
* @returns Whether to inject context into requests
Expand All @@ -21,7 +21,5 @@ export function shouldInjectContextForEndpoint(
return true;
}

return (
endpointTask === 'agent/v2/chat' || endpointTask === 'agent/v1/responses'
);
return endpointTask === 'agent/v2/chat';
}
4 changes: 0 additions & 4 deletions e2e-chatbot-app-next/tests/ai-sdk-provider/request-context.test.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -26,10 +26,6 @@ test.describe("Request Context Utils", () => {
expect(shouldInjectContextForEndpoint("agent/v2/chat")).toBe(true);
});

test("returns true for agent/v1/responses endpoint task", () => {
expect(shouldInjectContextForEndpoint("agent/v1/responses")).toBe(true);
});

test("returns false for llm/v1/chat endpoint task", () => {
expect(shouldInjectContextForEndpoint("llm/v1/chat")).toBe(false);
});
Expand Down
4 changes: 2 additions & 2 deletions e2e-chatbot-app-next/tests/api-mocking/api-mock-handlers.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -323,13 +323,13 @@ export const handlers = [
}),

// Mock fetching endpoint details
// Returns agent/v1/responses to enable context injection testing
// Returns agent/v2/chat to enable context injection testing
// Includes auth_policy to simulate an OBO-enabled endpoint
http.get(/\/api\/2\.0\/serving-endpoints\/([^/]+)$/, ({ params }) => {
const endpointName = (params as Record<string, string>)[0] ?? '';
return HttpResponse.json({
name: endpointName || 'test-endpoint',
task: 'agent/v1/responses',
task: 'agent/v2/chat',
auth_policy: {
user_auth_policy: {
api_scopes: ['serving.serving-endpoints'],
Expand Down
Loading