jcarr
/
gemini-cli
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

live output from shell tool (#573)

This commit is contained in:
Olcan 2025-05-27 15:40:18 -07:00 committed by GitHub
parent 0d5f7686d7
commit bfeaac8441
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
4 changed files with 87 additions and 23 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
packages/cli/src/ui/hooks/useGeminiStream.ts
View File

@ -70,21 +70,25 @@ export const useGeminiStream = (
const [pendingHistoryItemRef, setPendingHistoryItem] =
useStateAndRef<HistoryItemWithoutId | null>(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],

57
packages/cli/src/ui/hooks/useToolScheduler.ts
View File

@ -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)
.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,
};

7
packages/server/src/tools/shell.ts
View File

@ -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;

12
packages/server/src/tools/tools.ts
View File

@ -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 {