fix(cli): Prevent stdout/stderr patching for extension commands (#13600)

chrstnb · jacob314 · web-flow · commit bdf80ea7c046 · 2025-11-22T02:08:06.000Z
Co-authored-by: jacob314 &lt;jacob314@gmail.com&gt;
diff --git a/packages/cli/src/commands/extensions.tsx b/packages/cli/src/commands/extensions.tsx
@@ -14,13 +14,15 @@ import { enableCommand } from './extensions/enable.js';
 import { linkCommand } from './extensions/link.js';
 import { newCommand } from './extensions/new.js';
 import { validateCommand } from './extensions/validate.js';
+import { initializeOutputListenersAndFlush } from '../gemini.js';
 
 export const extensionsCommand: CommandModule = {
   command: 'extensions <command>',
   aliases: ['extension'],
   describe: 'Manage Gemini CLI extensions.',
   builder: (yargs) =>
     yargs
+      .middleware(() => initializeOutputListenersAndFlush())
       .command(installCommand)
       .command(uninstallCommand)
       .command(listCommand)
diff --git a/packages/cli/src/commands/extensions/disable.test.ts b/packages/cli/src/commands/extensions/disable.test.ts
@@ -56,6 +56,9 @@ vi.mock('../../config/extensions/consent.js', () => ({
 vi.mock('../../config/extensions/extensionSettings.js', () => ({
   promptForSetting: vi.fn(),
 }));
+vi.mock('../utils.js', () => ({
+  exitCli: vi.fn(),
+}));
 
 describe('extensions disable command', () => {
   const mockLoadSettings = vi.mocked(loadSettings);
diff --git a/packages/cli/src/commands/extensions/disable.ts b/packages/cli/src/commands/extensions/disable.ts
@@ -11,6 +11,7 @@ import { debugLogger } from '@google/gemini-cli-core';
 import { ExtensionManager } from '../../config/extension-manager.js';
 import { requestConsentNonInteractive } from '../../config/extensions/consent.js';
 import { promptForSetting } from '../../config/extensions/extensionSettings.js';
+import { exitCli } from '../utils.js';
 
 interface DisableArgs {
   name: string;
@@ -81,5 +82,6 @@ export const disableCommand: CommandModule = {
       name: argv['name'] as string,
       scope: argv['scope'] as string,
     });
+    await exitCli();
   },
 };
diff --git a/packages/cli/src/commands/extensions/enable.test.ts b/packages/cli/src/commands/extensions/enable.test.ts
@@ -58,6 +58,9 @@ vi.mock('../../config/extension-manager.js');
 vi.mock('../../config/settings.js');
 vi.mock('../../config/extensions/consent.js');
 vi.mock('../../config/extensions/extensionSettings.js');
+vi.mock('../utils.js', () => ({
+  exitCli: vi.fn(),
+}));
 
 describe('extensions enable command', () => {
   const mockLoadSettings = vi.mocked(loadSettings);
diff --git a/packages/cli/src/commands/extensions/enable.ts b/packages/cli/src/commands/extensions/enable.ts
@@ -14,6 +14,7 @@ import {
   getErrorMessage,
 } from '@google/gemini-cli-core';
 import { promptForSetting } from '../../config/extensions/extensionSettings.js';
+import { exitCli } from '../utils.js';
 
 interface EnableArgs {
   name: string;
@@ -86,5 +87,6 @@ export const enableCommand: CommandModule = {
       name: argv['name'] as string,
       scope: argv['scope'] as string,
     });
+    await exitCli();
   },
 };
diff --git a/packages/cli/src/commands/extensions/install.test.ts b/packages/cli/src/commands/extensions/install.test.ts
@@ -48,6 +48,10 @@ vi.mock('node:fs/promises', () => ({
   },
 }));
 
+vi.mock('../utils.js', () => ({
+  exitCli: vi.fn(),
+}));
+
 describe('extensions install command', () => {
   it('should fail if no source is provided', () => {
     const validationParser = yargs([]).command(installCommand).fail(false);
diff --git a/packages/cli/src/commands/extensions/install.ts b/packages/cli/src/commands/extensions/install.ts
@@ -18,6 +18,7 @@ import {
 import { ExtensionManager } from '../../config/extension-manager.js';
 import { loadSettings } from '../../config/settings.js';
 import { promptForSetting } from '../../config/extensions/extensionSettings.js';
+import { exitCli } from '../utils.js';
 
 interface InstallArgs {
   source: string;
@@ -130,5 +131,6 @@ export const installCommand: CommandModule = {
       allowPreRelease: argv['pre-release'] as boolean | undefined,
       consent: argv['consent'] as boolean | undefined,
     });
+    await exitCli();
   },
 };
diff --git a/packages/cli/src/commands/extensions/link.test.ts b/packages/cli/src/commands/extensions/link.test.ts
@@ -52,6 +52,9 @@ vi.mock('../../config/extensions/consent.js', () => ({
 vi.mock('../../config/extensions/extensionSettings.js', () => ({
   promptForSetting: vi.fn(),
 }));
+vi.mock('../utils.js', () => ({
+  exitCli: vi.fn(),
+}));
 
 describe('extensions link command', () => {
   const mockLoadSettings = vi.mocked(loadSettings);
diff --git a/packages/cli/src/commands/extensions/link.ts b/packages/cli/src/commands/extensions/link.ts
@@ -15,6 +15,7 @@ import { requestConsentNonInteractive } from '../../config/extensions/consent.js
 import { ExtensionManager } from '../../config/extension-manager.js';
 import { loadSettings } from '../../config/settings.js';
 import { promptForSetting } from '../../config/extensions/extensionSettings.js';
+import { exitCli } from '../utils.js';
 
 interface InstallArgs {
   path: string;
@@ -60,5 +61,6 @@ export const linkCommand: CommandModule = {
     await handleLink({
       path: argv['path'] as string,
     });
+    await exitCli();
   },
 };
diff --git a/packages/cli/src/commands/extensions/list.test.ts b/packages/cli/src/commands/extensions/list.test.ts
@@ -44,6 +44,9 @@ vi.mock('../../config/extensions/consent.js', () => ({
 vi.mock('../../config/extensions/extensionSettings.js', () => ({
   promptForSetting: vi.fn(),
 }));
+vi.mock('../utils.js', () => ({
+  exitCli: vi.fn(),
+}));
 
 describe('extensions list command', () => {
   const mockLoadSettings = vi.mocked(loadSettings);
diff --git a/packages/cli/src/commands/extensions/list.ts b/packages/cli/src/commands/extensions/list.ts
@@ -11,6 +11,7 @@ import { ExtensionManager } from '../../config/extension-manager.js';
 import { requestConsentNonInteractive } from '../../config/extensions/consent.js';
 import { loadSettings } from '../../config/settings.js';
 import { promptForSetting } from '../../config/extensions/extensionSettings.js';
+import { exitCli } from '../utils.js';
 
 export async function handleList() {
   try {
@@ -45,5 +46,6 @@ export const listCommand: CommandModule = {
   builder: (yargs) => yargs,
   handler: async () => {
     await handleList();
+    await exitCli();
   },
 };
diff --git a/packages/cli/src/commands/extensions/new.test.ts b/packages/cli/src/commands/extensions/new.test.ts
@@ -11,6 +11,9 @@ import * as fsPromises from 'node:fs/promises';
 import path from 'node:path';
 
 vi.mock('node:fs/promises');
+vi.mock('../utils.js', () => ({
+  exitCli: vi.fn(),
+}));
 
 const mockedFs = vi.mocked(fsPromises);
 
diff --git a/packages/cli/src/commands/extensions/new.ts b/packages/cli/src/commands/extensions/new.ts
@@ -9,6 +9,7 @@ import { join, dirname, basename } from 'node:path';
 import type { CommandModule } from 'yargs';
 import { fileURLToPath } from 'node:url';
 import { debugLogger } from '@google/gemini-cli-core';
+import { exitCli } from '../utils.js';
 
 interface NewArgs {
   path: string;
@@ -100,5 +101,6 @@ export const newCommand: CommandModule = {
       path: args['path'] as string,
       template: args['template'] as string | undefined,
     });
+    await exitCli();
   },
 };
diff --git a/packages/cli/src/commands/extensions/uninstall.test.ts b/packages/cli/src/commands/extensions/uninstall.test.ts
@@ -75,6 +75,9 @@ vi.mock('../../config/extensions/consent.js', () => ({
 vi.mock('../../config/extensions/extensionSettings.js', () => ({
   promptForSetting: vi.fn(),
 }));
+vi.mock('../utils.js', () => ({
+  exitCli: vi.fn(),
+}));
 
 describe('extensions uninstall command', () => {
   const mockLoadSettings = vi.mocked(loadSettings);
diff --git a/packages/cli/src/commands/extensions/uninstall.ts b/packages/cli/src/commands/extensions/uninstall.ts
@@ -11,6 +11,7 @@ import { requestConsentNonInteractive } from '../../config/extensions/consent.js
 import { ExtensionManager } from '../../config/extension-manager.js';
 import { loadSettings } from '../../config/settings.js';
 import { promptForSetting } from '../../config/extensions/extensionSettings.js';
+import { exitCli } from '../utils.js';
 
 interface UninstallArgs {
   names: string[]; // can be extension names or source URLs.
@@ -72,5 +73,6 @@ export const uninstallCommand: CommandModule = {
     await handleUninstall({
       names: argv['names'] as string[],
     });
+    await exitCli();
   },
 };
diff --git a/packages/cli/src/commands/extensions/update.test.ts b/packages/cli/src/commands/extensions/update.test.ts
@@ -56,6 +56,9 @@ vi.mock('../../config/extensions/consent.js', () => ({
 vi.mock('../../config/extensions/extensionSettings.js', () => ({
   promptForSetting: vi.fn(),
 }));
+vi.mock('../utils.js', () => ({
+  exitCli: vi.fn(),
+}));
 
 describe('extensions update command', () => {
   const mockLoadSettings = vi.mocked(loadSettings);
diff --git a/packages/cli/src/commands/extensions/update.ts b/packages/cli/src/commands/extensions/update.ts
@@ -19,6 +19,7 @@ import { ExtensionManager } from '../../config/extension-manager.js';
 import { requestConsentNonInteractive } from '../../config/extensions/consent.js';
 import { loadSettings } from '../../config/settings.js';
 import { promptForSetting } from '../../config/extensions/extensionSettings.js';
+import { exitCli } from '../utils.js';
 
 interface UpdateArgs {
   name?: string;
@@ -144,5 +145,6 @@ export const updateCommand: CommandModule = {
       name: argv['name'] as string | undefined,
       all: argv['all'] as boolean | undefined,
     });
+    await exitCli();
   },
 };
diff --git a/packages/cli/src/commands/extensions/validate.test.ts b/packages/cli/src/commands/extensions/validate.test.ts
@@ -13,6 +13,10 @@ import path from 'node:path';
 import * as os from 'node:os';
 import { debugLogger } from '@google/gemini-cli-core';
 
+vi.mock('../utils.js', () => ({
+  exitCli: vi.fn(),
+}));
+
 describe('extensions validate command', () => {
   it('should fail if no path is provided', () => {
     const validationParser = yargs([]).command(validateCommand).fail(false);
diff --git a/packages/cli/src/commands/extensions/validate.ts b/packages/cli/src/commands/extensions/validate.ts
@@ -15,6 +15,7 @@ import { ExtensionManager } from '../../config/extension-manager.js';
 import { requestConsentNonInteractive } from '../../config/extensions/consent.js';
 import { promptForSetting } from '../../config/extensions/extensionSettings.js';
 import { loadSettings } from '../../config/settings.js';
+import { exitCli } from '../utils.js';
 
 interface ValidateArgs {
   path: string;
@@ -101,5 +102,6 @@ export const validateCommand: CommandModule = {
     await handleValidate({
       path: args['path'] as string,
     });
+    await exitCli();
   },
 };
diff --git a/packages/cli/src/commands/mcp.test.ts b/packages/cli/src/commands/mcp.test.ts
@@ -56,6 +56,7 @@ describe('mcp command', () => {
       command: vi.fn().mockReturnThis(),
       demandCommand: vi.fn().mockReturnThis(),
       version: vi.fn().mockReturnThis(),
+      middleware: vi.fn().mockReturnThis(),
     };
 
     (mcpCommand.builder as (y: Argv) => Argv)(mockYargs as unknown as Argv);
diff --git a/packages/cli/src/commands/mcp.ts b/packages/cli/src/commands/mcp.ts
@@ -9,12 +9,14 @@ import type { CommandModule, Argv } from 'yargs';
 import { addCommand } from './mcp/add.js';
 import { removeCommand } from './mcp/remove.js';
 import { listCommand } from './mcp/list.js';
+import { initializeOutputListenersAndFlush } from '../gemini.js';
 
 export const mcpCommand: CommandModule = {
   command: 'mcp',
   describe: 'Manage MCP servers',
   builder: (yargs: Argv) =>
     yargs
+      .middleware(() => initializeOutputListenersAndFlush())
       .command(addCommand)
       .command(removeCommand)
       .command(listCommand)
diff --git a/packages/cli/src/commands/mcp/add.test.ts b/packages/cli/src/commands/mcp/add.test.ts
@@ -10,6 +10,10 @@ import { addCommand } from './add.js';
 import { loadSettings, SettingScope } from '../../config/settings.js';
 import { debugLogger } from '@google/gemini-cli-core';
 
+vi.mock('../utils.js', () => ({
+  exitCli: vi.fn(),
+}));
+
 vi.mock('fs/promises', () => ({
   readFile: vi.fn(),
   writeFile: vi.fn(),
diff --git a/packages/cli/src/commands/mcp/add.ts b/packages/cli/src/commands/mcp/add.ts
@@ -8,6 +8,7 @@
 import type { CommandModule } from 'yargs';
 import { loadSettings, SettingScope } from '../../config/settings.js';
 import { debugLogger, type MCPServerConfig } from '@google/gemini-cli-core';
+import { exitCli } from '../utils.js';
 
 async function addMcpServer(
   name: string,
@@ -230,5 +231,6 @@ export const addCommand: CommandModule = {
         excludeTools: argv['excludeTools'] as string[] | undefined,
       },
     );
+    await exitCli();
   },
 };
diff --git a/packages/cli/src/commands/mcp/list.test.ts b/packages/cli/src/commands/mcp/list.test.ts
@@ -44,6 +44,10 @@ vi.mock('@google/gemini-cli-core', async (importOriginal) => {
 });
 vi.mock('@modelcontextprotocol/sdk/client/index.js');
 
+vi.mock('../utils.js', () => ({
+  exitCli: vi.fn(),
+}));
+
 const mockedGetUserExtensionsDir =
   ExtensionStorage.getUserExtensionsDir as Mock;
 const mockedLoadSettings = loadSettings as Mock;
diff --git a/packages/cli/src/commands/mcp/list.ts b/packages/cli/src/commands/mcp/list.ts
@@ -17,6 +17,7 @@ import { Client } from '@modelcontextprotocol/sdk/client/index.js';
 import { ExtensionManager } from '../../config/extension-manager.js';
 import { requestConsentNonInteractive } from '../../config/extensions/consent.js';
 import { promptForSetting } from '../../config/extensions/extensionSettings.js';
+import { exitCli } from '../utils.js';
 
 const COLOR_GREEN = '\u001b[32m';
 const COLOR_YELLOW = '\u001b[33m';
@@ -145,5 +146,6 @@ export const listCommand: CommandModule = {
   describe: 'List all configured MCP servers',
   handler: async () => {
     await listMcpServers();
+    await exitCli();
   },
 };
diff --git a/packages/cli/src/commands/mcp/remove.test.ts b/packages/cli/src/commands/mcp/remove.test.ts
@@ -26,6 +26,10 @@ vi.mock('fs/promises', () => ({
   writeFile: vi.fn(),
 }));
 
+vi.mock('../utils.js', () => ({
+  exitCli: vi.fn(),
+}));
+
 describe('mcp remove command', () => {
   describe('unit tests with mocks', () => {
     let parser: Argv;
diff --git a/packages/cli/src/commands/mcp/remove.ts b/packages/cli/src/commands/mcp/remove.ts
@@ -8,6 +8,7 @@
 import type { CommandModule } from 'yargs';
 import { loadSettings, SettingScope } from '../../config/settings.js';
 import { debugLogger } from '@google/gemini-cli-core';
+import { exitCli } from '../utils.js';
 
 async function removeMcpServer(
   name: string,
@@ -57,5 +58,6 @@ export const removeCommand: CommandModule = {
     await removeMcpServer(argv['name'] as string, {
       scope: argv['scope'] as string,
     });
+    await exitCli();
   },
 };
diff --git a/packages/cli/src/commands/utils.ts b/packages/cli/src/commands/utils.ts
@@ -0,0 +1,12 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { runExitCleanup } from '../utils/cleanup.js';
+
+export async function exitCli(exitCode = 0) {
+  await runExitCleanup();
+  process.exit(exitCode);
+}
diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts
@@ -312,19 +312,6 @@ export async function parseArguments(settings: Settings): Promise<CliArgs> {
     process.exit(0);
   }
 
-  // If yargs handled --help/--version it will have exited; nothing to do here.
-
-  // Handle case where MCP subcommands are executed - they should exit the process
-  // and not return to main CLI logic
-  if (
-    result._.length > 0 &&
-    (result._[0] === 'mcp' || result._[0] === 'extensions')
-  ) {
-    // MCP commands handle their own execution and process exit
-    await runExitCleanup();
-    process.exit(0);
-  }
-
   // Normalize query args: handle both quoted "@path file" and unquoted @path file
   const queryArg = (result as { query?: string | string[] | undefined }).query;
   const q: string | undefined = Array.isArray(queryArg)
diff --git a/packages/cli/src/gemini.tsx b/packages/cli/src/gemini.tsx
@@ -634,7 +634,7 @@ function setWindowTitle(title: string, settings: LoadedSettings) {
   }
 }
 
-function initializeOutputListenersAndFlush() {
+export function initializeOutputListenersAndFlush() {
   // If there are no listeners for output, make sure we flush so output is not
   // lost.
   if (coreEvents.listenerCount(CoreEvent.Output) === 0) {
