Extensions MCP refactor (#12413)

2026-06-21 16:57:08 -07:00 · 2025-11-04 07:51:18 -08:00
parent 2b77c1ded4
commit da4fa5ad75
28 changed files with 877 additions and 478 deletions
@@ -21,27 +21,33 @@ vi.mock('../../config/extensions/storage.js', () => ({
  },
 }));
 vi.mock('../../config/extension-manager.js');
-vi.mock('@google/gemini-cli-core', () => ({
-  createTransport: vi.fn(),
-  MCPServerStatus: {
-    CONNECTED: 'CONNECTED',
-    CONNECTING: 'CONNECTING',
-    DISCONNECTED: 'DISCONNECTED',
-  },
-  Storage: vi.fn().mockImplementation((_cwd: string) => ({
-    getGlobalSettingsPath: () => '/tmp/gemini/settings.json',
-    getWorkspaceSettingsPath: () => '/tmp/gemini/workspace-settings.json',
-    getProjectTempDir: () => '/test/home/.gemini/tmp/mocked_hash',
-  })),
-  GEMINI_DIR: '.gemini',
-  getErrorMessage: (e: unknown) => (e instanceof Error ? e.message : String(e)),
-  debugLogger: {
-    log: vi.fn(),
-    warn: vi.fn(),
-    error: vi.fn(),
-    debug: vi.fn(),
-  },
-}));
+vi.mock('@google/gemini-cli-core', async (importOriginal) => {
+  const original =
+    await importOriginal<typeof import('@google/gemini-cli-core')>();
+  return {
+    ...original,
+    createTransport: vi.fn(),
+    MCPServerStatus: {
+      CONNECTED: 'CONNECTED',
+      CONNECTING: 'CONNECTING',
+      DISCONNECTED: 'DISCONNECTED',
+    },
+    Storage: vi.fn().mockImplementation((_cwd: string) => ({
+      getGlobalSettingsPath: () => '/tmp/gemini/settings.json',
+      getWorkspaceSettingsPath: () => '/tmp/gemini/workspace-settings.json',
+      getProjectTempDir: () => '/test/home/.gemini/tmp/mocked_hash',
+    })),
+    GEMINI_DIR: '.gemini',
+    getErrorMessage: (e: unknown) =>
+      e instanceof Error ? e.message : String(e),
+    debugLogger: {
+      log: vi.fn(),
+      warn: vi.fn(),
+      error: vi.fn(),
+      debug: vi.fn(),
+    },
+  };
+});
 vi.mock('@modelcontextprotocol/sdk/client/index.js');

 const mockedGetUserExtensionsDir =
