From bfeaac844186153698d3a7079b41214bbf1e4371 Mon Sep 17 00:00:00 2001 From: Olcan Date: Tue, 27 May 2025 15:40:18 -0700 Subject: [PATCH] live output from shell tool (#573) --- packages/cli/src/ui/hooks/useGeminiStream.ts | 34 ++++++----- packages/cli/src/ui/hooks/useToolScheduler.ts | 57 +++++++++++++++++-- packages/server/src/tools/shell.ts | 7 +++ packages/server/src/tools/tools.ts | 12 +++- 4 files changed, 87 insertions(+), 23 deletions(-) diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index f369d796..d91eea3d 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -70,21 +70,25 @@ export const useGeminiStream = ( const [pendingHistoryItemRef, setPendingHistoryItem] = useStateAndRef(null); const logger = useLogger(); - const [toolCalls, schedule, cancel] = useToolScheduler((tools) => { - if (tools.length) { - addItem(mapToDisplay(tools), Date.now()); - submitQuery( - tools - .filter( - (t) => - t.status === 'error' || - t.status === 'cancelled' || - t.status === 'success', - ) - .map((t) => t.response.responsePart), - ); - } - }, config); + const [toolCalls, schedule, cancel] = useToolScheduler( + (tools) => { + if (tools.length) { + addItem(mapToDisplay(tools), Date.now()); + submitQuery( + tools + .filter( + (t) => + t.status === 'error' || + t.status === 'cancelled' || + t.status === 'success', + ) + .map((t) => t.response.responsePart), + ); + } + }, + config, + setPendingHistoryItem, + ); const pendingToolCalls = useMemo( () => (toolCalls.length ? mapToDisplay(toolCalls) : undefined), [toolCalls], diff --git a/packages/cli/src/ui/hooks/useToolScheduler.ts b/packages/cli/src/ui/hooks/useToolScheduler.ts index 36493332..f1eee9fd 100644 --- a/packages/cli/src/ui/hooks/useToolScheduler.ts +++ b/packages/cli/src/ui/hooks/useToolScheduler.ts @@ -11,6 +11,7 @@ import { ToolConfirmationOutcome, Tool, ToolCallConfirmationDetails, + ToolResult, } from '@gemini-code/server'; import { Part } from '@google/genai'; import { useCallback, useEffect, useState } from 'react'; @@ -18,6 +19,7 @@ import { HistoryItemToolGroup, IndividualToolCallDisplay, ToolCallStatus, + HistoryItemWithoutId, } from '../types.js'; type ValidatingToolCall = { @@ -45,10 +47,11 @@ type SuccessfulToolCall = { response: ToolCallResponseInfo; }; -type ExecutingToolCall = { +export type ExecutingToolCall = { status: 'executing'; request: ToolCallRequestInfo; tool: Tool; + liveOutput?: string; }; type CancelledToolCall = { @@ -88,6 +91,9 @@ export type CompletedToolCall = export function useToolScheduler( onComplete: (tools: CompletedToolCall[]) => void, config: Config, + setPendingHistoryItem: React.Dispatch< + React.SetStateAction + >, ): [ToolCall[], ScheduleFn, CancelFn] { const [toolRegistry] = useState(() => config.getToolRegistry()); const [toolCalls, setToolCalls] = useState([]); @@ -224,9 +230,48 @@ export function useToolScheduler( .forEach((t) => { const callId = t.request.callId; setToolCalls(setStatus(t.request.callId, 'executing')); + + let accumulatedOutput = ''; + const onOutputChunk = + t.tool.name === 'execute_bash_command' + ? (chunk: string) => { + accumulatedOutput += chunk; + 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: accumulatedOutput, + } + : 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: accumulatedOutput } + : tc, + ), + ); + } + : undefined; + t.tool - .execute(t.request.args, signal) - .then((result) => { + .execute(t.request.args, signal, onOutputChunk) + .then((result: ToolResult) => { if (signal.aborted) { setToolCalls( setStatus(callId, 'cancelled', String(result.llmContent)), @@ -248,7 +293,7 @@ export function useToolScheduler( }; setToolCalls(setStatus(callId, 'success', response)); }) - .catch((e) => + .catch((e: Error) => setToolCalls( setStatus( callId, @@ -262,7 +307,7 @@ export function useToolScheduler( ); }); } - }, [toolCalls, toolRegistry, abortController.signal]); + }, [toolCalls, toolRegistry, abortController.signal, setPendingHistoryItem]); useEffect(() => { const allDone = toolCalls.every( @@ -480,7 +525,7 @@ export function mapToDisplay( callId: t.request.callId, name: t.tool.displayName, description: t.tool.getDescription(t.request.args), - resultDisplay: undefined, + resultDisplay: t.liveOutput ?? undefined, status: mapStatus(t.status), confirmationDetails: undefined, }; diff --git a/packages/server/src/tools/shell.ts b/packages/server/src/tools/shell.ts index 9e61717c..6ee36b8d 100644 --- a/packages/server/src/tools/shell.ts +++ b/packages/server/src/tools/shell.ts @@ -123,6 +123,7 @@ export class ShellTool extends BaseTool { async execute( params: ShellToolParams, abortSignal: AbortSignal, + onOutputChunk?: (chunk: string) => void, ): Promise { const validationError = this.validateToolParams(params); if (validationError) { @@ -157,6 +158,9 @@ export class ShellTool extends BaseTool { const str = data.toString(); stdout += str; output += str; + if (onOutputChunk) { + onOutputChunk(str); + } }); let stderr = ''; @@ -174,6 +178,9 @@ export class ShellTool extends BaseTool { } stderr += str; output += str; + if (onOutputChunk) { + onOutputChunk(str); + } }); let error: Error | null = null; diff --git a/packages/server/src/tools/tools.ts b/packages/server/src/tools/tools.ts index 2f5a4095..c57bbd39 100644 --- a/packages/server/src/tools/tools.ts +++ b/packages/server/src/tools/tools.ts @@ -64,7 +64,11 @@ export interface Tool< * @param params Parameters for the tool execution * @returns Result of the tool execution */ - execute(params: TParams, signal: AbortSignal): Promise; + execute( + params: TParams, + signal: AbortSignal, + onOutputChunk?: (chunk: string) => void, + ): Promise; } /** @@ -144,7 +148,11 @@ export abstract class BaseTool< * @param signal AbortSignal for tool cancellation * @returns Result of the tool execution */ - abstract execute(params: TParams, signal: AbortSignal): Promise; + abstract execute( + params: TParams, + signal: AbortSignal, + onOutputChunk?: (chunk: string) => void, + ): Promise; } export interface ToolResult {