From 1a5fd2ccb2c96931921dfddbc64639ddf946980c Mon Sep 17 00:00:00 2001 From: Olcan Date: Fri, 30 May 2025 13:59:05 -0700 Subject: [PATCH] add flags for markdown rendering and live updating to Tool to avoid special-casing shell tool by name, and open door for other tools to specify their rendering/updating (#629) --- packages/cli/src/ui/hooks/useToolScheduler.ts | 95 ++++++++----------- packages/server/src/tools/shell.ts | 2 + packages/server/src/tools/tools.ts | 14 +++ 3 files changed, 57 insertions(+), 54 deletions(-) diff --git a/packages/cli/src/ui/hooks/useToolScheduler.ts b/packages/cli/src/ui/hooks/useToolScheduler.ts index 7a5985c2..d4cb0d54 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useToolScheduler.ts @@ -288,41 +288,40 @@ export function useToolScheduler( const callId = t.request.callId; setToolCalls(setStatus(t.request.callId, 'executing')); - const updateOutput = - t.tool.name === 'execute_bash_command' - ? (output: string) => { - setPendingHistoryItem( - (prevItem: HistoryItemWithoutId | null) => { - if (prevItem?.type === 'tool_group') { - return { - ...prevItem, - tools: prevItem.tools.map( - (toolDisplay: IndividualToolCallDisplay) => - toolDisplay.callId === callId && - toolDisplay.status === ToolCallStatus.Executing - ? { - ...toolDisplay, - resultDisplay: output, - } - : toolDisplay, - ), - }; - } - return prevItem; - }, - ); - // Also update the toolCall itself so that mapToDisplay - // can pick up the live output if the item is not pending - // (e.g. if it's being re-rendered from history) - setToolCalls((prevToolCalls) => - prevToolCalls.map((tc) => - tc.request.callId === callId && tc.status === 'executing' - ? { ...tc, liveOutput: output } - : tc, - ), - ); - } - : undefined; + const updateOutput = t.tool.canUpdateOutput + ? (output: string) => { + setPendingHistoryItem( + (prevItem: HistoryItemWithoutId | null) => { + if (prevItem?.type === 'tool_group') { + return { + ...prevItem, + tools: prevItem.tools.map( + (toolDisplay: IndividualToolCallDisplay) => + toolDisplay.callId === callId && + toolDisplay.status === ToolCallStatus.Executing + ? { + ...toolDisplay, + resultDisplay: output, + } + : toolDisplay, + ), + }; + } + return prevItem; + }, + ); + // Also update the toolCall itself so that mapToDisplay + // can pick up the live output if the item is not pending + // (e.g. if it's being re-rendered from history) + setToolCalls((prevToolCalls) => + prevToolCalls.map((tc) => + tc.request.callId === callId && tc.status === 'executing' + ? { ...tc, liveOutput: output } + : tc, + ), + ); + } + : undefined; t.tool .execute(t.request.args, signal, updateOutput) @@ -541,18 +540,6 @@ export function mapToDisplay( ): HistoryItemToolGroup { const tools = Array.isArray(tool) ? tool : [tool]; const toolsDisplays = tools.map((t): IndividualToolCallDisplay => { - // Determine if markdown rendering should be skipped for this tool - let renderOutputAsMarkdown = true; // Default to true - if (t.status === 'error') { - // For errors, the tool object might not be available, so check t.request.name - if (t.request.name === 'execute_bash_command') { - renderOutputAsMarkdown = false; - } - } else if ('tool' in t && t.tool?.name === 'execute_bash_command') { - // For other statuses, check t.tool.name if tool exists - renderOutputAsMarkdown = false; - } - switch (t.status) { case 'success': return { @@ -562,7 +549,7 @@ export function mapToDisplay( resultDisplay: t.response.resultDisplay, status: mapStatus(t.status), confirmationDetails: undefined, - renderOutputAsMarkdown, + renderOutputAsMarkdown: t.tool.isOutputMarkdown, }; case 'error': return { @@ -572,7 +559,7 @@ export function mapToDisplay( resultDisplay: t.response.resultDisplay, status: mapStatus(t.status), confirmationDetails: undefined, - renderOutputAsMarkdown, + renderOutputAsMarkdown: false, }; case 'cancelled': return { @@ -582,7 +569,7 @@ export function mapToDisplay( resultDisplay: t.response.resultDisplay, status: mapStatus(t.status), confirmationDetails: undefined, - renderOutputAsMarkdown, + renderOutputAsMarkdown: t.tool.isOutputMarkdown, }; case 'awaiting_approval': return { @@ -592,7 +579,7 @@ export function mapToDisplay( resultDisplay: undefined, status: mapStatus(t.status), confirmationDetails: t.confirmationDetails, - renderOutputAsMarkdown, + renderOutputAsMarkdown: t.tool.isOutputMarkdown, }; case 'executing': return { @@ -602,7 +589,7 @@ export function mapToDisplay( resultDisplay: t.liveOutput ?? undefined, status: mapStatus(t.status), confirmationDetails: undefined, - renderOutputAsMarkdown, + renderOutputAsMarkdown: t.tool.isOutputMarkdown, }; case 'validating': // Add this case return { @@ -612,7 +599,7 @@ export function mapToDisplay( resultDisplay: undefined, status: mapStatus(t.status), confirmationDetails: undefined, - renderOutputAsMarkdown, + renderOutputAsMarkdown: t.tool.isOutputMarkdown, }; case 'scheduled': return { @@ -622,7 +609,7 @@ export function mapToDisplay( resultDisplay: undefined, status: mapStatus(t.status), confirmationDetails: undefined, - renderOutputAsMarkdown, + renderOutputAsMarkdown: t.tool.isOutputMarkdown, }; default: { // ensures every case is checked for above diff --git a/packages/server/src/tools/shell.ts b/packages/server/src/tools/shell.ts index 388c1543..4efc3500 100644 --- a/packages/server/src/tools/shell.ts +++ b/packages/server/src/tools/shell.ts @@ -42,6 +42,8 @@ export class ShellTool extends BaseTool { toolDisplayName, toolDescription, toolParameterSchema, + false, // output is not markdown + true, // output can be updated ); } diff --git a/packages/server/src/tools/tools.ts b/packages/server/src/tools/tools.ts index eb1da248..e5d0c7cf 100644 --- a/packages/server/src/tools/tools.ts +++ b/packages/server/src/tools/tools.ts @@ -33,6 +33,16 @@ export interface Tool< */ schema: FunctionDeclaration; + /** + * Whether the tool's output should be rendered as markdown + */ + isOutputMarkdown: boolean; + + /** + * Whether the tool supports live (streaming) output + */ + canUpdateOutput: boolean; + /** * Validates the parameters for the tool * Should be called from both `shouldConfirmExecute` and `execute` @@ -85,6 +95,8 @@ export abstract class BaseTool< * @param name Internal name of the tool (used for API calls) * @param displayName User-friendly display name of the tool * @param description Description of what the tool does + * @param isOutputMarkdown Whether the tool's output should be rendered as markdown + * @param canUpdateOutput Whether the tool supports live (streaming) output * @param parameterSchema JSON Schema defining the parameters */ constructor( @@ -92,6 +104,8 @@ export abstract class BaseTool< readonly displayName: string, readonly description: string, readonly parameterSchema: Record, + readonly isOutputMarkdown: boolean = true, + readonly canUpdateOutput: boolean = false, ) {} /**