@@ -1146,9 +1146,7 @@ describe('loadCliConfig with allowed-mcp-server-names', () => {
    ];
    const argv = await parseArguments({} as Settings);
    const config = await loadCliConfig(baseSettings, 'test-session', argv);
-    expect(config.getMcpServers()).toEqual({
-      server1: { url: 'http://localhost:8080' },
-    });
+    expect(config.getAllowedMcpServers()).toEqual(['server1']);
  });

  it('should allow multiple specified MCP servers', async () => {
@@ -1162,10 +1160,7 @@ describe('loadCliConfig with allowed-mcp-server-names', () => {
    ];
    const argv = await parseArguments({} as Settings);
    const config = await loadCliConfig(baseSettings, 'test-session', argv);
-    expect(config.getMcpServers()).toEqual({
-      server1: { url: 'http://localhost:8080' },
-      server3: { url: 'http://localhost:8082' },
-    });
+    expect(config.getAllowedMcpServers()).toEqual(['server1', 'server3']);
  });

  it('should handle server names that do not exist', async () => {
@@ -1179,16 +1174,14 @@ describe('loadCliConfig with allowed-mcp-server-names', () => {
    ];
    const argv = await parseArguments({} as Settings);
    const config = await loadCliConfig(baseSettings, 'test-session', argv);
-    expect(config.getMcpServers()).toEqual({
-      server1: { url: 'http://localhost:8080' },
-    });
+    expect(config.getAllowedMcpServers()).toEqual(['server1', 'server4']);
  });

  it('should allow no MCP servers if the flag is provided but empty', async () => {
    process.argv = ['node', 'script.js', '--allowed-mcp-server-names', ''];
    const argv = await parseArguments({} as Settings);
    const config = await loadCliConfig(baseSettings, 'test-session', argv);
-    expect(config.getMcpServers()).toEqual({});
+    expect(config.getAllowedMcpServers()).toEqual(['']);
  });

  it('should read allowMCPServers from settings', async () => {
@@ -1199,10 +1192,7 @@ describe('loadCliConfig with allowed-mcp-server-names', () => {
      mcp: { allowed: ['server1', 'server2'] },
    };
    const config = await loadCliConfig(settings, 'test-session', argv);
-    expect(config.getMcpServers()).toEqual({
-      server1: { url: 'http://localhost:8080' },
-      server2: { url: 'http://localhost:8081' },
-    });
+    expect(config.getAllowedMcpServers()).toEqual(['server1', 'server2']);
  });

  it('should read excludeMCPServers from settings', async () => {
@@ -1213,9 +1203,7 @@ describe('loadCliConfig with allowed-mcp-server-names', () => {
      mcp: { excluded: ['server1', 'server2'] },
    };
    const config = await loadCliConfig(settings, 'test-session', argv);
-    expect(config.getMcpServers()).toEqual({
-      server3: { url: 'http://localhost:8082' },
-    });
+    expect(config.getBlockedMcpServers()).toEqual(['server1', 'server2']);
  });

  it('should override allowMCPServers with excludeMCPServers if overlapping', async () => {
@@ -1229,9 +1217,8 @@ describe('loadCliConfig with allowed-mcp-server-names', () => {
      },
    };
    const config = await loadCliConfig(settings, 'test-session', argv);
-    expect(config.getMcpServers()).toEqual({
-      server2: { url: 'http://localhost:8081' },
-    });
+    expect(config.getAllowedMcpServers()).toEqual(['server1', 'server2']);
+    expect(config.getBlockedMcpServers()).toEqual(['server1']);
  });

  it('should prioritize mcp server flag if set', async () => {
@@ -1250,9 +1237,7 @@ describe('loadCliConfig with allowed-mcp-server-names', () => {
      },
    };
    const config = await loadCliConfig(settings, 'test-session', argv);
-    expect(config.getMcpServers()).toEqual({
-      server1: { url: 'http://localhost:8080' },
-    });
+    expect(config.getAllowedMcpServers()).toEqual(['server1']);
  });

  it('should prioritize CLI flag over both allowed and excluded settings', async () => {
@@ -1273,10 +1258,8 @@ describe('loadCliConfig with allowed-mcp-server-names', () => {
      },
    };
    const config = await loadCliConfig(settings, 'test-session', argv);
-    expect(config.getMcpServers()).toEqual({
-      server2: { url: 'http://localhost:8081' },
-      server3: { url: 'http://localhost:8082' },
-    });
+    expect(config.getAllowedMcpServers()).toEqual(['server2', 'server3']);
+    expect(config.getBlockedMcpServers()).toEqual([]);
  });
 });

@@ -13,9 +13,7 @@ import process from 'node:process';
 import { mcpCommand } from '../commands/mcp.js';
 import type {
  FileFilteringOptions,
-  MCPServerConfig,
  OutputFormat,
-  GeminiCLIExtension,
 } from '@google/gemini-cli-core';
 import { extensionsCommand } from '../commands/extensions.js';
 import {
@@ -49,9 +47,13 @@ import { appEvents } from '../utils/events.js';
 import { isWorkspaceTrusted } from './trustedFolders.js';
 import { createPolicyEngineConfig } from './policy.js';
 import { ExtensionManager } from './extension-manager.js';
-import type { ExtensionLoader } from '@google/gemini-cli-core/src/utils/extensionLoader.js';
+import type {
+  ExtensionEvents,
+  ExtensionLoader,
+} from '@google/gemini-cli-core/src/utils/extensionLoader.js';
 import { requestConsentNonInteractive } from './extensions/consent.js';
 import { promptForSetting } from './extensions/extensionSettings.js';
+import type { EventEmitter } from 'node:stream';

 export interface CliArgs {
  query: string | undefined;
@@ -429,6 +431,7 @@ export async function loadCliConfig(
    requestSetting: promptForSetting,
    workspaceDir: cwd,
    enabledExtensionOverrides: argv.extensions,
+    eventEmitter: appEvents as EventEmitter<ExtensionEvents>,
  });
  await extensionManager.loadExtensions();

@@ -448,7 +451,6 @@ export async function loadCliConfig(
      memoryFileFiltering,
    );

-  let mcpServers = mergeMcpServers(settings, extensionManager.getExtensions());
  const question = argv.promptInteractive || argv.prompt || '';

  // Determine approval mode with backward compatibility
@@ -565,37 +567,8 @@ export async function loadCliConfig(

  const excludeTools = mergeExcludeTools(
    settings,
-    extensionManager.getExtensions(),
    extraExcludes.length > 0 ? extraExcludes : undefined,
  );
-  const blockedMcpServers: Array<{ name: string; extensionName: string }> = [];
-
-  if (!argv.allowedMcpServerNames) {
-    if (settings.mcp?.allowed) {
-      mcpServers = allowedMcpServers(
-        mcpServers,
-        settings.mcp.allowed,
-        blockedMcpServers,
-      );
-    }
-
-    if (settings.mcp?.excluded) {
-      const excludedNames = new Set(settings.mcp.excluded.filter(Boolean));
-      if (excludedNames.size > 0) {
-        mcpServers = Object.fromEntries(
-          Object.entries(mcpServers).filter(([key]) => !excludedNames.has(key)),
-        );
-      }
-    }
-  }
-
-  if (argv.allowedMcpServerNames) {
-    mcpServers = allowedMcpServers(
-      mcpServers,
-      argv.allowedMcpServerNames,
-      blockedMcpServers,
-    );
-  }

  const useModelRouter = settings.experimental?.useModelRouter ?? true;
  const defaultModel = useModelRouter
@@ -633,7 +606,11 @@ export async function loadCliConfig(
    toolDiscoveryCommand: settings.tools?.discoveryCommand,
    toolCallCommand: settings.tools?.callCommand,
    mcpServerCommand: settings.mcp?.serverCommand,
-    mcpServers,
+    mcpServers: settings.mcpServers,
+    allowedMcpServers: argv.allowedMcpServerNames ?? settings.mcp?.allowed,
+    blockedMcpServers: argv.allowedMcpServerNames
+      ? [] // explicitly allowed servers overrides everything
+      : settings.mcp?.excluded,
    userMemory: memoryContent,
    geminiMdFileCount: fileCount,
    geminiMdFilePaths: filePaths,
@@ -663,7 +640,6 @@ export async function loadCliConfig(
    enabledExtensions: argv.extensions,
    extensionLoader: extensionManager,
    enableExtensionReloading: settings.experimental?.extensionReloading,
-    blockedMcpServers,
    noBrowser: !!process.env['NO_BROWSER'],
    summarizeToolOutput: settings.model?.summarizeToolOutput,
    ideMode,
@@ -699,75 +675,13 @@ export async function loadCliConfig(
  });
 }

-function allowedMcpServers(
-  mcpServers: { [x: string]: MCPServerConfig },
-  allowMCPServers: string[],
-  blockedMcpServers: Array<{ name: string; extensionName: string }>,
-) {
-  const allowedNames = new Set(allowMCPServers.filter(Boolean));
-  if (allowedNames.size > 0) {
-    mcpServers = Object.fromEntries(
-      Object.entries(mcpServers).filter(([key, server]) => {
-        const isAllowed = allowedNames.has(key);
-        if (!isAllowed) {
-          blockedMcpServers.push({
-            name: key,
-            extensionName: server.extension?.name || '',
-          });
-        }
-        return isAllowed;
-      }),
-    );
-  } else {
-    blockedMcpServers.push(
-      ...Object.entries(mcpServers).map(([key, server]) => ({
-        name: key,
-        extensionName: server.extension?.name || '',
-      })),
-    );
-    mcpServers = {};
-  }
-  return mcpServers;
-}
-
-function mergeMcpServers(settings: Settings, extensions: GeminiCLIExtension[]) {
-  const mcpServers = { ...(settings.mcpServers || {}) };
-  for (const extension of extensions) {
-    if (!extension.isActive) {
-      continue;
-    }
-    Object.entries(extension.mcpServers || {}).forEach(([key, server]) => {
-      if (mcpServers[key]) {
-        debugLogger.warn(
-          `Skipping extension MCP config for server with key "${key}" as it already exists.`,
-        );
-        return;
-      }
-      mcpServers[key] = {
-        ...server,
-        extension,
-      };
-    });
-  }
-  return mcpServers;
-}
-
 function mergeExcludeTools(
  settings: Settings,
-  extensions: GeminiCLIExtension[],
  extraExcludes?: string[] | undefined,
 ): string[] {
  const allExcludeTools = new Set([
    ...(settings.tools?.exclude || []),
    ...(extraExcludes || []),
  ]);
-  for (const extension of extensions) {
-    if (!extension.isActive) {
-      continue;
-    }
-    for (const tool of extension.excludeTools || []) {
-      allExcludeTools.add(tool);
-    }
-  }
  return [...allExcludeTools];
 }
@@ -28,6 +28,7 @@ import {
  ExtensionDisableEvent,
  ExtensionEnableEvent,
  ExtensionInstallEvent,
+  ExtensionLoader,
  ExtensionUninstallEvent,
  ExtensionUpdateEvent,
  getErrorMessage,
@@ -36,6 +37,7 @@ import {
  logExtensionInstallEvent,
  logExtensionUninstall,
  logExtensionUpdateEvent,
+  type ExtensionEvents,
  type MCPServerConfig,
  type ExtensionInstallMetadata,
  type GeminiCLIExtension,
@@ -54,11 +56,7 @@ import {
  maybePromptForSettings,
  type ExtensionSetting,
 } from './extensions/extensionSettings.js';
-import type {
-  ExtensionEvents,
-  ExtensionLoader,
-} from '@google/gemini-cli-core/src/utils/extensionLoader.js';
-import { EventEmitter } from 'node:events';
+import type { EventEmitter } from 'node:stream';

 interface ExtensionManagerParams {
  enabledExtensionOverrides?: string[];
@@ -66,6 +64,7 @@ interface ExtensionManagerParams {
  requestConsent: (consent: string) => Promise<boolean>;
  requestSetting: ((setting: ExtensionSetting) => Promise<string>) | null;
  workspaceDir: string;
+  eventEmitter?: EventEmitter<ExtensionEvents>;
 }

 /**
@@ -73,7 +72,7 @@ interface ExtensionManagerParams {
 *
 * You must call `loadExtensions` prior to calling other methods on this class.
 */
-export class ExtensionManager implements ExtensionLoader {
+export class ExtensionManager extends ExtensionLoader {
  private extensionEnablementManager: ExtensionEnablementManager;
  private settings: Settings;
  private requestConsent: (consent: string) => Promise<boolean>;
@@ -83,9 +82,9 @@ export class ExtensionManager implements ExtensionLoader {
  private telemetryConfig: Config;
  private workspaceDir: string;
  private loadedExtensions: GeminiCLIExtension[] | undefined;
-  private eventEmitter: EventEmitter<ExtensionEvents>;

  constructor(options: ExtensionManagerParams) {
+    super(options.eventEmitter);
    this.workspaceDir = options.workspaceDir;
    this.extensionEnablementManager = new ExtensionEnablementManager(
      options.enabledExtensionOverrides,
@@ -102,7 +101,6 @@ export class ExtensionManager implements ExtensionLoader {
    });
    this.requestConsent = options.requestConsent;
    this.requestSetting = options.requestSetting ?? undefined;
-    this.eventEmitter = new EventEmitter();
  }

  setRequestConsent(
@@ -126,10 +124,6 @@ export class ExtensionManager implements ExtensionLoader {
    return this.loadedExtensions!;
  }

-  extensionEvents(): EventEmitter<ExtensionEvents> {
-    return this.eventEmitter;
-  }
-
  async installOrUpdateExtension(
    installMetadata: ExtensionInstallMetadata,
    previousExtensionConfig?: ExtensionConfig,
@@ -303,7 +297,7 @@ export class ExtensionManager implements ExtensionLoader {
        await fs.promises.writeFile(metadataPath, metadataString);

        // TODO: Gracefully handle this call failing, we should back up the old
-        // extension prior to overwriting it and then restore it.
+        // extension prior to overwriting it and then restore and restart it.
        extension = await this.loadExtension(destinationPath)!;
        if (!extension) {
          throw new Error(`Extension not found`);
@@ -320,7 +314,6 @@ export class ExtensionManager implements ExtensionLoader {
              'success',
            ),
          );
-          this.eventEmitter.emit('extensionUpdated', { extension });
        } else {
          logExtensionInstallEvent(
            this.telemetryConfig,
@@ -332,7 +325,6 @@ export class ExtensionManager implements ExtensionLoader {
              'success',
            ),
          );
-          this.eventEmitter.emit('extensionInstalled', { extension });
          this.enableExtension(newExtensionConfig.name, SettingScope.User);
        }
      } finally {
@@ -397,7 +389,7 @@ export class ExtensionManager implements ExtensionLoader {
    if (!extension) {
      throw new Error(`Extension not found.`);
    }
-    this.unloadExtension(extension);
+    await this.unloadExtension(extension);
    const storage = new ExtensionStorage(extension.name);

    await fs.promises.rm(storage.getExtensionDir(), {
@@ -419,9 +411,11 @@ export class ExtensionManager implements ExtensionLoader {
        'success',
      ),
    );
-    this.eventEmitter.emit('extensionUninstalled', { extension });
  }

+  /**
+   * Loads all installed extensions, should only be called once.
+   */
  async loadExtensions(): Promise<GeminiCLIExtension[]> {
    if (this.loadedExtensions) {
      throw new Error('Extensions already loaded, only load extensions once.');
@@ -433,12 +427,14 @@ export class ExtensionManager implements ExtensionLoader {
    }
    for (const subdir of fs.readdirSync(extensionsDir)) {
      const extensionDir = path.join(extensionsDir, subdir);
-
      await this.loadExtension(extensionDir);
    }
    return this.loadedExtensions;
  }

+  /**
+   * Adds `extension` to the list of extensions and starts it if appropriate.
+   */
  private async loadExtension(
    extensionDir: string,
  ): Promise<GeminiCLIExtension | null> {
@@ -499,8 +495,9 @@ export class ExtensionManager implements ExtensionLoader {
        ),
        id: getExtensionId(config, installMetadata),
      };
-      this.eventEmitter.emit('extensionLoaded', { extension });
-      this.getExtensions().push(extension);
+      this.loadedExtensions = [...this.loadedExtensions, extension];
+
+      await this.maybeStartExtension(extension);
      return extension;
    } catch (e) {
      debugLogger.error(
@@ -512,11 +509,17 @@ export class ExtensionManager implements ExtensionLoader {
    }
  }

-  private unloadExtension(extension: GeminiCLIExtension) {
+  /**
+   * Removes `extension` from the list of extensions and stops it if
+   * appropriate.
+   */
+  private unloadExtension(
+    extension: GeminiCLIExtension,
+  ): Promise<void> | undefined {
    this.loadedExtensions = this.getExtensions().filter(
      (entry) => extension !== entry,
    );
-    this.eventEmitter.emit('extensionUnloaded', { extension });
+    return this.maybeStopExtension(extension);
  }

  loadExtensionConfig(extensionDir: string): ExtensionConfig {
@@ -616,14 +619,18 @@ export class ExtensionManager implements ExtensionLoader {
    const scopePath =
      scope === SettingScope.Workspace ? this.workspaceDir : os.homedir();
    this.extensionEnablementManager.disable(name, true, scopePath);
+    extension.isActive = false;
+    await this.maybeStopExtension(extension);
    logExtensionDisable(
      this.telemetryConfig,
      new ExtensionDisableEvent(hashValue(name), extension.id, scope),
    );
-    extension.isActive = false;
-    this.eventEmitter.emit('extensionDisabled', { extension });
  }

+  /**
+   * Enables an existing extension for a given scope, and starts it if
+   * appropriate.
+   */
  async enableExtension(name: string, scope: SettingScope) {
    if (
      scope === SettingScope.System ||
@@ -645,7 +652,7 @@ export class ExtensionManager implements ExtensionLoader {
      new ExtensionEnableEvent(hashValue(name), extension.id, scope),
    );
    extension.isActive = true;
-    this.eventEmitter.emit('extensionEnabled', { extension });
+    await this.maybeStartExtension(extension);
  }
 }

@@ -1666,7 +1666,10 @@ This extension will run the following MCP servers:
      });

      await extensionManager.loadExtensions();
-      extensionManager.disableExtension('my-extension', SettingScope.User);
+      await extensionManager.disableExtension(
+        'my-extension',
+        SettingScope.User,
+      );
      expect(
        isEnabled({
          name: 'my-extension',
@@ -1683,7 +1686,10 @@ This extension will run the following MCP servers:
      });

      await extensionManager.loadExtensions();
-      extensionManager.disableExtension('my-extension', SettingScope.Workspace);
+      await extensionManager.disableExtension(
+        'my-extension',
+        SettingScope.Workspace,
+      );
      expect(
        isEnabled({
          name: 'my-extension',
@@ -1706,8 +1712,14 @@ This extension will run the following MCP servers:
      });

      await extensionManager.loadExtensions();
-      extensionManager.disableExtension('my-extension', SettingScope.User);
-      extensionManager.disableExtension('my-extension', SettingScope.User);
+      await extensionManager.disableExtension(
+        'my-extension',
+        SettingScope.User,
+      );
+      await extensionManager.disableExtension(
+        'my-extension',
+        SettingScope.User,
+      );
      expect(
        isEnabled({
          name: 'my-extension',
@@ -1738,7 +1750,7 @@ This extension will run the following MCP servers:
      });

      await extensionManager.loadExtensions();
-      extensionManager.disableExtension('ext1', SettingScope.Workspace);
+      await extensionManager.disableExtension('ext1', SettingScope.Workspace);

      expect(mockLogExtensionDisable).toHaveBeenCalled();
      expect(ExtensionDisableEvent).toHaveBeenCalledWith(
@@ -1766,7 +1778,7 @@ This extension will run the following MCP servers:
        version: '1.0.0',
      });
      await extensionManager.loadExtensions();
-      extensionManager.disableExtension('ext1', SettingScope.User);
+      await extensionManager.disableExtension('ext1', SettingScope.User);
      let activeExtensions = getActiveExtensions();
      expect(activeExtensions).toHaveLength(0);

@@ -1783,7 +1795,7 @@ This extension will run the following MCP servers:
        version: '1.0.0',
      });
      await extensionManager.loadExtensions();
-      extensionManager.disableExtension('ext1', SettingScope.Workspace);
+      await extensionManager.disableExtension('ext1', SettingScope.Workspace);
      let activeExtensions = getActiveExtensions();
      expect(activeExtensions).toHaveLength(0);

@@ -1804,8 +1816,8 @@ This extension will run the following MCP servers:
        },
      });
      await extensionManager.loadExtensions();
-      extensionManager.disableExtension('ext1', SettingScope.Workspace);
-      extensionManager.enableExtension('ext1', SettingScope.Workspace);
+      await extensionManager.disableExtension('ext1', SettingScope.Workspace);
+      await extensionManager.enableExtension('ext1', SettingScope.Workspace);

      expect(mockLogExtensionEnable).toHaveBeenCalled();
      expect(ExtensionEnableEvent).toHaveBeenCalledWith(
@@ -190,6 +190,7 @@ describe('gemini.tsx main function', () => {
        getDebugMode: () => false,
        getListExtensions: () => false,
        getMcpServers: () => ({}),
+        getMcpClientManager: vi.fn(),
        initialize: vi.fn(),
        getIdeMode: () => false,
        getExperimentalZedIntegration: () => false,
@@ -339,6 +340,7 @@ describe('gemini.tsx main function kitty protocol', () => {
      getDebugMode: () => false,
      getListExtensions: () => false,
      getMcpServers: () => ({}),
+      getMcpClientManager: vi.fn(),
      initialize: vi.fn(),
      getIdeMode: () => false,
      getExperimentalZedIntegration: () => false,
@@ -159,8 +159,10 @@ describe('McpPromptLoader', () => {

  describe('loadCommands', () => {
    const mockConfigWithPrompts = {
-      getMcpServers: () => ({
-        'test-server': { httpUrl: 'https://test-server.com' },
+      getMcpClientManager: () => ({
+        getMcpServers: () => ({
+          'test-server': { httpUrl: 'https://test-server.com' },
+        }),
      }),
    } as unknown as Config;

@@ -34,7 +34,7 @@ export class McpPromptLoader implements ICommandLoader {
    if (!this.config) {
      return Promise.resolve([]);
    }
-    const mcpServers = this.config.getMcpServers() || {};
+    const mcpServers = this.config.getMcpClientManager()?.getMcpServers() || {};
    for (const serverName in mcpServers) {
      const prompts = getMCPServerPrompts(this.config, serverName) || [];
      for (const prompt of prompts) {
@@ -101,7 +101,8 @@ export class McpPromptLoader implements ICommandLoader {
            }

            try {
-              const mcpServers = this.config.getMcpServers() || {};
+              const mcpServers =
+                this.config.getMcpClientManager()?.getMcpServers() || {};
              const mcpServerConfig = mcpServers[serverName];
              if (!mcpServerConfig) {
                return {
@@ -295,6 +295,7 @@ describe('AppContainer State Management', () => {
      getExtensions: vi.fn().mockReturnValue([]),
      setRequestConsent: vi.fn(),
      setRequestSetting: vi.fn(),
+      start: vi.fn(),
    } as unknown as ExtensionManager);
    vi.spyOn(mockConfig, 'getExtensionLoader').mockReturnValue(
      mockExtensionManager,
@@ -62,6 +62,7 @@ describe('mcpCommand', () => {
    getBlockedMcpServers: ReturnType<typeof vi.fn>;
    getPromptRegistry: ReturnType<typeof vi.fn>;
    getGeminiClient: ReturnType<typeof vi.fn>;
+    getMcpClientManager: ReturnType<typeof vi.fn>;
  };

  beforeEach(() => {
@@ -88,6 +89,10 @@ describe('mcpCommand', () => {
        getPromptsByServer: vi.fn().mockReturnValue([]),
      }),
      getGeminiClient: vi.fn(),
+      getMcpClientManager: vi.fn().mockImplementation(() => ({
+        getBlockedMcpServers: vi.fn(),
+        getMcpServers: vi.fn(),
+      })),
    };

    mockContext = createMockCommandContext({
@@ -43,7 +43,7 @@ const authCommand: SlashCommand = {
      };
    }

-    const mcpServers = config.getMcpServers() || {};
+    const mcpServers = config.getMcpClientManager()?.getMcpServers() ?? {};

    if (!serverName) {
      // List servers that support OAuth
@@ -119,20 +119,20 @@ const authCommand: SlashCommand = {
      );

      // Trigger tool re-discovery to pick up authenticated server
-      const toolRegistry = config.getToolRegistry();
-      if (toolRegistry) {
+      const mcpClientManager = config.getMcpClientManager();
+      if (mcpClientManager) {
        context.ui.addItem(
          {
            type: 'info',
-            text: `Re-discovering tools from '${serverName}'...`,
+            text: `Restarting MCP server '${serverName}'...`,
          },
          Date.now(),
        );
-        await toolRegistry.discoverToolsForServer(serverName);
+        await mcpClientManager.restartServer(serverName);
      }
      // Update the client with the new tools
      const geminiClient = config.getGeminiClient();
-      if (geminiClient) {
+      if (geminiClient?.isInitialized()) {
        await geminiClient.setTools();
      }

@@ -158,7 +158,7 @@ const authCommand: SlashCommand = {
    const { config } = context.services;
    if (!config) return [];

-    const mcpServers = config.getMcpServers() || {};
+    const mcpServers = config.getMcpClientManager()?.getMcpServers() || {};
    return Object.keys(mcpServers).filter((name) =>
      name.startsWith(partialArg),
    );
@@ -188,9 +188,10 @@ const listAction = async (
    };
  }

-  const mcpServers = config.getMcpServers() || {};
+  const mcpServers = config.getMcpClientManager()?.getMcpServers() || {};
  const serverNames = Object.keys(mcpServers);
-  const blockedMcpServers = config.getBlockedMcpServers() || [];
+  const blockedMcpServers =
+    config.getMcpClientManager()?.getBlockedMcpServers() || [];

  const connectingServers = serverNames.filter(
    (name) => getMCPServerStatus(name) === MCPServerStatus.CONNECTING,
@@ -299,12 +300,12 @@ const refreshCommand: SlashCommand = {
      };
    }

-    const toolRegistry = config.getToolRegistry();
-    if (!toolRegistry) {
+    const mcpClientManager = config.getMcpClientManager();
+    if (!mcpClientManager) {
      return {
        type: 'message',
        messageType: 'error',
-        content: 'Could not retrieve tool registry.',
+        content: 'Could not retrieve mcp client manager.',
      };
    }

@@ -316,11 +317,11 @@ const refreshCommand: SlashCommand = {
      Date.now(),
    );

-    await toolRegistry.restartMcpServers();
+    await mcpClientManager.restart();

    // Update the client with the new tools
    const geminiClient = config.getGeminiClient();
-    if (geminiClient) {
+    if (geminiClient?.isInitialized()) {
      await geminiClient.setTools();
    }

@@ -144,7 +144,10 @@ const createMockConfig = (overrides = {}) => ({
  getDebugMode: vi.fn(() => false),
  getAccessibility: vi.fn(() => ({})),
  getMcpServers: vi.fn(() => ({})),
-  getBlockedMcpServers: vi.fn(() => []),
+  getMcpClientManager: vi.fn().mockImplementation(() => ({
+    getBlockedMcpServers: vi.fn(),
+    getMcpServers: vi.fn(),
+  })),
  ...overrides,
 });

@@ -101,8 +101,10 @@ export const Composer = () => {
                ideContext={uiState.ideContextState}
                geminiMdFileCount={uiState.geminiMdFileCount}
                contextFileNames={contextFileNames}
-                mcpServers={config.getMcpServers()}
-                blockedMcpServers={config.getBlockedMcpServers()}
+                mcpServers={config.getMcpClientManager()?.getMcpServers() ?? {}}
+                blockedMcpServers={
+                  config.getMcpClientManager()?.getBlockedMcpServers() ?? []
+                }
              />
            )
          )}
@@ -5,15 +5,13 @@
 */

 import { useEffect, useState } from 'react';
-import { appEvents } from './../../utils/events.js';
+import { AppEvent, appEvents } from './../../utils/events.js';
 import { Box, Text } from 'ink';
-import { useConfig } from '../contexts/ConfigContext.js';
 import { type McpClient, MCPServerStatus } from '@google/gemini-cli-core';
 import { GeminiSpinner } from './GeminiRespondingSpinner.js';
 import { theme } from '../semantic-colors.js';

 export const ConfigInitDisplay = () => {
-  const config = useConfig();
  const [message, setMessage] = useState('Initializing...');

  useEffect(() => {
@@ -31,11 +29,11 @@ export const ConfigInitDisplay = () => {
      setMessage(`Connecting to MCP servers... (${connected}/${clients.size})`);
    };

-    appEvents.on('mcp-client-update', onChange);
+    appEvents.on(AppEvent.McpClientUpdate, onChange);
    return () => {
-      appEvents.off('mcp-client-update', onChange);
+      appEvents.off(AppEvent.McpClientUpdate, onChange);
    };
-  }, [config]);
+  }, []);

  return (
    <Box marginTop={1}>
@@ -4,6 +4,7 @@
 * SPDX-License-Identifier: Apache-2.0
 */

+import type { ExtensionEvents, McpClient } from '@google/gemini-cli-core';
 import { EventEmitter } from 'node:events';

 export enum AppEvent {
@@ -11,6 +12,15 @@ export enum AppEvent {
  LogError = 'log-error',
  OauthDisplayMessage = 'oauth-display-message',
  Flicker = 'flicker',
+  McpClientUpdate = 'mcp-client-update',
 }

-export const appEvents = new EventEmitter();
+export interface AppEvents extends ExtensionEvents {
+  [AppEvent.OpenDebugConsole]: never[];
+  [AppEvent.LogError]: string[];
+  [AppEvent.OauthDisplayMessage]: string[];
+  [AppEvent.Flicker]: never[];
+  [AppEvent.McpClientUpdate]: Array<Map<string, McpClient> | never>;
+}
+
+export const appEvents = new EventEmitter<AppEvents>();