From 4b8d5e76243b26d99d361873d799f6ad422e1578 Mon Sep 17 00:00:00 2001 From: Coco Sheng Date: Tue, 28 Apr 2026 13:29:49 -0400 Subject: [PATCH] fix(cli): prevent ACP stdout pollution from SessionEnd hooks (#26125) --- packages/cli/src/gemini.tsx | 8 +- packages/cli/src/gemini_cleanup.test.tsx | 126 ++++++++++++++++++++++- 2 files changed, 132 insertions(+), 2 deletions(-) diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx index e7ec0b1a66..727a0af9c7 100644 --- a/packages/cli/src/gemini.tsx +++ b/packages/cli/src/gemini.tsx @@ -385,7 +385,6 @@ export async function main() { }, }); consolePatcher.patch(); - registerCleanup(consolePatcher.cleanup); dns.setDefaultResultOrder( validateDnsResolutionOrder(settings.merged.advanced.dnsResolutionOrder), @@ -411,6 +410,7 @@ export async function main() { const partialConfig = await loadCliConfig(settings.merged, sessionId, argv, { projectHooks: settings.workspace.settings.hooks, }); + adminControlsListner.setConfig(partialConfig); // Refresh auth to fetch remote admin settings from CCPA and before entering @@ -568,6 +568,12 @@ export async function main() { await config.getHookSystem()?.fireSessionEndEvent(SessionEndReason.Exit); }); + // Register ConsolePatcher cleanup last to ensure logs from shutdown hooks + // are correctly redirected to stderr (especially for non-interactive JSON output). + if (!config.getAcpMode()) { + registerCleanup(consolePatcher.cleanup); + } + // Launch cleanup expired sessions as a background task cleanupExpiredSessions(config, settings.merged).catch((e) => { debugLogger.error('Failed to cleanup expired sessions:', e); diff --git a/packages/cli/src/gemini_cleanup.test.tsx b/packages/cli/src/gemini_cleanup.test.tsx index b4c4f645b6..b76a132234 100644 --- a/packages/cli/src/gemini_cleanup.test.tsx +++ b/packages/cli/src/gemini_cleanup.test.tsx @@ -68,6 +68,13 @@ vi.mock('./config/settings.js', async (importOriginal) => { }; }); +vi.mock('./ui/utils/ConsolePatcher.js', () => ({ + ConsolePatcher: vi.fn().mockImplementation(() => ({ + patch: vi.fn(), + cleanup: vi.fn(), + })), +})); + vi.mock('./config/config.js', () => ({ loadCliConfig: vi.fn().mockResolvedValue({ getSandbox: vi.fn(() => false), @@ -150,6 +157,10 @@ vi.mock('./utils/cleanup.js', async (importOriginal) => { }; }); +vi.mock('./acp/acpClient.js', () => ({ + runAcpClient: vi.fn().mockResolvedValue(undefined), +})); + vi.mock('./zed-integration/zedIntegration.js', () => ({ runZedIntegration: vi.fn().mockResolvedValue(undefined), })); @@ -296,6 +307,120 @@ describe('gemini.tsx main function cleanup', () => { ); }); + it('should not register ConsolePatcher cleanup in ACP mode', async () => { + const { registerCleanup } = await import('./utils/cleanup.js'); + const { ConsolePatcher } = await import('./ui/utils/ConsolePatcher.js'); + const { loadCliConfig, parseArguments } = await import( + './config/config.js' + ); + const { loadSettings } = await import('./config/settings.js'); + + vi.mocked(parseArguments).mockResolvedValue({ + acp: true, + startupMessages: [], + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } as any); + + vi.mocked(loadSettings).mockReturnValue({ + merged: { + tools: { allowed: [], exclude: [] }, + advanced: { dnsResolutionOrder: 'ipv4first' }, + security: { auth: { selectedType: 'google' } }, + ui: { theme: 'default' }, + }, + workspace: { settings: {} }, + errors: [], + subscribe: vi.fn(), + getSnapshot: vi.fn(), + setValue: vi.fn(), + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } as any); + + vi.mocked(loadCliConfig).mockResolvedValue( + buildMockConfig({ + getAcpMode: () => true, + }), + ); + + let capturedCleanup: () => void; + vi.mocked(ConsolePatcher).mockImplementation(() => { + const instance = { + patch: vi.fn(), + cleanup: vi.fn(), + }; + capturedCleanup = instance.cleanup; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + return instance as any; + }); + + await main(); + + const registeredFunctions = vi + .mocked(registerCleanup) + .mock.calls.map((call) => call[0]); + expect(registeredFunctions).not.toContain(capturedCleanup!); + }); + + it('should register ConsolePatcher cleanup in non-ACP mode', async () => { + const { registerCleanup } = await import('./utils/cleanup.js'); + const { ConsolePatcher } = await import('./ui/utils/ConsolePatcher.js'); + const { loadCliConfig, parseArguments } = await import( + './config/config.js' + ); + const { loadSettings } = await import('./config/settings.js'); + + vi.mocked(parseArguments).mockResolvedValue({ + acp: false, + query: 'test', + startupMessages: [], + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } as any); + + vi.mocked(loadSettings).mockReturnValue({ + merged: { + tools: { allowed: [], exclude: [] }, + advanced: { dnsResolutionOrder: 'ipv4first' }, + security: { auth: { selectedType: 'google' } }, + ui: { theme: 'default' }, + }, + workspace: { settings: {} }, + errors: [], + subscribe: vi.fn(), + getSnapshot: vi.fn(), + setValue: vi.fn(), + // eslint-disable-next-line @typescript-eslint/no-explicit-any + } as any); + + vi.mocked(loadCliConfig).mockResolvedValue( + buildMockConfig({ + getAcpMode: () => false, + getQuestion: () => 'test', + }), + ); + + let capturedCleanup: () => void; + vi.mocked(ConsolePatcher).mockImplementation(() => { + const instance = { + patch: vi.fn(), + cleanup: vi.fn(), + }; + capturedCleanup = instance.cleanup; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + return instance as any; + }); + + try { + await main(); + } catch { + // Ignore errors from incomplete mocks in full main() execution + } + + const registeredFunctions = vi + .mocked(registerCleanup) + .mock.calls.map((call) => call[0]); + expect(registeredFunctions).toContain(capturedCleanup!); + }); + function buildMockConfig(overrides: Partial = {}): Config { return { isInteractive: vi.fn(() => false), @@ -319,7 +444,6 @@ describe('gemini.tsx main function cleanup', () => { getListExtensions: vi.fn(() => false), getListSessions: vi.fn(() => false), getDeleteSession: vi.fn(() => undefined), - getToolRegistry: vi.fn(), getExtensions: vi.fn(() => []), getModel: vi.fn(() => 'gemini-pro'), getEmbeddingModel: vi.fn(() => 'embedding-001'),