Only enable IDE integration if gemini-cli is running in the same path as open workspace (#5068)

2 files changed, 116 insertions, 44 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.';
-      }
+  getConnectionStatus(): IDEConnectionState {
+    return this.state;
+  }
+
+  private setState(status: IDEConnectionStatus, details?: string) {
+    this.state = { status, details };
+
+    if (status === IDEConnectionStatus.Disconnected) {
+      logger.debug('IDE integration is disconnected. ', details);
+      ideContext.clearIdeContext();
     }
-    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.',
+  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`),
       );
-      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.registerClientHandlers();
 
-      this.connectionStatus = IDEConnectionStatus.Connected;
+      await this.client.connect(transport);
+
+      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<void> {
+    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`);