1 files changed, 92 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);
+  }
 }