From 83c4dddb7ee7ba34d7dec09d00819972d2e1ff5f Mon Sep 17 00:00:00 2001 From: Shreya Keshive Date: Mon, 28 Jul 2025 16:55:00 -0400 Subject: [PATCH] Only enable IDE integration if gemini-cli is running in the same path as open workspace (#5068) --- packages/core/src/ide/ide-client.ts | 138 ++++++++++++------ .../vscode-ide-companion/src/extension.ts | 24 +++ 2 files changed, 117 insertions(+), 45 deletions(-) diff --git a/packages/core/src/ide/ide-client.ts b/packages/core/src/ide/ide-client.ts index 64264fd1..3c670d54 100644 --- a/packages/core/src/ide/ide-client.ts +++ b/packages/core/src/ide/ide-client.ts @@ -15,7 +15,7 @@ const logger = { export type IDEConnectionState = { status: IDEConnectionStatus; - details?: string; + details?: string; // User-facing }; export enum IDEConnectionStatus { @@ -29,41 +29,82 @@ export enum IDEConnectionStatus { */ export class IdeClient { client: Client | undefined = undefined; - connectionStatus: IDEConnectionStatus = IDEConnectionStatus.Disconnected; + private state: IDEConnectionState = { + status: IDEConnectionStatus.Disconnected, + }; constructor() { - this.connectToMcpServer().catch((err) => { + this.init().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, - }; + getConnectionStatus(): IDEConnectionState { + return this.state; } - async connectToMcpServer(): Promise { - 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.', + private setState(status: IDEConnectionStatus, details?: string) { + this.state = { status, details }; + + if (status === IDEConnectionStatus.Disconnected) { + logger.debug('IDE integration is disconnected. ', details); + ideContext.clearIdeContext(); + } + } + + private getPortFromEnv(): string | undefined { + const port = process.env['GEMINI_CLI_IDE_SERVER_PORT']; + if (!port) { + this.setState( + IDEConnectionStatus.Disconnected, + 'Gemini CLI Companion extension not found. Install via /ide install and restart the CLI in a fresh terminal window.', ); - this.connectionStatus = IDEConnectionStatus.Disconnected; + return undefined; + } + return port; + } + + private validateWorkspacePath(): boolean { + const ideWorkspacePath = process.env['GEMINI_CLI_IDE_WORKSPACE_PATH']; + if (!ideWorkspacePath) { + this.setState( + IDEConnectionStatus.Disconnected, + 'IDE integration requires a single workspace folder to be open in the IDE. Please ensure one folder is open and try again.', + ); + return false; + } + if (ideWorkspacePath !== process.cwd()) { + this.setState( + IDEConnectionStatus.Disconnected, + `Gemini CLI is running in a different directory (${process.cwd()}) from the IDE's open workspace (${ideWorkspacePath}). Please run Gemini CLI in the same directory.`, + ); + return false; + } + return true; + } + + private registerClientHandlers() { + if (!this.client) { return; } + this.client.setNotificationHandler( + IdeContextNotificationSchema, + (notification) => { + ideContext.setIdeContext(notification.params); + }, + ); + + this.client.onerror = (_error) => { + this.setState(IDEConnectionStatus.Disconnected, 'Client error.'); + }; + + this.client.onclose = () => { + this.setState(IDEConnectionStatus.Disconnected, 'Connection closed.'); + }; + } + + private async establishConnection(port: string) { let transport: StreamableHTTPClientTransport | undefined; try { this.client = new Client({ @@ -71,32 +112,21 @@ export class IdeClient { // TODO(#3487): use the CLI version here. version: '1.0.0', }); + transport = new StreamableHTTPClientTransport( - new URL(`http://localhost:${idePort}/mcp`), + new URL(`http://localhost:${port}/mcp`), ); + + this.registerClientHandlers(); + await this.client.connect(transport); - this.client.setNotificationHandler( - IdeContextNotificationSchema, - (notification) => { - ideContext.setIdeContext(notification.params); - }, - ); - this.client.onerror = (error) => { - logger.debug('IDE MCP client error:', error); - this.connectionStatus = IDEConnectionStatus.Disconnected; - ideContext.clearIdeContext(); - }; - this.client.onclose = () => { - logger.debug('IDE MCP client connection closed.'); - this.connectionStatus = IDEConnectionStatus.Disconnected; - ideContext.clearIdeContext(); - }; - - this.connectionStatus = IDEConnectionStatus.Connected; + this.setState(IDEConnectionStatus.Connected); } catch (error) { - this.connectionStatus = IDEConnectionStatus.Disconnected; - logger.debug('Failed to connect to MCP server:', error); + this.setState( + IDEConnectionStatus.Disconnected, + `Failed to connect to IDE server: ${error}`, + ); if (transport) { try { await transport.close(); @@ -106,4 +136,22 @@ export class IdeClient { } } } + + async init(): Promise { + if (this.state.status === IDEConnectionStatus.Connected) { + return; + } + this.setState(IDEConnectionStatus.Connecting); + + if (!this.validateWorkspacePath()) { + return; + } + + const port = this.getPortFromEnv(); + if (!port) { + return; + } + + await this.establishConnection(port); + } } diff --git a/packages/vscode-ide-companion/src/extension.ts b/packages/vscode-ide-companion/src/extension.ts index 647acae3..637b69e3 100644 --- a/packages/vscode-ide-companion/src/extension.ts +++ b/packages/vscode-ide-companion/src/extension.ts @@ -8,14 +8,35 @@ import * as vscode from 'vscode'; import { IDEServer } from './ide-server'; import { createLogger } from './utils/logger'; +const IDE_WORKSPACE_PATH_ENV_VAR = 'GEMINI_CLI_IDE_WORKSPACE_PATH'; + let ideServer: IDEServer; let logger: vscode.OutputChannel; let log: (message: string) => void = () => {}; +function updateWorkspacePath(context: vscode.ExtensionContext) { + const workspaceFolders = vscode.workspace.workspaceFolders; + if (workspaceFolders && workspaceFolders.length === 1) { + const workspaceFolder = workspaceFolders[0]; + context.environmentVariableCollection.replace( + IDE_WORKSPACE_PATH_ENV_VAR, + workspaceFolder.uri.fsPath, + ); + } else { + context.environmentVariableCollection.replace( + IDE_WORKSPACE_PATH_ENV_VAR, + '', + ); + } +} + export async function activate(context: vscode.ExtensionContext) { logger = vscode.window.createOutputChannel('Gemini CLI IDE Companion'); log = createLogger(context, logger); log('Extension activated'); + + updateWorkspacePath(context); + ideServer = new IDEServer(log); try { await ideServer.start(context); @@ -25,6 +46,9 @@ export async function activate(context: vscode.ExtensionContext) { } context.subscriptions.push( + vscode.workspace.onDidChangeWorkspaceFolders(() => { + updateWorkspacePath(context); + }), vscode.commands.registerCommand('gemini-cli.runGeminiCLI', () => { const geminiCmd = 'gemini'; const terminal = vscode.window.createTerminal(`Gemini CLI`);