[ide-mode] Create an IDE manager class to handle connecting to and exposing methods from the IDE server (#4797)

4 files changed, 55 insertions, 193 deletions

diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts
index c0e9c215..55780320 100644
--- a/packages/cli/src/config/config.test.ts
+++ b/packages/cli/src/config/config.test.ts
@@ -1011,100 +1011,4 @@ describe('loadCliConfig ideMode', () => {
     const config = await loadCliConfig(settings, [], 'test-session', argv);
     expect(config.getIdeMode()).toBe(false);
   });
-
-  it('should add _ide_server when ideMode is true', async () => {
-    process.argv = ['node', 'script.js', '--ide-mode'];
-    const argv = await parseArguments();
-    process.env.TERM_PROGRAM = 'vscode';
-    process.env.GEMINI_CLI_IDE_SERVER_PORT = '3000';
-    const settings: Settings = {};
-    const config = await loadCliConfig(settings, [], 'test-session', argv);
-    expect(config.getIdeMode()).toBe(true);
-    const mcpServers = config.getMcpServers();
-    expect(mcpServers['_ide_server']).toBeDefined();
-    expect(mcpServers['_ide_server'].httpUrl).toBe('http://localhost:3000/mcp');
-    expect(mcpServers['_ide_server'].description).toBe('IDE connection');
-    expect(mcpServers['_ide_server'].trust).toBe(false);
-  });
-
-  it('should warn if ideMode is true and no port is set', async () => {
-    const consoleWarnSpy = vi
-      .spyOn(console, 'warn')
-      .mockImplementation(() => {});
-    process.argv = ['node', 'script.js', '--ide-mode'];
-    const argv = await parseArguments();
-    process.env.TERM_PROGRAM = 'vscode';
-    const settings: Settings = {};
-    await loadCliConfig(settings, [], 'test-session', argv);
-    expect(consoleWarnSpy).toHaveBeenCalledWith(
-      '[WARN]',
-      'Could not connect to IDE. Make sure you have the companion VS Code extension installed from the marketplace or via /ide install.',
-    );
-    consoleWarnSpy.mockRestore();
-  });
-
-  it('should warn and overwrite if settings contain the reserved _ide_server name and ideMode is active', async () => {
-    const consoleWarnSpy = vi
-      .spyOn(console, 'warn')
-      .mockImplementation(() => {});
-
-    process.argv = ['node', 'script.js', '--ide-mode'];
-    const argv = await parseArguments();
-    process.env.TERM_PROGRAM = 'vscode';
-    process.env.GEMINI_CLI_IDE_SERVER_PORT = '3000';
-    const settings: Settings = {
-      mcpServers: {
-        _ide_server: new ServerConfig.MCPServerConfig(
-          undefined,
-          undefined,
-          undefined,
-          undefined,
-          'http://malicious:1234',
-        ),
-      },
-    };
-
-    const config = await loadCliConfig(settings, [], 'test-session', argv);
-
-    expect(consoleWarnSpy).toHaveBeenCalledWith(
-      '[WARN]',
-      'Ignoring user-defined MCP server config for "_ide_server" as it is a reserved name.',
-    );
-
-    const mcpServers = config.getMcpServers();
-    expect(mcpServers['_ide_server']).toBeDefined();
-    expect(mcpServers['_ide_server'].httpUrl).toBe('http://localhost:3000/mcp');
-
-    consoleWarnSpy.mockRestore();
-  });
-
-  it('should NOT warn if settings contain the reserved _ide_server name and ideMode is NOT active', async () => {
-    const consoleWarnSpy = vi
-      .spyOn(console, 'warn')
-      .mockImplementation(() => {});
-
-    process.argv = ['node', 'script.js'];
-    const argv = await parseArguments();
-    const settings: Settings = {
-      mcpServers: {
-        _ide_server: new ServerConfig.MCPServerConfig(
-          undefined,
-          undefined,
-          undefined,
-          undefined,
-          'http://malicious:1234',
-        ),
-      },
-    };
-
-    const config = await loadCliConfig(settings, [], 'test-session', argv);
-
-    expect(consoleWarnSpy).not.toHaveBeenCalled();
-
-    const mcpServers = config.getMcpServers();
-    expect(mcpServers['_ide_server']).toBeDefined();
-    expect(mcpServers['_ide_server'].url).toBe('http://malicious:1234');
-
-    consoleWarnSpy.mockRestore();
-  });
 });
diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts
index 650f3aa2..27e3ec09 100644
--- a/packages/cli/src/config/config.ts
+++ b/packages/cli/src/config/config.ts
@@ -19,8 +19,7 @@ import {
   FileDiscoveryService,
   TelemetryTarget,
   FileFilteringOptions,
-  MCPServerConfig,
-  IDE_SERVER_NAME,
+  IdeClient,
 } from '@google/gemini-cli-core';
 import { Settings } from './settings.js';
 
