Compare commits

...

20 Commits

Author SHA1 Message Date
Adam Weidman 3389a38d86 refactor: use pid as sole background execution data field
Align with PR1's standardization on pid as the canonical field name.
2026-03-09 01:12:22 -04:00
Adam Weidman bc76cbeb50 test: decouple remote background test from shell service 2026-03-09 01:11:15 -04:00
Adam Weidman 020fcc9674 refactor: use createExecution for remote lifecycle entries 2026-03-09 01:11:15 -04:00
Adam Weidman 8a90ecaf30 refactor: improve background execution readability 2026-03-09 01:11:15 -04:00
Adam Weidman db32a6185b refactor: align remote background wiring with execution ID naming 2026-03-09 01:11:15 -04:00
Adam Weidman b3850edb8b Support backgrounding remote agent executions via Ctrl+B 2026-03-09 01:11:15 -04:00
Adam Weidman c77fd3fc7a refactor: standardize on pid as canonical background execution field
Remove the dual executionId/pid fields from BackgroundExecutionData.
Use pid everywhere for consistency with existing types (ExecutingToolCall,
ExecutionHandle). A codebase-wide rename to executionId is planned as a
follow-up once all consumers are migrated.
2026-03-09 01:10:42 -04:00
Adam Weidman 013cfafbf9 refactor: rename registerExecution to attachExecution for clarity
The "attach" verb makes it clear that the caller brings an existing
process/handle, while "create" (createExecution) means the lifecycle
service allocates and owns the execution from scratch.
2026-03-09 00:43:41 -04:00
Adam Weidman e17e34d18f simplify background tool info extraction 2026-03-09 00:33:52 -04:00
Adam Weidman 228c978147 simplify background execution helper usage in stream hook 2026-03-09 00:26:58 -04:00
Adam Weidman 9b72826078 Refine execution lifecycle facade follow-ups 2026-03-09 00:14:48 -04:00
Adam Weidman 1e872ead2d refactor: name non-process execution ID seed constant 2026-03-08 23:40:17 -04:00
Adam Weidman fff69df93e refactor: remove virtual execution lifecycle aliases 2026-03-08 23:31:18 -04:00
Adam Weidman 0fc4247287 refactor: make createExecution the primary lifecycle API 2026-03-08 21:40:17 -04:00
Adam Weidman f2e2014894 refactor: generalize CLI background execution handling 2026-03-08 20:41:08 -04:00
Adam Weidman 6860284344 refactor: strengthen background execution abstraction tests 2026-03-08 20:24:11 -04:00
Adam Weidman 9a441ef927 refactor: generalize background execution contract 2026-03-08 19:45:50 -04:00
Adam Weidman 3940d6344a Harden execution lifecycle settling and simplify shell backgrounding API 2026-03-08 18:39:41 -04:00
Adam Weidman e9edd60615 Make execution lifecycle the owner of background state 2026-03-08 18:02:25 -04:00
Adam Weidman 6e291cfab8 Add neutral execution lifecycle facade 2026-03-08 17:32:16 -04:00
16 changed files with 1919 additions and 877 deletions
@@ -80,7 +80,7 @@ export const useShellCommandProcessor = (
setShellInputFocused: (value: boolean) => void,
terminalWidth?: number,
terminalHeight?: number,
activeToolPtyId?: number,
activeBackgroundExecutionId?: number,
isWaitingForConfirmation?: boolean,
) => {
const [state, dispatch] = useReducer(shellReducer, initialState);
@@ -103,7 +103,8 @@ export const useShellCommandProcessor = (
}
const m = manager.current;
const activePtyId = state.activeShellPtyId || activeToolPtyId;
const activePtyId =
state.activeShellPtyId ?? activeBackgroundExecutionId ?? undefined;
useEffect(() => {
const isForegroundActive = !!activePtyId || !!isWaitingForConfirmation;
@@ -191,7 +192,8 @@ export const useShellCommandProcessor = (
]);
const backgroundCurrentShell = useCallback(() => {
const pidToBackground = state.activeShellPtyId || activeToolPtyId;
const pidToBackground =
state.activeShellPtyId ?? activeBackgroundExecutionId;
if (pidToBackground) {
ShellExecutionService.background(pidToBackground);
m.backgroundedPids.add(pidToBackground);
@@ -202,7 +204,7 @@ export const useShellCommandProcessor = (
m.restoreTimeout = null;
}
}
}, [state.activeShellPtyId, activeToolPtyId, m]);
}, [state.activeShellPtyId, activeBackgroundExecutionId, m]);
const dismissBackgroundShell = useCallback(
(pid: number) => {
@@ -96,6 +96,24 @@ const MockedUserPromptEvent = vi.hoisted(() =>
vi.fn().mockImplementation(() => {}),
);
const mockParseAndFormatApiError = vi.hoisted(() => vi.fn());
const mockIsBackgroundExecutionData = vi.hoisted(
() => (data: unknown): data is { pid?: number } => {
if (typeof data !== 'object' || data === null) {
return false;
}
const value = data as {
pid?: unknown;
command?: unknown;
initialOutput?: unknown;
};
return (
(value.pid === undefined || typeof value.pid === 'number') &&
(value.command === undefined || typeof value.command === 'string') &&
(value.initialOutput === undefined ||
typeof value.initialOutput === 'string')
);
},
);
const MockValidationRequiredError = vi.hoisted(
() =>
@@ -121,6 +139,7 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
const actualCoreModule = (await importOriginal()) as any;
return {
...actualCoreModule,
isBackgroundExecutionData: mockIsBackgroundExecutionData,
GitService: vi.fn(),
GeminiClient: MockedGeminiClientClass,
UserPromptEvent: MockedUserPromptEvent,
@@ -599,6 +618,35 @@ describe('useGeminiStream', () => {
expect(mockSendMessageStream).not.toHaveBeenCalled(); // submitQuery uses this
});
it('should expose activePtyId for non-shell executing tools that report an execution ID', () => {
const remoteExecutingTool: TrackedExecutingToolCall = {
request: {
callId: 'remote-call-1',
name: 'remote_agent_call',
args: {},
isClientInitiated: false,
prompt_id: 'prompt-id-remote',
},
status: CoreToolCallStatus.Executing,
responseSubmittedToGemini: false,
tool: {
name: 'remote_agent_call',
displayName: 'Remote Agent',
description: 'Remote agent execution',
build: vi.fn(),
} as any,
invocation: {
getDescription: () => 'Calling remote agent',
} as unknown as AnyToolInvocation,
startTime: Date.now(),
liveOutput: 'working...',
pid: 4242,
};
const { result } = renderTestHook([remoteExecutingTool]);
expect(result.current.activePtyId).toBe(4242);
});
it('should submit tool responses when all tool calls are completed and ready', async () => {
const toolCall1ResponseParts: Part[] = [{ text: 'tool 1 final response' }];
const toolCall2ResponseParts: Part[] = [{ text: 'tool 2 final response' }];
+45 -38
View File
@@ -37,6 +37,7 @@ import {
buildUserSteeringHintPrompt,
GeminiCliOperation,
getPlanModeExitMessage,
isBackgroundExecutionData,
} from '@google/gemini-cli-core';
import type {
Config,
@@ -94,10 +95,10 @@ type ToolResponseWithParts = ToolCallResponseInfo & {
llmContent?: PartListUnion;
};
interface ShellToolData {
pid?: number;
command?: string;
initialOutput?: string;
interface BackgroundedToolInfo {
pid: number;
command: string;
initialOutput: string;
}
enum StreamProcessingStatus {
@@ -111,15 +112,32 @@ const SUPPRESSED_TOOL_ERRORS_NOTE =
const LOW_VERBOSITY_FAILURE_NOTE =
'This request failed. Press F12 for diagnostics, or run /settings and change "Error Verbosity" to full for full details.';
function isShellToolData(data: unknown): data is ShellToolData {
if (typeof data !== 'object' || data === null) {
return false;
function getBackgroundedToolInfo(
toolCall: TrackedCompletedToolCall | TrackedCancelledToolCall,
): BackgroundedToolInfo | undefined {
const response = toolCall.response as ToolResponseWithParts;
const rawData: unknown = response?.data;
if (!isBackgroundExecutionData(rawData)) {
return undefined;
}
const d = data as Partial<ShellToolData>;
if (rawData.pid === undefined) {
return undefined;
}
return {
pid: rawData.pid,
command: rawData.command ?? toolCall.request.name,
initialOutput: rawData.initialOutput ?? '',
};
}
function isBackgroundableExecutingToolCall(
toolCall: TrackedToolCall,
): toolCall is TrackedExecutingToolCall {
return (
(d.pid === undefined || typeof d.pid === 'number') &&
(d.command === undefined || typeof d.command === 'string') &&
(d.initialOutput === undefined || typeof d.initialOutput === 'string')
toolCall.status === CoreToolCallStatus.Executing &&
typeof toolCall.pid === 'number'
);
}
@@ -311,13 +329,11 @@ export const useGeminiStream = (
getPreferredEditor,
);
const activeToolPtyId = useMemo(() => {
const executingShellTool = toolCalls.find(
(tc) =>
tc.status === 'executing' && tc.request.name === 'run_shell_command',
const activeBackgroundExecutionId = useMemo(() => {
const executingBackgroundableTool = toolCalls.find(
isBackgroundableExecutingToolCall,
);
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
return (executingShellTool as TrackedExecutingToolCall | undefined)?.pid;
return executingBackgroundableTool?.pid;
}, [toolCalls]);
const onExec = useCallback(async (done: Promise<void>) => {
@@ -347,7 +363,7 @@ export const useGeminiStream = (
setShellInputFocused,
terminalWidth,
terminalHeight,
activeToolPtyId,
activeBackgroundExecutionId,
);
const streamingState = useMemo(
@@ -525,7 +541,8 @@ export const useGeminiStream = (
onComplete: (result: { userSelection: 'disable' | 'keep' }) => void;
} | null>(null);
const activePtyId = activeShellPtyId || activeToolPtyId;
const activePtyId =
activeShellPtyId ?? activeBackgroundExecutionId ?? undefined;
const prevActiveShellPtyIdRef = useRef<number | null>(null);
useEffect(() => {
@@ -1651,26 +1668,16 @@ export const useGeminiStream = (
!processedMemoryToolsRef.current.has(t.request.callId),
);
// Handle backgrounded shell tools
completedAndReadyToSubmitTools.forEach((t) => {
const isShell = t.request.name === 'run_shell_command';
// Access result from the tracked tool call response
const response = t.response as ToolResponseWithParts;
const rawData = response?.data;
const data = isShellToolData(rawData) ? rawData : undefined;
// Use data.pid for shell commands moved to the background.
const pid = data?.pid;
if (isShell && pid) {
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
const command = (data?.['command'] as string) ?? 'shell';
// eslint-disable-next-line @typescript-eslint/no-unsafe-type-assertion
const initialOutput = (data?.['initialOutput'] as string) ?? '';
registerBackgroundShell(pid, command, initialOutput);
for (const toolCall of completedAndReadyToSubmitTools) {
const backgroundedTool = getBackgroundedToolInfo(toolCall);
if (backgroundedTool) {
registerBackgroundShell(
backgroundedTool.pid,
backgroundedTool.command,
backgroundedTool.initialOutput,
);
}
});
}
if (newSuccessfulMemorySaves.length > 0) {
// Perform the refresh only if there are new ones.
@@ -22,6 +22,7 @@ import type { RemoteAgentDefinition } from './types.js';
import { createMockMessageBus } from '../test-utils/mock-message-bus.js';
import { A2AAuthProviderFactory } from './auth-provider/factory.js';
import type { A2AAuthProvider } from './auth-provider/types.js';
import { ExecutionLifecycleService } from '../services/executionLifecycleService.js';
// Mock A2AClientManager
vi.mock('./a2a-client-manager.js', () => ({
@@ -58,6 +59,7 @@ describe('RemoteAgentInvocation', () => {
beforeEach(() => {
vi.clearAllMocks();
ExecutionLifecycleService.resetForTest();
(A2AClientManager.getInstance as Mock).mockReturnValue(mockClientManager);
(
RemoteAgentInvocation as unknown as {
@@ -583,6 +585,94 @@ describe('RemoteAgentInvocation', () => {
'Generating...\n\nArtifact (Result):\nPart 1 Part 2',
);
});
it('should support Ctrl+B backgrounding through execution lifecycle IDs', async () => {
mockClientManager.getClient.mockReturnValue({});
let releaseSecondChunk: (() => void) | undefined;
const secondChunkGate = new Promise<void>((resolve) => {
releaseSecondChunk = resolve;
});
mockClientManager.sendMessageStream.mockImplementation(
async function* () {
yield {
kind: 'message',
messageId: 'msg-1',
role: 'agent',
parts: [{ kind: 'text', text: 'Chunk 1' }],
};
await secondChunkGate;
yield {
kind: 'message',
messageId: 'msg-2',
role: 'agent',
parts: [{ kind: 'text', text: 'Chunk 2' }],
};
},
);
let executionId: number | undefined;
const onExit = vi.fn();
let unsubscribeOnExit: (() => void) | undefined;
const streamedOutputChunks: string[] = [];
let unsubscribeStream: (() => void) | undefined;
const updateOutput = vi.fn((output: unknown) => {
if (output === 'Chunk 1' && executionId !== undefined) {
ExecutionLifecycleService.background(executionId);
unsubscribeStream = ExecutionLifecycleService.subscribe(
executionId,
(event) => {
if (event.type === 'data' && typeof event.chunk === 'string') {
streamedOutputChunks.push(event.chunk);
}
},
);
}
});
const invocation = new RemoteAgentInvocation(
mockDefinition,
{ query: 'long task' },
mockMessageBus,
);
const resultPromise = invocation.execute(
new AbortController().signal,
updateOutput,
undefined,
(newExecutionId) => {
executionId = newExecutionId;
unsubscribeOnExit = ExecutionLifecycleService.onExit(
newExecutionId,
onExit,
);
},
);
const result = await resultPromise;
expect(executionId).toBeDefined();
expect(result.returnDisplay).toContain(
'Remote agent moved to background',
);
expect(result.data).toMatchObject({
pid: executionId,
initialOutput: 'Chunk 1',
});
releaseSecondChunk?.();
await vi.waitFor(() => {
expect(onExit).toHaveBeenCalledWith(0, undefined);
});
await vi.waitFor(() => {
expect(streamedOutputChunks.join('')).toContain('Chunk 2');
});
unsubscribeStream?.();
unsubscribeOnExit?.();
});
});
describe('Confirmations', () => {
+208 -84
View File
@@ -6,6 +6,7 @@
import {
BaseToolInvocation,
type BackgroundExecutionData,
type ToolConfirmationOutcome,
type ToolResult,
type ToolCallConfirmationDetails,
@@ -27,6 +28,15 @@ import type { AuthenticationHandler } from '@a2a-js/sdk/client';
import { debugLogger } from '../utils/debugLogger.js';
import type { AnsiOutput } from '../utils/terminalSerializer.js';
import { A2AAuthProviderFactory } from './auth-provider/factory.js';
import {
ExecutionLifecycleService,
type ExecutionResult,
} from '../services/executionLifecycleService.js';
const OUTPUT_RESYNC_MARKER = '\n\n[Output updated]\n';
const REMOTE_AGENT_ERROR_PREFIX = 'Error calling remote agent: ';
const MISSING_EXECUTION_ID_MESSAGE =
'Error calling remote agent: missing execution ID.';
/**
* Authentication handler implementation using Google Application Default Credentials (ADC).
@@ -145,102 +155,216 @@ export class RemoteAgentInvocation extends BaseToolInvocation<
};
}
async execute(
_signal: AbortSignal,
updateOutput?: (output: string | AnsiOutput) => void,
): Promise<ToolResult> {
// 1. Ensure the agent is loaded (cached by manager)
// We assume the user has provided an access token via some mechanism (TODO),
// or we rely on ADC.
const reassembler = new A2AResultReassembler();
try {
const priorState = RemoteAgentInvocation.sessionState.get(
this.definition.name,
);
if (priorState) {
this.contextId = priorState.contextId;
this.taskId = priorState.taskId;
}
private restoreSessionState(): void {
const priorState = RemoteAgentInvocation.sessionState.get(
this.definition.name,
);
if (!priorState) {
return;
}
this.contextId = priorState.contextId;
this.taskId = priorState.taskId;
}
const authHandler = await this.getAuthHandler();
private persistSessionState(): void {
// Persist state even on partial failures or aborts to maintain conversational continuity.
RemoteAgentInvocation.sessionState.set(this.definition.name, {
contextId: this.contextId,
taskId: this.taskId,
});
}
if (!this.clientManager.getClient(this.definition.name)) {
await this.clientManager.loadAgent(
this.definition.name,
this.definition.agentCardUrl,
authHandler,
);
}
private updateSessionStateFromResponseChunk(chunk: SendMessageResult): void {
const { contextId: newContextId, taskId: newTaskId, clearTaskId } =
extractIdsFromResponse(chunk);
const message = this.params.query;
if (newContextId) {
this.contextId = newContextId;
}
const stream = this.clientManager.sendMessageStream(
this.definition.name,
message,
{
contextId: this.contextId,
taskId: this.taskId,
signal: _signal,
},
);
this.taskId = clearTaskId ? undefined : (newTaskId ?? this.taskId);
}
let finalResponse: SendMessageResult | undefined;
private buildBackgroundResult(
executionId: number,
output: string,
): ToolResult {
const command = `${this.getDescription()}: ${this.params.query}`;
const backgroundMessage = `Remote agent moved to background (Execution ID: ${executionId}). Output hidden. Press Ctrl+B to view.`;
const data: BackgroundExecutionData = {
pid: executionId,
command,
initialOutput: output,
};
for await (const chunk of stream) {
if (_signal.aborted) {
throw new Error('Operation aborted');
}
finalResponse = chunk;
reassembler.update(chunk);
return {
llmContent: [{ text: backgroundMessage }],
returnDisplay: backgroundMessage,
data,
};
}
if (updateOutput) {
updateOutput(reassembler.toString());
}
private buildCompletionResult(
executionResult: ExecutionResult,
executionId: number,
): ToolResult {
if (executionResult.backgrounded) {
return this.buildBackgroundResult(executionId, executionResult.output);
}
const {
contextId: newContextId,
taskId: newTaskId,
clearTaskId,
} = extractIdsFromResponse(chunk);
if (newContextId) {
this.contextId = newContextId;
}
this.taskId = clearTaskId ? undefined : (newTaskId ?? this.taskId);
}
if (!finalResponse) {
throw new Error('No response from remote agent.');
}
const finalOutput = reassembler.toString();
debugLogger.debug(
`[RemoteAgent] Final response from ${this.definition.name}:\n${JSON.stringify(finalResponse, null, 2)}`,
);
return {
llmContent: [{ text: finalOutput }],
returnDisplay: finalOutput,
};
} catch (error: unknown) {
const partialOutput = reassembler.toString();
const errorMessage = `Error calling remote agent: ${error instanceof Error ? error.message : String(error)}`;
const fullDisplay = partialOutput
? `${partialOutput}\n\n${errorMessage}`
: errorMessage;
if (executionResult.error) {
const fullDisplay = executionResult.output
? `${executionResult.output}\n\n${executionResult.error.message}`
: executionResult.error.message;
return {
llmContent: [{ text: fullDisplay }],
returnDisplay: fullDisplay,
error: { message: errorMessage },
error: { message: executionResult.error.message },
};
} finally {
// Persist state even on partial failures or aborts to maintain conversational continuity.
RemoteAgentInvocation.sessionState.set(this.definition.name, {
contextId: this.contextId,
taskId: this.taskId,
});
}
return {
llmContent: [{ text: executionResult.output }],
returnDisplay: executionResult.output,
};
}
private formatRemoteAgentError(error: unknown): string {
return `${REMOTE_AGENT_ERROR_PREFIX}${
error instanceof Error ? error.message : String(error)
}`;
}
private publishBackgroundDelta(
executionId: number,
previousOutput: string,
nextOutput: string,
): string {
if (nextOutput.length === 0 || nextOutput === previousOutput) {
return previousOutput;
}
if (nextOutput.startsWith(previousOutput)) {
ExecutionLifecycleService.appendOutput(
executionId,
nextOutput.slice(previousOutput.length),
);
return nextOutput;
}
// If the reassembled output changes non-monotonically, resync by appending
// the full latest snapshot with a clear separator.
ExecutionLifecycleService.appendOutput(
executionId,
`${OUTPUT_RESYNC_MARKER}${nextOutput}`,
);
return nextOutput;
}
async execute(
signal: AbortSignal,
updateOutput?: (output: string | AnsiOutput) => void,
_shellExecutionConfig?: unknown,
setExecutionIdCallback?: (executionId: number) => void,
): Promise<ToolResult> {
const reassembler = new A2AResultReassembler();
const executionController = new AbortController();
const onAbort = () => executionController.abort();
signal.addEventListener('abort', onAbort, { once: true });
const { pid: executionId, result } = ExecutionLifecycleService.createExecution(
'',
() => executionController.abort(),
'remote_agent',
);
if (executionId === undefined) {
signal.removeEventListener('abort', onAbort);
return {
llmContent: [{ text: MISSING_EXECUTION_ID_MESSAGE }],
returnDisplay: MISSING_EXECUTION_ID_MESSAGE,
error: {
message: MISSING_EXECUTION_ID_MESSAGE,
},
};
}
setExecutionIdCallback?.(executionId);
const streamRemoteAgentExecution = async () => {
let lastOutput = '';
try {
this.restoreSessionState();
const authHandler = await this.getAuthHandler();
if (!this.clientManager.getClient(this.definition.name)) {
await this.clientManager.loadAgent(
this.definition.name,
this.definition.agentCardUrl,
authHandler,
);
}
const stream = this.clientManager.sendMessageStream(
this.definition.name,
this.params.query,
{
contextId: this.contextId,
taskId: this.taskId,
signal: executionController.signal,
},
);
let finalResponse: SendMessageResult | undefined;
for await (const chunk of stream) {
if (executionController.signal.aborted) {
throw new Error('Operation aborted');
}
finalResponse = chunk;
reassembler.update(chunk);
const currentOutput = reassembler.toString();
lastOutput = this.publishBackgroundDelta(
executionId,
lastOutput,
currentOutput,
);
updateOutput?.(currentOutput);
this.updateSessionStateFromResponseChunk(chunk);
}
if (!finalResponse) {
throw new Error('No response from remote agent.');
}
debugLogger.debug(
`[RemoteAgent] Final response from ${this.definition.name}:\n${JSON.stringify(finalResponse, null, 2)}`,
);
ExecutionLifecycleService.completeExecution(executionId, {
exitCode: 0,
});
} catch (error: unknown) {
const partialOutput = reassembler.toString();
lastOutput = this.publishBackgroundDelta(
executionId,
lastOutput,
partialOutput,
);
ExecutionLifecycleService.completeExecution(executionId, {
error: new Error(this.formatRemoteAgentError(error)),
aborted: executionController.signal.aborted,
exitCode: executionController.signal.aborted ? 130 : 1,
});
} finally {
signal.removeEventListener('abort', onAbort);
this.persistSessionState();
}
};
void streamRemoteAgentExecution();
const executionResult = await result;
return this.buildCompletionResult(executionResult, executionId);
}
}
@@ -11,6 +11,7 @@ import {
BaseToolInvocation,
type ToolResult,
type AnyDeclarativeTool,
type ToolLiveOutput,
} from '../tools/tools.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import type { HookSystem } from '../hooks/hookSystem.js';
@@ -37,6 +38,30 @@ class MockInvocation extends BaseToolInvocation<{ key?: string }, ToolResult> {
}
}
class MockBackgroundableInvocation extends BaseToolInvocation<
{ key?: string },
ToolResult
> {
constructor(params: { key?: string }, messageBus: MessageBus) {
super(params, messageBus);
}
getDescription() {
return 'mock-pid';
}
async execute(
_signal: AbortSignal,
_updateOutput?: (output: ToolLiveOutput) => void,
_shellExecutionConfig?: unknown,
setExecutionIdCallback?: (executionId: number) => void,
) {
setExecutionIdCallback?.(4242);
return {
llmContent: 'pid',
returnDisplay: 'pid',
};
}
}
describe('executeToolWithHooks', () => {
let messageBus: MessageBus;
let mockTool: AnyDeclarativeTool;
@@ -258,4 +283,26 @@ describe('executeToolWithHooks', () => {
expect(invocation.params.key).toBe('original');
expect(mockTool.build).not.toHaveBeenCalled();
});
it('should pass execution ID callback through for non-shell invocations', async () => {
const invocation = new MockBackgroundableInvocation({}, messageBus);
const abortSignal = new AbortController().signal;
const setExecutionIdCallback = vi.fn();
vi.mocked(mockHookSystem.fireBeforeToolEvent).mockResolvedValue(undefined);
vi.mocked(mockHookSystem.fireAfterToolEvent).mockResolvedValue(undefined);
await executeToolWithHooks(
invocation,
'test_tool',
abortSignal,
mockTool,
undefined,
undefined,
setExecutionIdCallback,
mockConfig,
);
expect(setExecutionIdCallback).toHaveBeenCalledWith(4242);
});
});
+12 -21
View File
@@ -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<ToolResult> {
@@ -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) {
@@ -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,
@@ -534,6 +534,59 @@ describe('ToolExecutor', () => {
);
});
it('should report execution ID updates for non-shell backgroundable tools', async () => {
const mockTool = new MockTool({
name: 'remote_agent_call',
description: 'Remote agent call',
});
const invocation = mockTool.build({});
const testExecutionId = 67890;
vi.mocked(coreToolHookTriggers.executeToolWithHooks).mockImplementation(
async (
_inv,
_name,
_sig,
_tool,
_liveCb,
_shellCfg,
setExecutionIdCallback,
) => {
setExecutionIdCallback?.(testExecutionId);
return { llmContent: 'done', returnDisplay: 'done' };
},
);
const scheduledCall: ScheduledToolCall = {
status: CoreToolCallStatus.Scheduled,
request: {
callId: 'call-remote-pid',
name: 'remote_agent_call',
args: {},
isClientInitiated: false,
prompt_id: 'prompt-remote-pid',
},
tool: mockTool,
invocation: invocation as unknown as AnyToolInvocation,
startTime: Date.now(),
};
const onUpdateToolCall = vi.fn();
await executor.execute({
call: scheduledCall,
signal: new AbortController().signal,
onUpdateToolCall,
});
expect(onUpdateToolCall).toHaveBeenCalledWith(
expect.objectContaining({
status: CoreToolCallStatus.Executing,
pid: testExecutionId,
}),
);
});
it('should return cancelled result with partial output when signal is aborted', async () => {
const mockTool = new MockTool({
name: 'slowTool',
+22 -37
View File
@@ -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<ToolResult>;
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;
@@ -0,0 +1,289 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { beforeEach, describe, expect, it, vi } from 'vitest';
import {
ExecutionLifecycleService,
type ExecutionHandle,
type ExecutionResult,
} from './executionLifecycleService.js';
function createResult(
overrides: Partial<ExecutionResult> = {},
): ExecutionResult {
return {
rawOutput: Buffer.from(''),
output: '',
exitCode: 0,
signal: null,
error: null,
aborted: false,
pid: 123,
executionMethod: 'child_process',
...overrides,
};
}
describe('ExecutionLifecycleService', () => {
beforeEach(() => {
ExecutionLifecycleService.resetForTest();
});
it('completes managed executions in the foreground and notifies exit subscribers', async () => {
const handle = ExecutionLifecycleService.createExecution();
if (handle.pid === undefined) {
throw new Error('Expected execution ID.');
}
const onExit = vi.fn();
const unsubscribe = ExecutionLifecycleService.onExit(handle.pid, onExit);
ExecutionLifecycleService.appendOutput(handle.pid, 'Hello');
ExecutionLifecycleService.appendOutput(handle.pid, ' World');
ExecutionLifecycleService.completeExecution(handle.pid, {
exitCode: 0,
});
const result = await handle.result;
expect(result.output).toBe('Hello World');
expect(result.executionMethod).toBe('none');
expect(result.backgrounded).toBeUndefined();
await vi.waitFor(() => {
expect(onExit).toHaveBeenCalledWith(0, undefined);
});
unsubscribe();
});
it('supports explicit execution methods for managed executions', async () => {
const handle = ExecutionLifecycleService.createExecution(
'',
undefined,
'remote_agent',
);
if (handle.pid === undefined) {
throw new Error('Expected execution ID.');
}
ExecutionLifecycleService.completeExecution(handle.pid, {
exitCode: 0,
});
const result = await handle.result;
expect(result.executionMethod).toBe('remote_agent');
});
it('supports backgrounding managed executions and continues streaming updates', async () => {
const handle = ExecutionLifecycleService.createExecution();
if (handle.pid === undefined) {
throw new Error('Expected execution ID.');
}
const chunks: string[] = [];
const onExit = vi.fn();
const unsubscribeStream = ExecutionLifecycleService.subscribe(
handle.pid,
(event) => {
if (event.type === 'data' && typeof event.chunk === 'string') {
chunks.push(event.chunk);
}
},
);
const unsubscribeExit = ExecutionLifecycleService.onExit(handle.pid, onExit);
ExecutionLifecycleService.appendOutput(handle.pid, 'Chunk 1');
ExecutionLifecycleService.background(handle.pid);
const backgroundResult = await handle.result;
expect(backgroundResult.backgrounded).toBe(true);
expect(backgroundResult.output).toBe('Chunk 1');
ExecutionLifecycleService.appendOutput(handle.pid, '\nChunk 2');
ExecutionLifecycleService.completeExecution(handle.pid, {
exitCode: 0,
});
await vi.waitFor(() => {
expect(chunks.join('')).toContain('Chunk 2');
expect(onExit).toHaveBeenCalledWith(0, undefined);
});
unsubscribeStream();
unsubscribeExit();
});
it('kills managed executions and resolves with aborted result', async () => {
const onKill = vi.fn();
const handle = ExecutionLifecycleService.createExecution('', onKill);
if (handle.pid === undefined) {
throw new Error('Expected execution ID.');
}
ExecutionLifecycleService.appendOutput(handle.pid, 'work');
ExecutionLifecycleService.kill(handle.pid);
const result = await handle.result;
expect(onKill).toHaveBeenCalledTimes(1);
expect(result.aborted).toBe(true);
expect(result.exitCode).toBe(130);
expect(result.error?.message).toContain('Operation cancelled by user');
});
it('does not probe OS process state for completed non-process execution IDs', async () => {
const handle = ExecutionLifecycleService.createExecution();
if (handle.pid === undefined) {
throw new Error('Expected execution ID.');
}
ExecutionLifecycleService.completeExecution(handle.pid, { exitCode: 0 });
await handle.result;
const processKillSpy = vi.spyOn(process, 'kill');
expect(ExecutionLifecycleService.isActive(handle.pid)).toBe(false);
expect(processKillSpy).not.toHaveBeenCalled();
processKillSpy.mockRestore();
});
it('manages external executions through registration hooks', async () => {
const writeInput = vi.fn();
const isActive = vi.fn().mockReturnValue(true);
const exitListener = vi.fn();
const chunks: string[] = [];
let output = 'seed';
const handle: ExecutionHandle = ExecutionLifecycleService.attachExecution(
4321,
{
executionMethod: 'child_process',
getBackgroundOutput: () => output,
getSubscriptionSnapshot: () => output,
writeInput,
isActive,
},
);
const unsubscribe = ExecutionLifecycleService.subscribe(4321, (event) => {
if (event.type === 'data' && typeof event.chunk === 'string') {
chunks.push(event.chunk);
}
});
ExecutionLifecycleService.onExit(4321, exitListener);
ExecutionLifecycleService.writeInput(4321, 'stdin');
expect(writeInput).toHaveBeenCalledWith('stdin');
expect(ExecutionLifecycleService.isActive(4321)).toBe(true);
const firstChunk = { type: 'data', chunk: ' +delta' } as const;
ExecutionLifecycleService.emitEvent(4321, firstChunk);
output += firstChunk.chunk;
ExecutionLifecycleService.background(4321);
const backgroundResult = await handle.result;
expect(backgroundResult.backgrounded).toBe(true);
expect(backgroundResult.output).toBe('seed +delta');
expect(backgroundResult.executionMethod).toBe('child_process');
ExecutionLifecycleService.completeWithResult(
4321,
createResult({
pid: 4321,
output: 'seed +delta done',
rawOutput: Buffer.from('seed +delta done'),
executionMethod: 'child_process',
}),
);
await vi.waitFor(() => {
expect(exitListener).toHaveBeenCalledWith(0, undefined);
});
const lateExit = vi.fn();
ExecutionLifecycleService.onExit(4321, lateExit);
expect(lateExit).toHaveBeenCalledWith(0, undefined);
unsubscribe();
});
it('supports late subscription catch-up after backgrounding an external execution', async () => {
let output = 'seed';
const onExit = vi.fn();
const handle = ExecutionLifecycleService.attachExecution(4322, {
executionMethod: 'child_process',
getBackgroundOutput: () => output,
getSubscriptionSnapshot: () => output,
});
ExecutionLifecycleService.onExit(4322, onExit);
ExecutionLifecycleService.background(4322);
const backgroundResult = await handle.result;
expect(backgroundResult.backgrounded).toBe(true);
expect(backgroundResult.output).toBe('seed');
output += ' +late';
ExecutionLifecycleService.emitEvent(4322, { type: 'data', chunk: ' +late' });
const chunks: string[] = [];
const unsubscribe = ExecutionLifecycleService.subscribe(4322, (event) => {
if (event.type === 'data' && typeof event.chunk === 'string') {
chunks.push(event.chunk);
}
});
expect(chunks[0]).toBe('seed +late');
output += ' +live';
ExecutionLifecycleService.emitEvent(4322, { type: 'data', chunk: ' +live' });
expect(chunks[chunks.length - 1]).toBe(' +live');
ExecutionLifecycleService.completeWithResult(
4322,
createResult({
pid: 4322,
output,
rawOutput: Buffer.from(output),
executionMethod: 'child_process',
}),
);
await vi.waitFor(() => {
expect(onExit).toHaveBeenCalledWith(0, undefined);
});
unsubscribe();
});
it('kills external executions and settles pending promises', async () => {
const terminate = vi.fn();
const onExit = vi.fn();
const handle = ExecutionLifecycleService.attachExecution(4323, {
executionMethod: 'child_process',
initialOutput: 'running',
kill: terminate,
});
ExecutionLifecycleService.onExit(4323, onExit);
ExecutionLifecycleService.kill(4323);
const result = await handle.result;
expect(terminate).toHaveBeenCalledTimes(1);
expect(result.aborted).toBe(true);
expect(result.exitCode).toBe(130);
expect(result.output).toBe('running');
expect(result.error?.message).toContain('Operation cancelled by user');
expect(onExit).toHaveBeenCalledWith(130, undefined);
});
it('rejects duplicate execution registration for active execution IDs', () => {
ExecutionLifecycleService.attachExecution(4324, {
executionMethod: 'child_process',
});
expect(() => {
ExecutionLifecycleService.attachExecution(4324, {
executionMethod: 'child_process',
});
}).toThrow('Execution 4324 is already attached.');
});
});
@@ -0,0 +1,451 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import type { AnsiOutput } from '../utils/terminalSerializer.js';
export type ExecutionMethod =
| 'lydell-node-pty'
| 'node-pty'
| 'child_process'
| 'remote_agent'
| 'none';
export interface ExecutionResult {
rawOutput: Buffer;
output: string;
exitCode: number | null;
signal: number | null;
error: Error | null;
aborted: boolean;
pid: number | undefined;
executionMethod: ExecutionMethod;
backgrounded?: boolean;
}
export interface ExecutionHandle {
pid: number | undefined;
result: Promise<ExecutionResult>;
}
export type ExecutionOutputEvent =
| {
type: 'data';
chunk: string | AnsiOutput;
}
| {
type: 'binary_detected';
}
| {
type: 'binary_progress';
bytesReceived: number;
}
| {
type: 'exit';
exitCode: number | null;
signal: number | null;
};
export interface ExecutionCompletionOptions {
exitCode?: number | null;
signal?: number | null;
error?: Error | null;
aborted?: boolean;
}
export interface ExternalExecutionRegistration {
executionMethod: ExecutionMethod;
initialOutput?: string;
getBackgroundOutput?: () => string;
getSubscriptionSnapshot?: () => string | AnsiOutput | undefined;
writeInput?: (input: string) => void;
kill?: () => void;
isActive?: () => boolean;
}
interface ManagedExecutionBase {
executionMethod: ExecutionMethod;
output: string;
getBackgroundOutput?: () => string;
getSubscriptionSnapshot?: () => string | AnsiOutput | undefined;
}
interface VirtualExecutionState extends ManagedExecutionBase {
kind: 'virtual';
onKill?: () => void;
}
interface ExternalExecutionState extends ManagedExecutionBase {
kind: 'external';
writeInput?: (input: string) => void;
kill?: () => void;
isActive?: () => boolean;
}
type ManagedExecutionState = VirtualExecutionState | ExternalExecutionState;
const NON_PROCESS_EXECUTION_ID_START = 2_000_000_000;
/**
* Central owner for execution backgrounding lifecycle across shell and tools.
*/
export class ExecutionLifecycleService {
private static readonly EXIT_INFO_TTL_MS = 5 * 60 * 1000;
private static nextExecutionId = NON_PROCESS_EXECUTION_ID_START;
private static activeExecutions = new Map<number, ManagedExecutionState>();
private static activeResolvers = new Map<
number,
(result: ExecutionResult) => void
>();
private static activeListeners = new Map<
number,
Set<(event: ExecutionOutputEvent) => void>
>();
private static exitedExecutionInfo = new Map<
number,
{ exitCode: number; signal?: number }
>();
private static storeExitInfo(
executionId: number,
exitCode: number,
signal?: number,
): void {
this.exitedExecutionInfo.set(executionId, {
exitCode,
signal,
});
setTimeout(() => {
this.exitedExecutionInfo.delete(executionId);
}, this.EXIT_INFO_TTL_MS).unref();
}
private static allocateExecutionId(): number {
let executionId = ++this.nextExecutionId;
while (this.activeExecutions.has(executionId)) {
executionId = ++this.nextExecutionId;
}
return executionId;
}
private static createPendingResult(
executionId: number,
): Promise<ExecutionResult> {
return new Promise<ExecutionResult>((resolve) => {
this.activeResolvers.set(executionId, resolve);
});
}
private static createAbortedResult(
executionId: number,
execution: ManagedExecutionState,
): ExecutionResult {
const output = execution.getBackgroundOutput?.() ?? execution.output;
return {
rawOutput: Buffer.from(output, 'utf8'),
output,
exitCode: 130,
signal: null,
error: new Error('Operation cancelled by user.'),
aborted: true,
pid: executionId,
executionMethod: execution.executionMethod,
};
}
/**
* Resets lifecycle state for isolated unit tests.
*/
static resetForTest(): void {
this.activeExecutions.clear();
this.activeResolvers.clear();
this.activeListeners.clear();
this.exitedExecutionInfo.clear();
this.nextExecutionId = NON_PROCESS_EXECUTION_ID_START;
}
static attachExecution(
executionId: number,
registration: ExternalExecutionRegistration,
): ExecutionHandle {
if (this.activeExecutions.has(executionId) || this.activeResolvers.has(executionId)) {
throw new Error(`Execution ${executionId} is already attached.`);
}
this.exitedExecutionInfo.delete(executionId);
this.activeExecutions.set(executionId, {
executionMethod: registration.executionMethod,
output: registration.initialOutput ?? '',
kind: 'external',
getBackgroundOutput: registration.getBackgroundOutput,
getSubscriptionSnapshot: registration.getSubscriptionSnapshot,
writeInput: registration.writeInput,
kill: registration.kill,
isActive: registration.isActive,
});
return {
pid: executionId,
result: this.createPendingResult(executionId),
};
}
static createExecution(
initialOutput = '',
onKill?: () => void,
executionMethod: ExecutionMethod = 'none',
): ExecutionHandle {
const executionId = this.allocateExecutionId();
this.activeExecutions.set(executionId, {
executionMethod,
output: initialOutput,
kind: 'virtual',
onKill,
getBackgroundOutput: () => {
const state = this.activeExecutions.get(executionId);
return state?.output ?? initialOutput;
},
getSubscriptionSnapshot: () => {
const state = this.activeExecutions.get(executionId);
return state?.output ?? initialOutput;
},
});
return {
pid: executionId,
result: this.createPendingResult(executionId),
};
}
static appendOutput(executionId: number, chunk: string): void {
const execution = this.activeExecutions.get(executionId);
if (!execution || chunk.length === 0) {
return;
}
execution.output += chunk;
this.emitEvent(executionId, { type: 'data', chunk });
}
static emitEvent(executionId: number, event: ExecutionOutputEvent): void {
const listeners = this.activeListeners.get(executionId);
if (listeners) {
listeners.forEach((listener) => listener(event));
}
}
private static resolvePending(
executionId: number,
result: ExecutionResult,
): void {
const resolve = this.activeResolvers.get(executionId);
if (!resolve) {
return;
}
resolve(result);
this.activeResolvers.delete(executionId);
}
private static settleExecution(
executionId: number,
result: ExecutionResult,
): void {
if (!this.activeExecutions.has(executionId)) {
return;
}
this.resolvePending(executionId, result);
this.emitEvent(executionId, {
type: 'exit',
exitCode: result.exitCode,
signal: result.signal,
});
this.activeListeners.delete(executionId);
this.activeExecutions.delete(executionId);
this.storeExitInfo(
executionId,
result.exitCode ?? 0,
result.signal ?? undefined,
);
}
static completeExecution(
executionId: number,
options?: ExecutionCompletionOptions,
): void {
const execution = this.activeExecutions.get(executionId);
if (!execution) {
return;
}
const {
error = null,
aborted = false,
exitCode = error ? 1 : 0,
signal = null,
} = options ?? {};
const output = execution.getBackgroundOutput?.() ?? execution.output;
this.settleExecution(executionId, {
rawOutput: Buffer.from(output, 'utf8'),
output,
exitCode,
signal,
error,
aborted,
pid: executionId,
executionMethod: execution.executionMethod,
});
}
static completeWithResult(
executionId: number,
result: ExecutionResult,
): void {
this.settleExecution(executionId, result);
}
static background(executionId: number): void {
const resolve = this.activeResolvers.get(executionId);
if (!resolve) {
return;
}
const execution = this.activeExecutions.get(executionId);
if (!execution) {
return;
}
const output = execution.getBackgroundOutput?.() ?? execution.output;
resolve({
rawOutput: Buffer.from(''),
output,
exitCode: null,
signal: null,
error: null,
aborted: false,
pid: executionId,
executionMethod: execution.executionMethod,
backgrounded: true,
});
this.activeResolvers.delete(executionId);
}
static subscribe(
executionId: number,
listener: (event: ExecutionOutputEvent) => void,
): () => void {
if (!this.activeListeners.has(executionId)) {
this.activeListeners.set(executionId, new Set());
}
this.activeListeners.get(executionId)?.add(listener);
const execution = this.activeExecutions.get(executionId);
if (execution) {
const snapshot =
execution.getSubscriptionSnapshot?.() ??
(execution.output.length > 0 ? execution.output : undefined);
if (snapshot && (typeof snapshot !== 'string' || snapshot.length > 0)) {
listener({ type: 'data', chunk: snapshot });
}
}
return () => {
this.activeListeners.get(executionId)?.delete(listener);
if (this.activeListeners.get(executionId)?.size === 0) {
this.activeListeners.delete(executionId);
}
};
}
static onExit(
executionId: number,
callback: (exitCode: number, signal?: number) => void,
): () => void {
if (this.activeExecutions.has(executionId)) {
const listener = (event: ExecutionOutputEvent) => {
if (event.type === 'exit') {
callback(event.exitCode ?? 0, event.signal ?? undefined);
unsubscribe();
}
};
const unsubscribe = this.subscribe(executionId, listener);
return unsubscribe;
}
const exitedInfo = this.exitedExecutionInfo.get(executionId);
if (exitedInfo) {
callback(exitedInfo.exitCode, exitedInfo.signal);
}
return () => {};
}
static kill(executionId: number): void {
const execution = this.activeExecutions.get(executionId);
if (!execution) {
return;
}
if (execution.kind === 'virtual') {
execution.onKill?.();
}
if (execution.kind === 'external') {
execution.kill?.();
}
this.completeWithResult(
executionId,
this.createAbortedResult(executionId, execution),
);
}
static isActive(executionId: number): boolean {
const execution = this.activeExecutions.get(executionId);
if (!execution) {
if (executionId >= NON_PROCESS_EXECUTION_ID_START) {
return false;
}
try {
return process.kill(executionId, 0);
} catch {
return false;
}
}
if (execution.kind === 'virtual') {
return true;
}
if (execution.kind === 'external' && execution.isActive) {
try {
return execution.isActive();
} catch {
return false;
}
}
try {
return process.kill(executionId, 0);
} catch {
return false;
}
}
static writeInput(executionId: number, input: string): void {
const execution = this.activeExecutions.get(executionId);
if (execution?.kind === 'external') {
execution.writeInput?.(input);
}
}
}
@@ -21,6 +21,7 @@ import {
type ShellOutputEvent,
type ShellExecutionConfig,
} from './shellExecutionService.js';
import { ExecutionLifecycleService } from './executionLifecycleService.js';
import type { AnsiOutput, AnsiToken } from '../utils/terminalSerializer.js';
// Hoisted Mocks
@@ -166,6 +167,7 @@ describe('ShellExecutionService', () => {
beforeEach(() => {
vi.clearAllMocks();
ExecutionLifecycleService.resetForTest();
mockSerializeTerminalToObject.mockReturnValue([]);
mockIsBinary.mockReturnValue(false);
mockPlatform.mockReturnValue('linux');
@@ -432,13 +434,21 @@ describe('ShellExecutionService', () => {
});
describe('pty interaction', () => {
let activePtysGetSpy: { mockRestore: () => void };
beforeEach(() => {
vi.spyOn(ShellExecutionService['activePtys'], 'get').mockReturnValue({
// eslint-disable-next-line @typescript-eslint/no-explicit-any
ptyProcess: mockPtyProcess as any,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
headlessTerminal: mockHeadlessTerminal as any,
});
activePtysGetSpy = vi
.spyOn(ShellExecutionService['activePtys'], 'get')
.mockReturnValue({
// eslint-disable-next-line @typescript-eslint/no-explicit-any
ptyProcess: mockPtyProcess as any,
// eslint-disable-next-line @typescript-eslint/no-explicit-any
headlessTerminal: mockHeadlessTerminal as any,
});
});
afterEach(() => {
activePtysGetSpy.mockRestore();
});
it('should write to the pty and trigger a render', async () => {
File diff suppressed because it is too large Load Diff
+5 -4
View File
@@ -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<ToolResult> {
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<string, unknown> | undefined;
let data: BackgroundExecutionData | undefined;
let llmContent = '';
let timeoutMessage = '';
+36
View File
@@ -61,15 +61,51 @@ 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<TResult>;
}
/**
* Structured payload used by tools to surface background execution metadata to
* the CLI UI.
*
* NOTE: `pid` is used as the canonical identifier for now to stay consistent
* with existing types (ExecutingToolCall.pid, ExecutionHandle.pid, etc.).
* A future rename to `executionId` is planned once the codebase is fully
* migrated not done in this PR to keep the diff focused on the abstraction.
*/
export interface BackgroundExecutionData extends Record<string, unknown> {
pid?: number;
command?: string;
initialOutput?: string;
}
export function isBackgroundExecutionData(
data: unknown,
): data is BackgroundExecutionData {
if (typeof data !== 'object' || data === null) {
return false;
}
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 (
(pid === undefined || typeof pid === 'number') &&
(command === undefined || typeof command === 'string') &&
(initialOutput === undefined || typeof initialOutput === 'string')
);
}
/**
* Options for policy updates that can be customized by tool invocations.
*/
+87
View File
@@ -0,0 +1,87 @@
## Summary
Extract a generic `ExecutionLifecycleService` that owns background-execution state (create, stream, background, complete, kill) so that **any tool** — not just shell — can be backgrounded via Ctrl+B. This is a pure abstraction/renaming refactor with no new features; the follow-up PR (stacked on this branch) wires remote agents into the same lifecycle.
## Details
### Why this change
Background execution was deeply coupled to `ShellExecutionService`: resolver maps, listener maps, exit-info TTL, and the `background()` / `subscribe()` / `onExit()` API all lived inside the shell service. Adding Ctrl+B support for remote agents (or MCP tools, local agents, etc.) would have required either duplicating that machinery or reaching into shell internals. This PR pulls the lifecycle out into a standalone service that shell delegates to.
### Recommended reading order
This is a large diff (13 files, ~1500 insertions) but the bulk is mechanical delegation. Read in this order:
1. **`executionLifecycleService.ts`** — the new abstraction. Two execution kinds:
- *virtual*: tool calls `createExecution()`, gets an ID + result promise, streams via `appendOutput()`, completes via `completeExecution()`. Used by the upcoming remote-agent wiring.
- *external*: shell calls `attachExecution(pid, {...hooks})` to attach lifecycle tracking to a real OS process. Same lifecycle, but with write/kill/isActive hooks delegated back to the PTY/child_process owner.
- Shared: `background()` resolves the result promise early with `backgrounded: true` but keeps the execution active for continued streaming. `subscribe()` replays a snapshot for late joiners. `onExit()` fires on completion or replays from a 5-minute TTL cache.
2. **`shellExecutionService.ts`** — the main deletion site. All background-state maps (`activeResolvers`, `activeListeners`, `exitedPidInfo`) are removed. Shell now calls `attachExecution()` at spawn time and `completeWithResult()` at exit. The public `background()`, `subscribe()`, `onExit()`, `kill()`, `isActive()`, `writeInput()` methods become one-line delegates.
3. **`tool-executor.ts` + `coreToolHookTriggers.ts`** — the `setPidCallback``setExecutionIdCallback` rename chain. `ToolExecutor` no longer checks `instanceof ShellToolInvocation`; every tool gets the callback. `executeToolWithHooks` passes it through to `invocation.execute()`.
4. **`tools.ts` + `shell.ts`** — `ToolInvocation.execute()` signature gains optional `setExecutionIdCallback`. `BackgroundExecutionData` interface + `isBackgroundExecutionData` / `getBackgroundExecutionId` helpers added. Shell tool populates `data.executionId` alongside the existing `data.pid`.
5. **`useGeminiStream.ts` + `shellCommandProcessor.ts`** — UI generalization. The `activePtyId` computation no longer filters on `request.name === 'run_shell_command'`; any executing tool with a numeric `pid` (execution ID) qualifies for Ctrl+B. Background registration uses the core `isBackgroundExecutionData` / `getBackgroundExecutionId` helpers. Parameter renamed `activeToolPtyId``activeBackgroundExecutionId`.
### Key design decisions
- **Static class, not singleton instance**: `ExecutionLifecycleService` uses static methods and maps, matching the existing `ShellExecutionService` pattern. This avoids DI plumbing changes across the codebase.
- **`pid` field as backward-compat alias**: `ExecutionHandle.pid` and `BackgroundExecutionData.pid` remain for existing shell consumers. New code should use `executionId` / `getBackgroundExecutionId()`.
- **`NON_PROCESS_EXECUTION_ID_START = 2_000_000_000`**: Virtual execution IDs start above any realistic OS PID. `isActive()` short-circuits for IDs in this range to avoid spurious `process.kill()` syscalls.
- **No `finalizeExecution`**: The deprecated alias was removed (zero callers).
## Related Issues
<!-- Use keywords to auto-close issues (Closes #123, Fixes #456). If this PR is
only related to an issue or is a partial fix, simply reference the issue number
without a keyword (Related to #123). -->
## How to Validate
1. **Read the new service first**: `packages/core/src/services/executionLifecycleService.ts` — verify the create → stream → background → complete lifecycle makes sense in isolation.
2. **Verify shell delegation**: `shellExecutionService.ts` should have no background-state maps of its own. Every public lifecycle method should be a one-liner delegating to `ExecutionLifecycleService`.
3. **Run the targeted test suites**:
```bash
npm run test --workspace @google/gemini-cli-core -- --run \
src/services/executionLifecycleService.test.ts \
src/services/shellExecutionService.test.ts \
src/core/coreToolHookTriggers.test.ts \
src/scheduler/tool-executor.test.ts
npm run test --workspace @google/gemini-cli -- --run \
src/ui/hooks/useGeminiStream.test.tsx
```
4. **Typecheck and lint**:
```bash
npm run typecheck --workspace @google/gemini-cli-core
npm run lint --workspace @google/gemini-cli-core
npm run lint --workspace @google/gemini-cli
```
5. **Verify no behavioral change**: Existing shell backgrounding (Ctrl+B during a `run_shell_command`) should work identically. The UI background panel, `onExit` notifications, and kill behavior are unchanged.
## Pre-Merge Checklist
- [ ] Updated relevant documentation and README (if needed)
- [x] Added/updated tests (if needed)
- [ ] Noted breaking changes (if any)
- [ ] Validated on required platforms/methods:
- [ ] MacOS
- [ ] npm run
- [ ] npx
- [ ] Docker
- [ ] Podman
- [ ] Seatbelt
- [ ] Windows
- [ ] npm run
- [ ] npx
- [ ] Docker
- [ ] Linux
- [ ] npm run
- [ ] npx
- [ ] Docker