fix(cli): prevent ACP stdout pollution from SessionEnd hooks (#26125)

Author: cocosheng-g

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);

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> = {}): 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'),