@@ -264,6 +263,11 @@ export async function loadCliConfig(
     process.env.TERM_PROGRAM === 'vscode' &&
     !process.env.SANDBOX;
 
+  let ideClient: IdeClient | undefined;
+  if (ideMode) {
+    ideClient = new IdeClient();
+  }
+
   const allExtensions = annotateActiveExtensions(
     extensions,
     argv.extensions || [],
@@ -355,37 +359,6 @@ export async function loadCliConfig(
     }
   }
 
-  if (ideMode) {
-    if (mcpServers[IDE_SERVER_NAME]) {
-      logger.warn(
-        `Ignoring user-defined MCP server config for "${IDE_SERVER_NAME}" as it is a reserved name.`,
-      );
-    }
-    const companionPort = process.env.GEMINI_CLI_IDE_SERVER_PORT;
-    if (companionPort) {
-      const httpUrl = `http://localhost:${companionPort}/mcp`;
-      mcpServers[IDE_SERVER_NAME] = new MCPServerConfig(
-        undefined, // command
-        undefined, // args
-        undefined, // env
-        undefined, // cwd
-        undefined, // url
-        httpUrl, // httpUrl
-        undefined, // headers
-        undefined, // tcp
-        undefined, // timeout
-        false, // trust
-        'IDE connection', // description
-        undefined, // includeTools
-        undefined, // excludeTools
-      );
-    } else {
-      logger.warn(
-        'Could not connect to IDE. Make sure you have the companion VS Code extension installed from the marketplace or via /ide install.',
-      );
-    }
-  }
-
   const sandboxConfig = await loadSandboxConfig(settings, argv);
 
   return new Config({
@@ -450,6 +423,7 @@ export async function loadCliConfig(
     noBrowser: !!process.env.NO_BROWSER,
     summarizeToolOutput: settings.summarizeToolOutput,
     ideMode,
+    ideClient,
   });
 }
 
diff --git a/packages/cli/src/ui/commands/ideCommand.test.ts b/packages/cli/src/ui/commands/ideCommand.test.ts
index 51322843..d1d72466 100644
--- a/packages/cli/src/ui/commands/ideCommand.test.ts
+++ b/packages/cli/src/ui/commands/ideCommand.test.ts
@@ -19,25 +19,10 @@ import { type Config } from '@google/gemini-cli-core';
 import * as child_process from 'child_process';
 import { glob } from 'glob';
 
-import {
-  getMCPDiscoveryState,
-  getMCPServerStatus,
-  IDE_SERVER_NAME,
-  MCPDiscoveryState,
-  MCPServerStatus,
-} from '@google/gemini-cli-core';
+import { IDEConnectionStatus } from '@google/gemini-cli-core/index.js';
 
 vi.mock('child_process');
 vi.mock('glob');
-vi.mock('@google/gemini-cli-core', async (importOriginal) => {
-  const original =
-    await importOriginal<typeof import('@google/gemini-cli-core')>();
-  return {
-    ...original,
-    getMCPServerStatus: vi.fn(),
-    getMCPDiscoveryState: vi.fn(),
-  };
-});
 
 function regexEscape(value: string) {
   return value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
@@ -49,8 +34,6 @@ describe('ideCommand', () => {
   let execSyncSpy: MockInstance;
   let globSyncSpy: MockInstance;
   let platformSpy: MockInstance;
-  let getMCPServerStatusSpy: MockInstance;
-  let getMCPDiscoveryStateSpy: MockInstance;
 
   beforeEach(() => {
     mockContext = {
@@ -61,13 +44,12 @@ describe('ideCommand', () => {
 
     mockConfig = {
       getIdeMode: vi.fn(),
+      getIdeClient: vi.fn(),
     } as unknown as Config;
 
     execSyncSpy = vi.spyOn(child_process, 'execSync');
     globSyncSpy = vi.spyOn(glob, 'sync');
     platformSpy = vi.spyOn(process, 'platform', 'get');
-    getMCPServerStatusSpy = vi.mocked(getMCPServerStatus);
-    getMCPDiscoveryStateSpy = vi.mocked(getMCPDiscoveryState);
   });
 
   afterEach(() => {
@@ -91,15 +73,21 @@ describe('ideCommand', () => {
   });
 
   describe('status subcommand', () => {
+    const mockGetConnectionStatus = vi.fn();
     beforeEach(() => {
       vi.mocked(mockConfig.getIdeMode).mockReturnValue(true);
+      vi.mocked(mockConfig.getIdeClient).mockReturnValue({
+        getConnectionStatus: mockGetConnectionStatus,
+      } as ReturnType<Config['getIdeClient']>);
     });
 
     it('should show connected status', () => {
-      getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.CONNECTED);
+      mockGetConnectionStatus.mockReturnValue({
+        status: IDEConnectionStatus.Connected,
+      });
       const command = ideCommand(mockConfig);
       const result = command!.subCommands![0].action!(mockContext, '');
-      expect(getMCPServerStatusSpy).toHaveBeenCalledWith(IDE_SERVER_NAME);
+      expect(mockGetConnectionStatus).toHaveBeenCalled();
       expect(result).toEqual({
         type: 'message',
         messageType: 'info',
@@ -108,37 +96,45 @@ describe('ideCommand', () => {
     });
 
     it('should show connecting status', () => {
-      getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.CONNECTING);
+      mockGetConnectionStatus.mockReturnValue({
+        status: IDEConnectionStatus.Connecting,
+      });
       const command = ideCommand(mockConfig);
       const result = command!.subCommands![0].action!(mockContext, '');
+      expect(mockGetConnectionStatus).toHaveBeenCalled();
       expect(result).toEqual({
         type: 'message',
         messageType: 'info',
-        content: '🔄 Initializing...',
+        content: `🟡 Connecting...`,
       });
     });
-
-    it('should show discovery in progress status', () => {
-      getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.DISCONNECTED);
-      getMCPDiscoveryStateSpy.mockReturnValue(MCPDiscoveryState.IN_PROGRESS);
+    it('should show disconnected status', () => {
+      mockGetConnectionStatus.mockReturnValue({
+        status: IDEConnectionStatus.Disconnected,
+      });
       const command = ideCommand(mockConfig);
       const result = command!.subCommands![0].action!(mockContext, '');
+      expect(mockGetConnectionStatus).toHaveBeenCalled();
       expect(result).toEqual({
         type: 'message',
-        messageType: 'info',
-        content: '🔄 Initializing...',
+        messageType: 'error',
+        content: `🔴 Disconnected`,
       });
     });
 
-    it('should show disconnected status', () => {
-      getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.DISCONNECTED);
-      getMCPDiscoveryStateSpy.mockReturnValue(MCPDiscoveryState.COMPLETED);
+    it('should show disconnected status with details', () => {
+      const details = 'Something went wrong';
+      mockGetConnectionStatus.mockReturnValue({
+        status: IDEConnectionStatus.Disconnected,
+        details,
+      });
       const command = ideCommand(mockConfig);
       const result = command!.subCommands![0].action!(mockContext, '');
+      expect(mockGetConnectionStatus).toHaveBeenCalled();
       expect(result).toEqual({
         type: 'message',
         messageType: 'error',
-        content: '🔴 Disconnected',
+        content: `🔴 Disconnected: ${details}`,
       });
     });
   });
diff --git a/packages/cli/src/ui/commands/ideCommand.ts b/packages/cli/src/ui/commands/ideCommand.ts
index 6fc4f50b..31f2371f 100644
--- a/packages/cli/src/ui/commands/ideCommand.ts
+++ b/packages/cli/src/ui/commands/ideCommand.ts
@@ -5,14 +5,7 @@
  */
 
 import { fileURLToPath } from 'url';
-import {
-  Config,
-  getMCPDiscoveryState,
-  getMCPServerStatus,
-  IDE_SERVER_NAME,
-  MCPDiscoveryState,
-  MCPServerStatus,
-} from '@google/gemini-cli-core';
+import { Config, IDEConnectionStatus } from '@google/gemini-cli-core';
 import {
   CommandContext,
   SlashCommand,
@@ -56,36 +49,31 @@ export const ideCommand = (config: Config | null): SlashCommand | null => {
         description: 'check status of IDE integration',
         kind: CommandKind.BUILT_IN,
         action: (_context: CommandContext): SlashCommandActionReturn => {
-          const status = getMCPServerStatus(IDE_SERVER_NAME);
-          const discoveryState = getMCPDiscoveryState();
-          switch (status) {
-            case MCPServerStatus.CONNECTED:
+          const connection = config.getIdeClient()?.getConnectionStatus();
+          switch (connection?.status) {
+            case IDEConnectionStatus.Connected:
               return {
                 type: 'message',
                 messageType: 'info',
                 content: `🟢 Connected`,
-              };
-            case MCPServerStatus.CONNECTING:
+              } as const;
+            case IDEConnectionStatus.Connecting:
               return {
                 type: 'message',
                 messageType: 'info',
-                content: `🔄 Initializing...`,
-              };
-            case MCPServerStatus.DISCONNECTED:
-            default:
-              if (discoveryState === MCPDiscoveryState.IN_PROGRESS) {
-                return {
-                  type: 'message',
-                  messageType: 'info',
-                  content: `🔄 Initializing...`,
-                };
-              } else {
-                return {
-                  type: 'message',
-                  messageType: 'error',
-                  content: `🔴 Disconnected`,
-                };
+                content: `🟡 Connecting...`,
+              } as const;
+            default: {
+              let content = `🔴 Disconnected`;
+              if (connection?.details) {
+                content += `: ${connection.details}`;
               }
+              return {
+                type: 'message',
+                messageType: 'error',
+                content,
+              } as const;
+            }
           }
         },
       },