diff --git a/packages/core/src/services/executionLifecycleService.test.ts b/packages/core/src/services/executionLifecycleService.test.ts new file mode 100644 index 0000000000..8de1bc6ba6 --- /dev/null +++ b/packages/core/src/services/executionLifecycleService.test.ts @@ -0,0 +1,108 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, expect, it, vi } from 'vitest'; +import { + ShellExecutionService, + type ShellExecutionResult, +} from './shellExecutionService.js'; +import { + ExecutionLifecycleService, + type ExecutionCompletionOptions, +} from './executionLifecycleService.js'; + +const createResult = (): ShellExecutionResult => ({ + rawOutput: Buffer.from(''), + output: '', + exitCode: 0, + signal: null, + error: null, + aborted: false, + pid: 123, + executionMethod: 'none', +}); + +describe('ExecutionLifecycleService', () => { + it('creates executions through ShellExecutionService virtual execution API', () => { + const onKill = vi.fn(); + const handle = { + pid: 123, + result: Promise.resolve(createResult()), + }; + const createSpy = vi + .spyOn(ShellExecutionService, 'createVirtualExecution') + .mockReturnValue(handle); + + const created = ExecutionLifecycleService.createExecution('seed', onKill); + + expect(createSpy).toHaveBeenCalledWith('seed', onKill); + expect(created).toBe(handle); + }); + + it('delegates append and completion to ShellExecutionService virtual APIs', () => { + const appendSpy = vi.spyOn(ShellExecutionService, 'appendVirtualOutput'); + const completeSpy = vi.spyOn( + ShellExecutionService, + 'completeVirtualExecution', + ); + const options: ExecutionCompletionOptions = { + exitCode: 0, + signal: null, + }; + + ExecutionLifecycleService.appendOutput(123, 'delta'); + ExecutionLifecycleService.completeExecution(123, options); + + expect(appendSpy).toHaveBeenCalledWith(123, 'delta'); + expect(completeSpy).toHaveBeenCalledWith(123, options); + }); + + it('delegates backgrounding, subscriptions, exit callbacks, and kill', () => { + const unsubscribe = vi.fn(); + const backgroundSpy = vi + .spyOn(ShellExecutionService, 'background') + .mockImplementation(() => {}); + const subscribeSpy = vi + .spyOn(ShellExecutionService, 'subscribe') + .mockReturnValue(unsubscribe); + const onExitSpy = vi + .spyOn(ShellExecutionService, 'onExit') + .mockReturnValue(unsubscribe); + const killSpy = vi + .spyOn(ShellExecutionService, 'kill') + .mockImplementation(() => {}); + + const listener = vi.fn(); + const onExit = vi.fn(); + const returnedSub = ExecutionLifecycleService.subscribe(123, listener); + const returnedExit = ExecutionLifecycleService.onExit(123, onExit); + ExecutionLifecycleService.background(123); + ExecutionLifecycleService.kill(123); + + expect(subscribeSpy).toHaveBeenCalledWith(123, listener); + expect(onExitSpy).toHaveBeenCalledWith(123, onExit); + expect(backgroundSpy).toHaveBeenCalledWith(123); + expect(killSpy).toHaveBeenCalledWith(123); + expect(returnedSub).toBe(unsubscribe); + expect(returnedExit).toBe(unsubscribe); + }); + + it('delegates active checks and input writes', () => { + const isActiveSpy = vi + .spyOn(ShellExecutionService, 'isPtyActive') + .mockReturnValue(true); + const writeSpy = vi + .spyOn(ShellExecutionService, 'writeToPty') + .mockImplementation(() => {}); + + const isActive = ExecutionLifecycleService.isActive(123); + ExecutionLifecycleService.writeInput(123, 'input'); + + expect(isActiveSpy).toHaveBeenCalledWith(123); + expect(writeSpy).toHaveBeenCalledWith(123, 'input'); + expect(isActive).toBe(true); + }); +}); diff --git a/packages/core/src/services/executionLifecycleService.ts b/packages/core/src/services/executionLifecycleService.ts new file mode 100644 index 0000000000..1fe94402d1 --- /dev/null +++ b/packages/core/src/services/executionLifecycleService.ts @@ -0,0 +1,74 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { + ShellExecutionService, + type ShellExecutionHandle, + type ShellOutputEvent, +} from './shellExecutionService.js'; + +export interface ExecutionCompletionOptions { + exitCode?: number | null; + signal?: number | null; + error?: Error | null; + aborted?: boolean; +} + +/** + * Generic lifecycle facade for backgroundable executions. + * + * This wraps ShellExecutionService so non-shell executors (remote/local agents) + * can use neutral lifecycle naming without duplicating process-management logic. + */ +export class ExecutionLifecycleService { + static createExecution( + initialOutput = '', + onKill?: () => void, + ): ShellExecutionHandle { + return ShellExecutionService.createVirtualExecution(initialOutput, onKill); + } + + static appendOutput(executionId: number, chunk: string): void { + ShellExecutionService.appendVirtualOutput(executionId, chunk); + } + + static completeExecution( + executionId: number, + options?: ExecutionCompletionOptions, + ): void { + ShellExecutionService.completeVirtualExecution(executionId, options); + } + + static background(executionId: number): void { + ShellExecutionService.background(executionId); + } + + static subscribe( + executionId: number, + listener: (event: ShellOutputEvent) => void, + ): () => void { + return ShellExecutionService.subscribe(executionId, listener); + } + + static onExit( + executionId: number, + callback: (exitCode: number, signal?: number) => void, + ): () => void { + return ShellExecutionService.onExit(executionId, callback); + } + + static kill(executionId: number): void { + ShellExecutionService.kill(executionId); + } + + static isActive(executionId: number): boolean { + return ShellExecutionService.isPtyActive(executionId); + } + + static writeInput(executionId: number, input: string): void { + ShellExecutionService.writeToPty(executionId, input); + } +} diff --git a/packages/core/src/services/shellExecutionService.test.ts b/packages/core/src/services/shellExecutionService.test.ts index 77de13de3a..0d6194888f 100644 --- a/packages/core/src/services/shellExecutionService.test.ts +++ b/packages/core/src/services/shellExecutionService.test.ts @@ -432,13 +432,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 () => { @@ -1633,3 +1641,85 @@ describe('ShellExecutionService environment variables', () => { await new Promise(process.nextTick); }); }); + +describe('ShellExecutionService virtual executions', () => { + it('completes a virtual execution in the foreground', async () => { + const { pid, result } = ShellExecutionService.createVirtualExecution(); + if (pid === undefined) { + throw new Error('Expected virtual pid to be defined.'); + } + const onExit = vi.fn(); + const unsubscribe = ShellExecutionService.onExit(pid, onExit); + + ShellExecutionService.appendVirtualOutput(pid, 'Hello'); + ShellExecutionService.appendVirtualOutput(pid, ' World'); + ShellExecutionService.completeVirtualExecution(pid, { exitCode: 0 }); + + const executionResult = await result; + + expect(executionResult.output).toBe('Hello World'); + expect(executionResult.backgrounded).toBeUndefined(); + expect(executionResult.exitCode).toBe(0); + expect(executionResult.error).toBeNull(); + + await vi.waitFor(() => { + expect(onExit).toHaveBeenCalledWith(0, undefined); + }); + + unsubscribe(); + }); + + it('supports backgrounding virtual executions and streaming additional output', async () => { + const { pid, result } = ShellExecutionService.createVirtualExecution(); + if (pid === undefined) { + throw new Error('Expected virtual pid to be defined.'); + } + const chunks: string[] = []; + const onExit = vi.fn(); + + const unsubscribeStream = ShellExecutionService.subscribe(pid, (event) => { + if (event.type === 'data' && typeof event.chunk === 'string') { + chunks.push(event.chunk); + } + }); + const unsubscribeExit = ShellExecutionService.onExit(pid, onExit); + + ShellExecutionService.appendVirtualOutput(pid, 'Chunk 1'); + ShellExecutionService.background(pid); + + const backgroundResult = await result; + expect(backgroundResult.backgrounded).toBe(true); + expect(backgroundResult.output).toBe('Chunk 1'); + + ShellExecutionService.appendVirtualOutput(pid, '\nChunk 2'); + ShellExecutionService.completeVirtualExecution(pid, { exitCode: 0 }); + + await vi.waitFor(() => { + expect(chunks.join('')).toContain('Chunk 2'); + expect(onExit).toHaveBeenCalledWith(0, undefined); + }); + + unsubscribeStream(); + unsubscribeExit(); + }); + + it('kills virtual executions via the existing kill API', async () => { + const onKill = vi.fn(); + const { pid, result } = ShellExecutionService.createVirtualExecution( + '', + onKill, + ); + if (pid === undefined) { + throw new Error('Expected virtual pid to be defined.'); + } + + ShellExecutionService.appendVirtualOutput(pid, 'work'); + ShellExecutionService.kill(pid); + + const killResult = await result; + expect(onKill).toHaveBeenCalledTimes(1); + expect(killResult.aborted).toBe(true); + expect(killResult.exitCode).toBe(130); + expect(killResult.error?.message).toContain('Operation cancelled by user'); + }); +}); diff --git a/packages/core/src/services/shellExecutionService.ts b/packages/core/src/services/shellExecutionService.ts index fdb2ca79b5..397dba1d9b 100644 --- a/packages/core/src/services/shellExecutionService.ts +++ b/packages/core/src/services/shellExecutionService.ts @@ -152,6 +152,22 @@ interface ActiveChildProcess { }; } +interface ActiveVirtualProcessState { + output: string; + onKill?: () => void; +} + +type ActiveManagedProcess = + | { + kind: 'child'; + process: ChildProcess; + state: ActiveChildProcess['state']; + } + | { + kind: 'virtual'; + state: ActiveVirtualProcessState; + }; + const getFullBufferText = (terminal: pkg.Terminal): string => { const buffer = terminal.buffer.active; const lines: string[] = []; @@ -197,7 +213,7 @@ const getFullBufferText = (terminal: pkg.Terminal): string => { export class ShellExecutionService { private static activePtys = new Map(); - private static activeChildProcesses = new Map(); + private static activeProcesses = new Map(); private static exitedPtyInfo = new Map< number, { exitCode: number; signal?: number } @@ -210,6 +226,7 @@ export class ShellExecutionService { number, Set<(event: ShellOutputEvent) => void> >(); + private static nextVirtualPid = 2_000_000_000; /** * Executes a shell command using `node-pty`, capturing all output and lifecycle events. * @@ -292,6 +309,119 @@ export class ShellExecutionService { } } + private static getActiveChildProcess( + pid: number, + ): ActiveChildProcess | undefined { + const activeProcess = this.activeProcesses.get(pid); + if (activeProcess?.kind !== 'child') { + return undefined; + } + return { process: activeProcess.process, state: activeProcess.state }; + } + + private static getActiveVirtualProcess( + pid: number, + ): ActiveVirtualProcessState | undefined { + const activeProcess = this.activeProcesses.get(pid); + if (activeProcess?.kind !== 'virtual') { + return undefined; + } + return activeProcess.state; + } + + private static allocateVirtualPid(): number { + let pid = ++this.nextVirtualPid; + while (this.activePtys.has(pid) || this.activeProcesses.has(pid)) { + pid = ++this.nextVirtualPid; + } + return pid; + } + + static createVirtualExecution( + initialOutput = '', + onKill?: () => void, + ): ShellExecutionHandle { + const pid = this.allocateVirtualPid(); + this.activeProcesses.set(pid, { + kind: 'virtual', + state: { + output: initialOutput, + onKill, + }, + }); + + const result = new Promise((resolve) => { + this.activeResolvers.set(pid, resolve); + }); + + return { pid, result }; + } + + static appendVirtualOutput(pid: number, chunk: string): void { + const virtual = this.getActiveVirtualProcess(pid); + if (!virtual || chunk.length === 0) { + return; + } + virtual.output += chunk; + this.emitEvent(pid, { type: 'data', chunk }); + } + + static completeVirtualExecution( + pid: number, + options?: { + exitCode?: number | null; + signal?: number | null; + error?: Error | null; + aborted?: boolean; + }, + ): void { + const virtual = this.getActiveVirtualProcess(pid); + if (!virtual) { + return; + } + + const { + error = null, + aborted = false, + exitCode = error ? 1 : 0, + signal = null, + } = options ?? {}; + + const resolve = this.activeResolvers.get(pid); + if (resolve) { + resolve({ + rawOutput: Buffer.from(virtual.output, 'utf8'), + output: virtual.output, + exitCode, + signal, + error, + aborted, + pid, + executionMethod: 'none', + }); + this.activeResolvers.delete(pid); + } + + this.emitEvent(pid, { + type: 'exit', + exitCode, + signal, + }); + this.activeListeners.delete(pid); + this.activeProcesses.delete(pid); + + this.exitedPtyInfo.set(pid, { + exitCode: exitCode ?? 0, + signal: signal ?? undefined, + }); + setTimeout( + () => { + this.exitedPtyInfo.delete(pid); + }, + 5 * 60 * 1000, + ).unref(); + } + private static childProcessFallback( commandToExecute: string, cwd: string, @@ -328,7 +458,8 @@ export class ShellExecutionService { }; if (child.pid) { - this.activeChildProcesses.set(child.pid, { + this.activeProcesses.set(child.pid, { + kind: 'child', process: child, state, }); @@ -438,7 +569,7 @@ export class ShellExecutionService { onOutputEvent(event); ShellExecutionService.emitEvent(child.pid, event); - this.activeChildProcesses.delete(child.pid); + this.activeProcesses.delete(child.pid); this.activeResolvers.delete(child.pid); this.activeListeners.delete(child.pid); } @@ -914,11 +1045,9 @@ export class ShellExecutionService { * @param input The string to write to the terminal. */ static writeToPty(pid: number, input: string): void { - if (this.activeChildProcesses.has(pid)) { - const activeChild = this.activeChildProcesses.get(pid); - if (activeChild) { - activeChild.process.stdin?.write(input); - } + const activeChild = this.getActiveChildProcess(pid); + if (activeChild) { + activeChild.process.stdin?.write(input); return; } @@ -933,7 +1062,11 @@ export class ShellExecutionService { } static isPtyActive(pid: number): boolean { - if (this.activeChildProcesses.has(pid)) { + if (this.getActiveVirtualProcess(pid)) { + return true; + } + + if (this.getActiveChildProcess(pid)) { try { return process.kill(pid, 0); } catch { @@ -971,8 +1104,10 @@ export class ShellExecutionService { }, ); return () => disposable.dispose(); - } else if (this.activeChildProcesses.has(pid)) { - const activeChild = this.activeChildProcesses.get(pid); + } + + const activeChild = this.getActiveChildProcess(pid); + if (activeChild) { const listener = (code: number | null, signal: NodeJS.Signals | null) => { let signalNumber: number | undefined; if (signal) { @@ -984,14 +1119,25 @@ export class ShellExecutionService { return () => { activeChild?.process.removeListener('exit', listener); }; - } else { - // Check if it already exited recently - const exitedInfo = this.exitedPtyInfo.get(pid); - if (exitedInfo) { - callback(exitedInfo.exitCode, exitedInfo.signal); - } - return () => {}; } + + if (this.getActiveVirtualProcess(pid)) { + const listener = (event: ShellOutputEvent) => { + if (event.type === 'exit') { + callback(event.exitCode ?? 0, event.signal ?? undefined); + unsubscribe(); + } + }; + const unsubscribe = this.subscribe(pid, listener); + return unsubscribe; + } + + // Check if it already exited recently + const exitedInfo = this.exitedPtyInfo.get(pid); + if (exitedInfo) { + callback(exitedInfo.exitCode, exitedInfo.signal); + } + return () => {}; } /** @@ -1001,11 +1147,20 @@ export class ShellExecutionService { */ static kill(pid: number): void { const activePty = this.activePtys.get(pid); - const activeChild = this.activeChildProcesses.get(pid); + const activeChild = this.getActiveChildProcess(pid); + const activeVirtual = this.getActiveVirtualProcess(pid); - if (activeChild) { + if (activeVirtual) { + activeVirtual.onKill?.(); + this.completeVirtualExecution(pid, { + error: new Error('Operation cancelled by user.'), + aborted: true, + exitCode: 130, + }); + return; + } else if (activeChild) { killProcessGroup({ pid }).catch(() => {}); - this.activeChildProcesses.delete(pid); + this.activeProcesses.delete(pid); } else if (activePty) { killProcessGroup({ pid, pty: activePty.ptyProcess }).catch(() => {}); this.activePtys.delete(pid); @@ -1028,7 +1183,8 @@ export class ShellExecutionService { const rawOutput = Buffer.from(''); const activePty = this.activePtys.get(pid); - const activeChild = this.activeChildProcesses.get(pid); + const activeChild = this.getActiveChildProcess(pid); + const activeVirtual = this.getActiveVirtualProcess(pid); if (activePty) { output = getFullBufferText(activePty.headlessTerminal); @@ -1057,6 +1213,19 @@ export class ShellExecutionService { executionMethod: 'child_process', backgrounded: true, }); + } else if (activeVirtual) { + output = activeVirtual.output; + resolve({ + rawOutput, + output, + exitCode: null, + signal: null, + error: null, + aborted: false, + pid, + executionMethod: 'none', + backgrounded: true, + }); } this.activeResolvers.delete(pid); @@ -1074,7 +1243,8 @@ export class ShellExecutionService { // Send current buffer content immediately const activePty = this.activePtys.get(pid); - const activeChild = this.activeChildProcesses.get(pid); + const activeChild = this.getActiveChildProcess(pid); + const activeVirtual = this.getActiveVirtualProcess(pid); if (activePty) { // Use serializeTerminalToObject to preserve colors and structure @@ -1096,6 +1266,8 @@ export class ShellExecutionService { if (output) { listener({ type: 'data', chunk: output }); } + } else if (activeVirtual?.output) { + listener({ type: 'data', chunk: activeVirtual.output }); } return () => {