diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index be14b3b195..d30bc78e21 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -94,18 +94,19 @@ type ToolResponseWithParts = ToolCallResponseInfo & { llmContent?: PartListUnion; }; -interface BackgroundExecutionData { - pid?: number; - command?: string; - initialOutput?: string; -} - interface BackgroundedShellInfo { pid: number; command: string; initialOutput: string; } +interface BackgroundExecutionData { + executionId?: number; + pid?: number; + command?: string; + initialOutput?: string; +} + enum StreamProcessingStatus { Completed, UserCancelled, @@ -123,14 +124,33 @@ function isBackgroundExecutionData( if (typeof data !== 'object' || data === null) { return false; } - const d = data as Partial; + + const executionId = 'executionId' in data ? data.executionId : undefined; + const pid = 'pid' in data ? data.pid : undefined; + const command = 'command' in data ? data.command : undefined; + const initialOutput = + 'initialOutput' in data ? data.initialOutput : undefined; + return ( - (d.pid === undefined || typeof d.pid === 'number') && - (d.command === undefined || typeof d.command === 'string') && - (d.initialOutput === undefined || typeof d.initialOutput === 'string') + (executionId === undefined || typeof executionId === 'number') && + (pid === undefined || typeof pid === 'number') && + (command === undefined || typeof command === 'string') && + (initialOutput === undefined || typeof initialOutput === 'string') ); } +function getBackgroundExecutionId( + data: BackgroundExecutionData, +): number | undefined { + if (typeof data.executionId === 'number') { + return data.executionId; + } + if (typeof data.pid === 'number') { + return data.pid; + } + return undefined; +} + function getBackgroundedShellInfo( toolCall: TrackedCompletedToolCall | TrackedCancelledToolCall, ): BackgroundedShellInfo | undefined { @@ -141,13 +161,14 @@ function getBackgroundedShellInfo( const response = toolCall.response as ToolResponseWithParts; const rawData = response?.data; const data = isBackgroundExecutionData(rawData) ? rawData : undefined; + const executionId = data ? getBackgroundExecutionId(data) : undefined; - if (!data?.pid) { + if (!executionId) { return undefined; } return { - pid: data.pid, + pid: executionId, command: data.command ?? 'shell', initialOutput: data.initialOutput ?? '', }; diff --git a/packages/core/src/core/coreToolHookTriggers.ts b/packages/core/src/core/coreToolHookTriggers.ts index cbd90e8039..464cfc5f04 100644 --- a/packages/core/src/core/coreToolHookTriggers.ts +++ b/packages/core/src/core/coreToolHookTriggers.ts @@ -15,7 +15,6 @@ import type { import { ToolErrorType } from '../tools/tool-error.js'; import { debugLogger } from '../utils/debugLogger.js'; import type { ShellExecutionConfig } from '../index.js'; -import { ShellToolInvocation } from '../tools/shell.js'; import { DiscoveredMCPToolInvocation } from '../tools/mcp-tool.js'; /** @@ -26,7 +25,7 @@ import { DiscoveredMCPToolInvocation } from '../tools/mcp-tool.js'; * @returns MCP context if this is an MCP tool, undefined otherwise */ function extractMcpContext( - invocation: ShellToolInvocation | AnyToolInvocation, + invocation: AnyToolInvocation, config: Config, ): McpToolContext | undefined { if (!(invocation instanceof DiscoveredMCPToolInvocation)) { @@ -63,18 +62,18 @@ function extractMcpContext( * @param signal Abort signal for cancellation * @param liveOutputCallback Optional callback for live output updates * @param shellExecutionConfig Optional shell execution config - * @param setPidCallback Optional callback to set the PID for shell invocations + * @param setExecutionIdCallback Optional callback to set an execution ID for backgroundable invocations * @param config Config to look up MCP server details for hook context * @returns The tool result */ export async function executeToolWithHooks( - invocation: ShellToolInvocation | AnyToolInvocation, + invocation: AnyToolInvocation, toolName: string, signal: AbortSignal, tool: AnyDeclarativeTool, liveOutputCallback?: (outputChunk: ToolLiveOutput) => void, shellExecutionConfig?: ShellExecutionConfig, - setPidCallback?: (pid: number) => void, + setExecutionIdCallback?: (executionId: number) => void, config?: Config, originalRequestName?: string, ): Promise { @@ -154,22 +153,14 @@ export async function executeToolWithHooks( } } - // Execute the actual tool - let toolResult: ToolResult; - if (setPidCallback && invocation instanceof ShellToolInvocation) { - toolResult = await invocation.execute( - signal, - liveOutputCallback, - shellExecutionConfig, - setPidCallback, - ); - } else { - toolResult = await invocation.execute( - signal, - liveOutputCallback, - shellExecutionConfig, - ); - } + // Execute the actual tool. Tools that support backgrounding can optionally + // surface an execution ID via the callback. + const toolResult: ToolResult = await invocation.execute( + signal, + liveOutputCallback, + shellExecutionConfig, + setExecutionIdCallback, + ); // Append notification if parameters were modified if (inputWasModified) { diff --git a/packages/core/src/scheduler/tool-executor.test.ts b/packages/core/src/scheduler/tool-executor.test.ts index e744738341..b382b2b208 100644 --- a/packages/core/src/scheduler/tool-executor.test.ts +++ b/packages/core/src/scheduler/tool-executor.test.ts @@ -469,7 +469,7 @@ describe('ToolExecutor', () => { expect(result.status).toBe(CoreToolCallStatus.Success); }); - it('should report PID updates for shell tools', async () => { + it('should report execution ID updates for backgroundable tools', async () => { // 1. Setup ShellToolInvocation const messageBus = createMockMessageBus(); const shellInvocation = new ShellToolInvocation( @@ -480,7 +480,7 @@ describe('ToolExecutor', () => { // We need a dummy tool that matches the invocation just for structure const mockTool = new MockTool({ name: SHELL_TOOL_NAME }); - // 2. Mock executeToolWithHooks to trigger the PID callback + // 2. Mock executeToolWithHooks to trigger the execution ID callback const testPid = 12345; vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockImplementation( async ( @@ -490,13 +490,13 @@ describe('ToolExecutor', () => { _tool, _liveCb, _shellCfg, - setPidCallback, + setExecutionIdCallback, _config, _originalRequestName, ) => { - // Simulate the shell tool reporting a PID - if (setPidCallback) { - setPidCallback(testPid); + // Simulate the tool reporting an execution ID + if (setExecutionIdCallback) { + setExecutionIdCallback(testPid); } return { llmContent: 'done', returnDisplay: 'done' }; }, @@ -525,7 +525,7 @@ describe('ToolExecutor', () => { onUpdateToolCall, }); - // 4. Verify PID was reported + // 4. Verify execution ID was reported expect(onUpdateToolCall).toHaveBeenCalledWith( expect.objectContaining({ status: CoreToolCallStatus.Executing, diff --git a/packages/core/src/scheduler/tool-executor.ts b/packages/core/src/scheduler/tool-executor.ts index 8269f1fc41..c367d30d72 100644 --- a/packages/core/src/scheduler/tool-executor.ts +++ b/packages/core/src/scheduler/tool-executor.ts @@ -16,7 +16,6 @@ import { type ToolLiveOutput, } from '../index.js'; import { SHELL_TOOL_NAME } from '../tools/tool-names.js'; -import { ShellToolInvocation } from '../tools/shell.js'; import { DiscoveredMCPTool } from '../tools/mcp-tool.js'; import { executeToolWithHooks } from '../core/coreToolHookTriggers.js'; import { @@ -89,43 +88,29 @@ export class ToolExecutor { let completedToolCall: CompletedToolCall; try { - let promise: Promise; - if (invocation instanceof ShellToolInvocation) { - const setPidCallback = (pid: number) => { - const executingCall: ExecutingToolCall = { - ...call, - status: CoreToolCallStatus.Executing, - tool, - invocation, - pid, - startTime: 'startTime' in call ? call.startTime : undefined, - }; - onUpdateToolCall(executingCall); + const setExecutionIdCallback = (executionId: number) => { + const executingCall: ExecutingToolCall = { + ...call, + status: CoreToolCallStatus.Executing, + tool, + invocation, + pid: executionId, + startTime: 'startTime' in call ? call.startTime : undefined, }; - promise = executeToolWithHooks( - invocation, - toolName, - signal, - tool, - liveOutputCallback, - shellExecutionConfig, - setPidCallback, - this.config, - request.originalRequestName, - ); - } else { - promise = executeToolWithHooks( - invocation, - toolName, - signal, - tool, - liveOutputCallback, - shellExecutionConfig, - undefined, - this.config, - request.originalRequestName, - ); - } + onUpdateToolCall(executingCall); + }; + + const promise = executeToolWithHooks( + invocation, + toolName, + signal, + tool, + liveOutputCallback, + shellExecutionConfig, + setExecutionIdCallback, + this.config, + request.originalRequestName, + ); const toolResult: ToolResult = await promise; diff --git a/packages/core/src/tools/shell.ts b/packages/core/src/tools/shell.ts index 4ea83b0af4..701d2df713 100644 --- a/packages/core/src/tools/shell.ts +++ b/packages/core/src/tools/shell.ts @@ -18,6 +18,7 @@ import { Kind, type ToolInvocation, type ToolResult, + type BackgroundExecutionData, type ToolCallConfirmationDetails, type ToolExecuteConfirmationDetails, type PolicyUpdateOptions, @@ -150,7 +151,7 @@ export class ShellToolInvocation extends BaseToolInvocation< signal: AbortSignal, updateOutput?: (output: ToolLiveOutput) => void, shellExecutionConfig?: ShellExecutionConfig, - setPidCallback?: (pid: number) => void, + setExecutionIdCallback?: (executionId: number) => void, ): Promise { const strippedCommand = stripShellWrapper(this.params.command); @@ -281,8 +282,8 @@ export class ShellToolInvocation extends BaseToolInvocation< ); if (pid) { - if (setPidCallback) { - setPidCallback(pid); + if (setExecutionIdCallback) { + setExecutionIdCallback(pid); } // If the model requested to run in the background, do so after a short delay. @@ -324,7 +325,7 @@ export class ShellToolInvocation extends BaseToolInvocation< } } - let data: Record | undefined; + let data: BackgroundExecutionData | undefined; let llmContent = ''; let timeoutMessage = ''; @@ -346,6 +347,7 @@ export class ShellToolInvocation extends BaseToolInvocation< } else if (this.params.is_background || result.backgrounded) { llmContent = `Command moved to background (PID: ${result.pid}). Output hidden. Press Ctrl+B to view.`; data = { + executionId: result.pid, pid: result.pid, command: this.params.command, initialOutput: result.output, diff --git a/packages/core/src/tools/tools.ts b/packages/core/src/tools/tools.ts index 0a82cc1510..da0e03d10b 100644 --- a/packages/core/src/tools/tools.ts +++ b/packages/core/src/tools/tools.ts @@ -61,15 +61,57 @@ export interface ToolInvocation< * Executes the tool with the validated parameters. * @param signal AbortSignal for tool cancellation. * @param updateOutput Optional callback to stream output. + * @param setExecutionIdCallback Optional callback for tools that expose a background execution handle. * @returns Result of the tool execution. */ execute( signal: AbortSignal, updateOutput?: (output: ToolLiveOutput) => void, shellExecutionConfig?: ShellExecutionConfig, + setExecutionIdCallback?: (executionId: number) => void, ): Promise; } +/** + * Structured payload used by tools to surface background execution metadata to + * the CLI UI. + */ +export interface BackgroundExecutionData extends Record { + /** + * Neutral execution identifier for background lifecycle tracking. + */ + executionId?: number; + /** + * Backwards-compatible alias for executionId. + */ + pid?: number; + command?: string; + initialOutput?: string; +} + +export function isBackgroundExecutionData( + data: unknown, +): data is BackgroundExecutionData { + if (typeof data !== 'object' || data === null) { + return false; + } + + const value = data as Partial; + return ( + (value.executionId === undefined || typeof value.executionId === 'number') && + (value.pid === undefined || typeof value.pid === 'number') && + (value.command === undefined || typeof value.command === 'string') && + (value.initialOutput === undefined || + typeof value.initialOutput === 'string') + ); +} + +export function getBackgroundExecutionId( + data: BackgroundExecutionData, +): number | undefined { + return data.executionId ?? data.pid; +} + /** * Options for policy updates that can be customized by tool invocations. */