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

This commit is contained in:
christine betts 2025-07-25 17:46:55 +00:00 committed by GitHub
parent 3c16429fc4
commit 1b8ba5ca6b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
14 changed files with 178 additions and 256 deletions

View File

@ -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();
});
});

View File

@ -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,
});
}

View File

@ -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);
const command = ideCommand(mockConfig);
const result = command!.subCommands![0].action!(mockContext, '');
expect(result).toEqual({
type: 'message',
messageType: 'info',
content: '🔄 Initializing...',
});
});
it('should show disconnected status', () => {
getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.DISCONNECTED);
getMCPDiscoveryStateSpy.mockReturnValue(MCPDiscoveryState.COMPLETED);
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: 'error',
content: '🔴 Disconnected',
content: `🔴 Disconnected`,
});
});
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: ${details}`,
});
});
});

View File

@ -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;
}
}
},
},

View File

@ -31,6 +31,7 @@
"@types/glob": "^8.1.0",
"@types/html-to-text": "^9.0.4",
"ajv": "^8.17.1",
"chardet": "^2.1.0",
"diff": "^7.0.0",
"dotenv": "^17.1.0",
"glob": "^10.4.5",
@ -44,8 +45,7 @@
"simple-git": "^3.28.0",
"strip-ansi": "^7.1.0",
"undici": "^7.10.0",
"ws": "^8.18.0",
"chardet": "^2.1.0"
"ws": "^8.18.0"
},
"devDependencies": {
"@types/diff": "^7.0.2",

View File

@ -45,6 +45,7 @@ import {
import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js';
import { shouldAttemptBrowserLaunch } from '../utils/browser.js';
import { MCPOAuthConfig } from '../mcp/oauth-provider.js';
import { IdeClient } from '../ide/ide-client.js';
// Re-export OAuth config type
export type { MCPOAuthConfig };
@ -180,6 +181,7 @@ export interface ConfigParameters {
noBrowser?: boolean;
summarizeToolOutput?: Record<string, SummarizeToolOutputSettings>;
ideMode?: boolean;
ideClient?: IdeClient;
}
export class Config {
@ -221,6 +223,7 @@ export class Config {
private readonly extensionContextFilePaths: string[];
private readonly noBrowser: boolean;
private readonly ideMode: boolean;
private readonly ideClient: IdeClient | undefined;
private modelSwitchedDuringSession: boolean = false;
private readonly maxSessionTurns: number;
private readonly listExtensions: boolean;
@ -286,6 +289,7 @@ export class Config {
this.noBrowser = params.noBrowser ?? false;
this.summarizeToolOutput = params.summarizeToolOutput;
this.ideMode = params.ideMode ?? false;
this.ideClient = params.ideClient;
if (params.contextFileName) {
setGeminiMdFilename(params.contextFileName);
@ -574,6 +578,10 @@ export class Config {
return this.ideMode;
}
getIdeClient(): IdeClient | undefined {
return this.ideClient;
}
async getGitService(): Promise<GitService> {
if (!this.gitService) {
this.gitService = new GitService(this.targetDir);

View File

@ -23,7 +23,7 @@ import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
import { FileDiscoveryService } from '../services/fileDiscoveryService.js';
import { setSimulate429 } from '../utils/testUtils.js';
import { tokenLimit } from './tokenLimits.js';
import { ideContext } from '../services/ideContext.js';
import { ideContext } from '../ide/ideContext.js';
// --- Mocks ---
const mockChatCreateFn = vi.fn();
@ -72,7 +72,7 @@ vi.mock('../telemetry/index.js', () => ({
logApiResponse: vi.fn(),
logApiError: vi.fn(),
}));
vi.mock('../services/ideContext.js');
vi.mock('../ide/ideContext.js');
describe('findIndexAfterFraction', () => {
const history: Content[] = [

View File

@ -42,7 +42,7 @@ import {
import { ProxyAgent, setGlobalDispatcher } from 'undici';
import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
import { LoopDetectionService } from '../services/loopDetectionService.js';
import { ideContext } from '../services/ideContext.js';
import { ideContext } from '../ide/ideContext.js';
import { logFlashDecidedToContinue } from '../telemetry/loggers.js';
import { FlashDecidedToContinueEvent } from '../telemetry/types.js';

View File

@ -0,0 +1,100 @@
/**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { ideContext, OpenFilesNotificationSchema } from '../ide/ideContext.js';
import { Client } from '@modelcontextprotocol/sdk/client/index.js';
import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js';
const logger = {
// eslint-disable-next-line @typescript-eslint/no-explicit-any
debug: (...args: any[]) =>
console.debug('[DEBUG] [ImportProcessor]', ...args),
};
export type IDEConnectionState = {
status: IDEConnectionStatus;
details?: string;
};
export enum IDEConnectionStatus {
Connected = 'connected',
Disconnected = 'disconnected',
Connecting = 'connecting',
}
/**
* Manages the connection to and interaction with the IDE server.
*/
export class IdeClient {
client: Client | undefined = undefined;
connectionStatus: IDEConnectionStatus = IDEConnectionStatus.Disconnected;
constructor() {
this.connectToMcpServer().catch((err) => {
logger.debug('Failed to initialize IdeClient:', err);
});
}
getConnectionStatus(): {
status: IDEConnectionStatus;
details?: string;
} {
let details: string | undefined;
if (this.connectionStatus === IDEConnectionStatus.Disconnected) {
if (!process.env['GEMINI_CLI_IDE_SERVER_PORT']) {
details = 'GEMINI_CLI_IDE_SERVER_PORT environment variable is not set.';
}
}
return {
status: this.connectionStatus,
details,
};
}
async connectToMcpServer(): Promise<void> {
this.connectionStatus = IDEConnectionStatus.Connecting;
const idePort = process.env['GEMINI_CLI_IDE_SERVER_PORT'];
if (!idePort) {
logger.debug(
'Unable to connect to IDE mode MCP server. GEMINI_CLI_IDE_SERVER_PORT environment variable is not set.',
);
this.connectionStatus = IDEConnectionStatus.Disconnected;
return;
}
try {
this.client = new Client({
name: 'streamable-http-client',
// TODO(#3487): use the CLI version here.
version: '1.0.0',
});
const transport = new StreamableHTTPClientTransport(
new URL(`http://localhost:${idePort}/mcp`),
);
await this.client.connect(transport);
this.client.setNotificationHandler(
OpenFilesNotificationSchema,
(notification) => {
ideContext.setOpenFilesContext(notification.params);
},
);
this.client.onerror = (error) => {
logger.debug('IDE MCP client error:', error);
this.connectionStatus = IDEConnectionStatus.Disconnected;
ideContext.clearOpenFilesContext();
};
this.client.onclose = () => {
logger.debug('IDE MCP client connection closed.');
this.connectionStatus = IDEConnectionStatus.Disconnected;
ideContext.clearOpenFilesContext();
};
this.connectionStatus = IDEConnectionStatus.Connected;
} catch (error) {
this.connectionStatus = IDEConnectionStatus.Disconnected;
logger.debug('Failed to connect to MCP server:', error);
}
}
}

