diff --git a/dev/src/cli/cli.ts b/dev/src/cli/cli.ts index f227c979f..bfb98ee15 100644 --- a/dev/src/cli/cli.ts +++ b/dev/src/cli/cli.ts @@ -239,6 +239,7 @@ export function createProgram(): Command { host: options['host'], port: parseInt(options['port'], 10), serveDebugUI: true, + installProcessHandlers: true, allowOrigins: options['allow_origins'], sessionService: getSessionServiceFromOptions(options), artifactService: getArtifactServiceFromOptions(options), @@ -285,6 +286,7 @@ export function createProgram(): Command { host: options['host'], port: parseInt(options['port'], 10), serveDebugUI: false, + installProcessHandlers: true, allowOrigins: options['allow_origins'], sessionService: getSessionServiceFromOptions(options), artifactService: getArtifactServiceFromOptions(options), diff --git a/dev/src/cli/deploy/cli_deploy_agent_engine.ts b/dev/src/cli/deploy/cli_deploy_agent_engine.ts index 434ed04d1..79ee3205c 100644 --- a/dev/src/cli/deploy/cli_deploy_agent_engine.ts +++ b/dev/src/cli/deploy/cli_deploy_agent_engine.ts @@ -66,6 +66,7 @@ export async function deployToAgentEngine(options: DeployToAgentEngineOptions) { options.agentPath, options.agentFileLoadOptions, ); + agentLoader.installProcessHandlers(); const isFileProvided = await isFile(options.agentPath); const agentDir = isFileProvided diff --git a/dev/src/cli/deploy/cli_deploy_cloud_run.ts b/dev/src/cli/deploy/cli_deploy_cloud_run.ts index ad55ab8fd..5e6337438 100644 --- a/dev/src/cli/deploy/cli_deploy_cloud_run.ts +++ b/dev/src/cli/deploy/cli_deploy_cloud_run.ts @@ -169,6 +169,7 @@ export async function deployToCloudRun(options: DeployToCloudRunOptions) { options.agentPath, options.agentFileLoadOptions, ); + agentLoader.installProcessHandlers(); const isFileProvided = await isFile(options.agentPath); const agentDir = isFileProvided diff --git a/dev/src/server/adk_api_server.ts b/dev/src/server/adk_api_server.ts index a52ad63ab..4f5fdd7b4 100644 --- a/dev/src/server/adk_api_server.ts +++ b/dev/src/server/adk_api_server.ts @@ -58,6 +58,13 @@ interface ServerOptions { memoryService?: BaseMemoryService; artifactService?: BaseArtifactService; agentLoader?: AgentLoader; + /** + * Installs process exit and signal handlers on the agent loader. Only a CLI + * entrypoint that owns the process should enable this: the handlers call + * `process.exit()`. Defaults to false so an embedded server never mutates + * the host process. + */ + installProcessHandlers?: boolean; agentFileLoadOptions?: AgentFileOptions; serveDebugUI?: boolean; allowOrigins?: string; @@ -146,6 +153,10 @@ export class AdkApiServer { this.a2aAuthToken = options.a2aAuthToken || process.env[A2A_AUTH_TOKEN_ENV_VAR] || undefined; this.app = express(); + + if (options.installProcessHandlers) { + this.agentLoader.installProcessHandlers(); + } } private async setupTelemetry(): Promise { diff --git a/dev/src/utils/agent_loader.ts b/dev/src/utils/agent_loader.ts index 448a32aeb..fa821a84b 100644 --- a/dev/src/utils/agent_loader.ts +++ b/dev/src/utils/agent_loader.ts @@ -49,6 +49,11 @@ const FILE_MODULE_TYPE_EXTENSION_MAP = { [FileModuleType.ESM]: '.mjs', }; +/** + * Termination signals wired up by {@link AgentLoader.installProcessHandlers}. + */ +const TERMINATION_SIGNALS = ['SIGINT', 'SIGUSR1', 'SIGUSR2'] as const; + /** * Metadata for a file. */ @@ -361,34 +366,44 @@ export class AgentLoader { private agentsAlreadyPreloaded = false; private readonly preloadedAgents: Record = {}; private watcher?: fs.FSWatcher; + private removeProcessHandlers?: () => void; constructor( private readonly agentsDirPath: string = process.cwd(), private readonly options = DEFAULT_AGENT_FILE_OPTIONS, private readonly watchForChanges = false, - ) { - // Do cleanups on exit - const exitHandler = async ({ - exit, - cleanup, - }: { - exit?: boolean; - cleanup?: boolean; - }) => { - if (cleanup) { - await this.disposeAll(); - } + ) {} + + /** + * Wires process exit and termination signals to this loader's cleanup. + * + * This mutates global process state and terminates the process on a + * termination signal, so it is only appropriate for a CLI entrypoint that + * owns the process. Library and test consumers must instead call + * {@link disposeAll} when they are done. Calling this twice is a no-op, and + * {@link disposeAll} removes the listeners again. + */ + installProcessHandlers(): void { + if (this.removeProcessHandlers) { + return; + } + + // An `exit` listener cannot await, so this cleanup is best-effort. + const onExit = () => void this.disposeAll(); + const onSignal = () => process.exit(); - if (exit) { - process.exit(); + process.on('exit', onExit); + for (const signal of TERMINATION_SIGNALS) { + process.on(signal, onSignal); + } + + this.removeProcessHandlers = () => { + process.removeListener('exit', onExit); + for (const signal of TERMINATION_SIGNALS) { + process.removeListener(signal, onSignal); } + this.removeProcessHandlers = undefined; }; - - process.on('exit', () => exitHandler({cleanup: true})); - process.on('SIGINT', () => exitHandler({exit: true})); - process.on('SIGUSR1', () => exitHandler({exit: true})); - process.on('SIGUSR2', () => exitHandler({exit: true})); - process.on('uncaughtException', () => exitHandler({exit: true})); } /** @@ -471,6 +486,7 @@ export class AgentLoader { } async disposeAll(): Promise { + this.removeProcessHandlers?.(); this.watcher?.close(); this.watcher = undefined; await Promise.all( diff --git a/dev/test/cli/cli_deploy_agent_engine_test.ts b/dev/test/cli/cli_deploy_agent_engine_test.ts index f20519350..d2f82e430 100644 --- a/dev/test/cli/cli_deploy_agent_engine_test.ts +++ b/dev/test/cli/cli_deploy_agent_engine_test.ts @@ -126,6 +126,7 @@ vi.mock('../../src/utils/agent_loader.js', () => ({ getFilePath: vi.fn().mockReturnValue('path/to/agent1.ts'), }), disposeAll: vi.fn().mockResolvedValue(undefined), + installProcessHandlers: vi.fn(), })), })); @@ -196,6 +197,7 @@ describe('deployToAgentEngine', () => { getFilePath: vi.fn().mockReturnValue('path/to/agent1.ts'), }), disposeAll: vi.fn().mockResolvedValue(undefined), + installProcessHandlers: vi.fn(), })); execMock.mockImplementation((cmd: string, callback: Callback) => { @@ -300,6 +302,13 @@ describe('deployToAgentEngine', () => { expect(exists).toBe(false); }); + it('installs process handlers on the agent loader', async () => { + await deployToAgentEngine(defaultOptions); + + const agentLoader = vi.mocked(AgentLoader).mock.results[0].value; + expect(agentLoader.installProcessHandlers).toHaveBeenCalledOnce(); + }); + it('should deploy successfully with all optional parameters', async () => { const optionsWithAll: DeployToAgentEngineOptions = { ...defaultOptions, diff --git a/dev/test/cli/cli_deploy_cloud_run_test.ts b/dev/test/cli/cli_deploy_cloud_run_test.ts index e79bed848..bb8d0f7ea 100644 --- a/dev/test/cli/cli_deploy_cloud_run_test.ts +++ b/dev/test/cli/cli_deploy_cloud_run_test.ts @@ -53,6 +53,7 @@ vi.mock('../../src/utils/agent_loader.js', () => ({ getFilePath: vi.fn().mockReturnValue('path/to/agent1.ts'), }), disposeAll: vi.fn().mockResolvedValue(undefined), + installProcessHandlers: vi.fn(), })), })); @@ -151,6 +152,7 @@ describe('deployToCloudRun', () => { getFilePath: vi.fn().mockReturnValue('path/to/agent1.ts'), }), disposeAll: vi.fn().mockResolvedValue(undefined), + installProcessHandlers: vi.fn(), })); execMock.mockImplementation((cmd: string, callback: Callback) => { @@ -198,6 +200,13 @@ describe('deployToCloudRun', () => { }); }); + it('installs process handlers on the agent loader', async () => { + await deployToCloudRun(defaultOptions); + + const agentLoader = vi.mocked(AgentLoader).mock.results[0].value; + expect(agentLoader.installProcessHandlers).toHaveBeenCalledOnce(); + }); + it('should resolve default project and region from gcloud if not provided', async () => { const optionsWithoutProjectRegion = { ...defaultOptions, diff --git a/dev/test/cli/cli_test.ts b/dev/test/cli/cli_test.ts index 75d003802..54b249eb6 100644 --- a/dev/test/cli/cli_test.ts +++ b/dev/test/cli/cli_test.ts @@ -154,6 +154,13 @@ describe('CLI Entrypoint', () => { const args = (AdkApiServer as unknown as Mock).mock.calls[0][0]; expect(args.a2aAuthToken).toBe('tok'); }); + + it('should opt the server into process handlers', async () => { + await parse(['web']); + + const args = vi.mocked(AdkApiServer).mock.calls[0][0]; + expect(args.installProcessHandlers).toBe(true); + }); }); describe('command: api_server', () => { @@ -181,6 +188,13 @@ describe('CLI Entrypoint', () => { const args = (AdkApiServer as unknown as Mock).mock.calls[0][0]; expect(args.a2aAuthToken).toBe('tok'); }); + + it('should opt the server into process handlers', async () => { + await parse(['api_server']); + + const args = vi.mocked(AdkApiServer).mock.calls[0][0]; + expect(args.installProcessHandlers).toBe(true); + }); }); describe('command: create', () => { diff --git a/dev/test/server/adk_api_server_test.ts b/dev/test/server/adk_api_server_test.ts index d6db712d2..9540ca4a4 100644 --- a/dev/test/server/adk_api_server_test.ts +++ b/dev/test/server/adk_api_server_test.ts @@ -1255,4 +1255,32 @@ describe('AdkWebServer', () => { } }); }); + + describe('process handlers', () => { + it('does not install agent loader process handlers by default', () => { + const loader = new AgentLoader(process.cwd()); + const installProcessHandlers = vi.spyOn(loader, 'installProcessHandlers'); + const signalListeners = process.listenerCount('SIGINT'); + + new AdkApiServer({agentLoader: loader}); + + expect(installProcessHandlers).not.toHaveBeenCalled(); + expect(process.listenerCount('SIGINT')).toBe(signalListeners); + }); + + it('installs agent loader process handlers when opted in', async () => { + const loader = new AgentLoader(process.cwd()); + const installProcessHandlers = vi.spyOn(loader, 'installProcessHandlers'); + const signalListeners = process.listenerCount('SIGINT'); + + try { + new AdkApiServer({agentLoader: loader, installProcessHandlers: true}); + + expect(installProcessHandlers).toHaveBeenCalledOnce(); + expect(process.listenerCount('SIGINT')).toBe(signalListeners + 1); + } finally { + await loader.disposeAll(); + } + }); + }); }); diff --git a/dev/test/utils/agent_loader_test.ts b/dev/test/utils/agent_loader_test.ts index d53160d9d..c250aade6 100644 --- a/dev/test/utils/agent_loader_test.ts +++ b/dev/test/utils/agent_loader_test.ts @@ -791,4 +791,175 @@ describe('AgentLoader', () => { await loader.disposeAll(); }); }); + + describe('process handlers', () => { + // Never `process.emit` any of these events: that also fires Vitest's and + // Tinypool's own listeners and can take the worker down. These tests + // capture the listener the loader installed and invoke it directly. + const PROCESS_EVENTS = [ + 'exit', + 'SIGINT', + 'SIGUSR1', + 'SIGUSR2', + 'uncaughtException', + ] as const; + + const counts = () => + PROCESS_EVENTS.map((event) => process.listenerCount(event)); + + const afterInstall = (before: number[]) => + before.map((count, index) => + PROCESS_EVENTS[index] === 'uncaughtException' ? count : count + 1, + ); + + const listenerAddedBy = ( + listenersOf: () => T[], + install: () => void, + ) => { + const before = new Set(listenersOf()); + install(); + const added = listenersOf().filter((l) => !before.has(l)); + expect(added).toHaveLength(1); + return added[0]; + }; + + it('does not register process listeners in the constructor', async () => { + const before = counts(); + const loaders = [ + new AgentLoader(tempAgentsDir), + new AgentLoader(tempAgentsDir), + new AgentLoader(tempAgentsDir), + ]; + + try { + expect(counts()).toEqual(before); + } finally { + await Promise.all(loaders.map((loader) => loader.disposeAll())); + } + }); + + it('installs exit and termination signal listeners on demand', async () => { + const loader = new AgentLoader(tempAgentsDir); + const before = counts(); + + try { + loader.installProcessHandlers(); + + expect(counts()).toEqual(afterInstall(before)); + } finally { + await loader.disposeAll(); + } + }); + + it('never installs an uncaughtException listener', async () => { + const before = process.listenerCount('uncaughtException'); + const loader = new AgentLoader(tempAgentsDir); + + try { + loader.installProcessHandlers(); + + expect(process.listenerCount('uncaughtException')).toBe(before); + } finally { + await loader.disposeAll(); + } + }); + + it('is idempotent', async () => { + const loader = new AgentLoader(tempAgentsDir); + const before = counts(); + + try { + loader.installProcessHandlers(); + loader.installProcessHandlers(); + + expect(counts()).toEqual(afterInstall(before)); + } finally { + await loader.disposeAll(); + } + }); + + it('removes installed listeners on disposeAll', async () => { + const loader = new AgentLoader(tempAgentsDir); + const before = counts(); + + try { + loader.installProcessHandlers(); + await loader.disposeAll(); + + expect(counts()).toEqual(before); + } finally { + await loader.disposeAll(); + } + }); + + it('can reinstall after disposeAll', async () => { + const loader = new AgentLoader(tempAgentsDir); + const before = counts(); + + try { + loader.installProcessHandlers(); + await loader.disposeAll(); + loader.installProcessHandlers(); + + expect(counts()).toEqual(afterInstall(before)); + + await loader.disposeAll(); + + expect(counts()).toEqual(before); + } finally { + await loader.disposeAll(); + } + }); + + it('disposeAll is a no-op for handlers that were never installed', async () => { + const before = counts(); + const loader = new AgentLoader(tempAgentsDir); + + await loader.disposeAll(); + + expect(counts()).toEqual(before); + }); + + it('the exit listener disposes cached agents', async () => { + const loader = new AgentLoader(tempAgentsDir); + const disposeAll = vi.spyOn(loader, 'disposeAll'); + + try { + const exitListener = listenerAddedBy( + () => process.listeners('exit'), + () => { + loader.installProcessHandlers(); + }, + ); + exitListener(0); + + expect(disposeAll).toHaveBeenCalled(); + } finally { + disposeAll.mockRestore(); + await loader.disposeAll(); + } + }); + + it('a termination signal listener exits the process', async () => { + const loader = new AgentLoader(tempAgentsDir); + const exit = vi.spyOn(process, 'exit').mockImplementation(() => { + throw new Error('process.exit'); + }); + + try { + const signalListener = listenerAddedBy( + () => process.listeners('SIGINT'), + () => { + loader.installProcessHandlers(); + }, + ); + + expect(() => signalListener('SIGINT')).toThrow('process.exit'); + expect(exit).toHaveBeenCalled(); + } finally { + exit.mockRestore(); + await loader.disposeAll(); + } + }); + }); });