live output from shell tool (#573)

packages/cli/src/ui/hooks/useGeminiStream.ts

@@ -70,21 +70,25 @@ export const useGeminiStream = (
   const [pendingHistoryItemRef, setPendingHistoryItem] =
     useStateAndRef<HistoryItemWithoutId | null>(null);
   const logger = useLogger();
-  const [toolCalls, schedule, cancel] = useToolScheduler((tools) => {
+  const [toolCalls, schedule, cancel] = useToolScheduler(
-    if (tools.length) {
+    (tools) => {
-      addItem(mapToDisplay(tools), Date.now());
+      if (tools.length) {
-      submitQuery(
+        addItem(mapToDisplay(tools), Date.now());
-        tools
+        submitQuery(
-          .filter(
+          tools
-            (t) =>
+            .filter(
-              t.status === 'error' ||
+              (t) =>
-              t.status === 'cancelled' ||
+                t.status === 'error' ||
-              t.status === 'success',
+                t.status === 'cancelled' ||
-          )
+                t.status === 'success',
-          .map((t) => t.response.responsePart),
+            )
-      );
+            .map((t) => t.response.responsePart),
-    }
+        );
-  }, config);
+      }
+    },
+    config,
+    setPendingHistoryItem,
+  );
   const pendingToolCalls = useMemo(
     () => (toolCalls.length ? mapToDisplay(toolCalls) : undefined),
     [toolCalls],

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<HistoryItemWithoutId | null>
+  >,
 ): [ToolCall[], ScheduleFn, CancelFn] {
   const [toolRegistry] = useState(() => config.getToolRegistry());
   const [toolCalls, setToolCalls] = useState<ToolCall[]>([]);
@@ -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)
+            .execute(t.request.args, signal, onOutputChunk)
-            .then((result) => {
+            .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,
         };

packages/server/src/tools/shell.ts

@@ -123,6 +123,7 @@ export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
   async execute(
     params: ShellToolParams,
     abortSignal: AbortSignal,
+    onOutputChunk?: (chunk: string) => void,
   ): Promise<ToolResult> {
     const validationError = this.validateToolParams(params);
     if (validationError) {
@@ -157,6 +158,9 @@ export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
       const str = data.toString();
       stdout += str;
       output += str;
+      if (onOutputChunk) {
+        onOutputChunk(str);
+      }
     });
 
     let stderr = '';
@@ -174,6 +178,9 @@ export class ShellTool extends BaseTool<ShellToolParams, ToolResult> {
       }
       stderr += str;
       output += str;
+      if (onOutputChunk) {
+        onOutputChunk(str);
+      }
     });
 
     let error: Error | null = null;

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<TResult>;
+  execute(
+    params: TParams,
+    signal: AbortSignal,
+    onOutputChunk?: (chunk: string) => void,
+  ): Promise<TResult>;
 }
 
 /**
@@ -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<TResult>;
+  abstract execute(
+    params: TParams,
+    signal: AbortSignal,
+    onOutputChunk?: (chunk: string) => void,
+  ): Promise<TResult>;
 }
 
 export interface ToolResult {