View File

@ -6,10 +6,6 @@
import { z } from 'zod';
/**
* The reserved server name for the IDE's MCP server.
*/
export const IDE_SERVER_NAME = '_ide_server';
/**
* Zod schema for validating a cursor position.
*/

View File

@ -40,7 +40,10 @@ export * from './utils/systemEncoding.js';
// Export services
export * from './services/fileDiscoveryService.js';
export * from './services/gitService.js';
export * from './services/ideContext.js';
// Export IDE specific logic
export * from './ide/ide-client.js';
export * from './ide/ideContext.js';
// Export base tool definitions
export * from './tools/tools.js';

View File

@ -24,11 +24,6 @@ import { ToolRegistry } from './tool-registry.js';
import { MCPOAuthProvider } from '../mcp/oauth-provider.js';
import { OAuthUtils } from '../mcp/oauth-utils.js';
import { MCPOAuthTokenStorage } from '../mcp/oauth-token-storage.js';
import {
OpenFilesNotificationSchema,
IDE_SERVER_NAME,
ideContext,
} from '../services/ideContext.js';
import { getErrorMessage } from '../utils/errors.js';
export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes
@ -379,24 +374,11 @@ export async function connectAndDiscover(
);
try {
updateMCPServerStatus(mcpServerName, MCPServerStatus.CONNECTED);
mcpClient.onerror = (error) => {
console.error(`MCP ERROR (${mcpServerName}):`, error.toString());
updateMCPServerStatus(mcpServerName, MCPServerStatus.DISCONNECTED);
if (mcpServerName === IDE_SERVER_NAME) {
ideContext.clearOpenFilesContext();
}
};
if (mcpServerName === IDE_SERVER_NAME) {
mcpClient.setNotificationHandler(
OpenFilesNotificationSchema,
(notification) => {
ideContext.setOpenFilesContext(notification.params);
},
);
}
const tools = await discoverTools(
mcpServerName,
mcpServerConfig,

View File

@ -244,34 +244,5 @@ const createMcpServer = () => {
},
{ capabilities: { logging: {} } },
);
server.registerTool(
'getOpenFiles',
{
description:
'(IDE Tool) Get the path of the file currently active in VS Code.',
inputSchema: {},
},
async () => {
const editor = vscode.window.activeTextEditor;
const filePath =
editor && editor.document.uri.scheme === 'file'
? editor.document.uri.fsPath
: '';
if (filePath) {
return {
content: [{ type: 'text', text: `Active file: ${filePath}` }],
};
} else {
return {
content: [
{
type: 'text',
text: 'No file is currently active in the editor.',
},
],
};
}
},
);
return server;
};