mirror of
https://github.com/google-gemini/gemini-cli.git
synced 2026-08-02 21:21:09 -07:00
Compare commits
20 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 3389a38d86 | |||
| bc76cbeb50 | |||
| 020fcc9674 | |||
| 8a90ecaf30 | |||
| db32a6185b | |||
| b3850edb8b | |||
| c77fd3fc7a | |||
| 013cfafbf9 | |||
| e17e34d18f | |||
| 228c978147 | |||
| 9b72826078 | |||
| 1e872ead2d | |||
| fff69df93e | |||
| 0fc4247287 | |||
| f2e2014894 | |||
| 6860284344 | |||
| 9a441ef927 | |||
| 3940d6344a | |||
| e9edd60615 | |||
| 6e291cfab8 |
@@ -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' }];
|
||||
|
||||
@@ -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', () => {
|
||||
|
||||
@@ -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);
|
||||
});
|
||||
});
|
||||
|
||||
@@ -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',
|
||||
|
||||
@@ -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
@@ -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 = '';
|
||||
|
||||
@@ -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.
|
||||
*/
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user