Refactor IDE client state management, improve user-facing error messages, and add logging of connection events (#5591)

Co-authored-by: matt korwel <[email protected]>

1 files changed, 80 insertions, 79 deletions

diff --git a/packages/core/src/ide/ide-client.ts b/packages/core/src/ide/ide-client.ts
index be24db3e..8f967147 100644
--- a/packages/core/src/ide/ide-client.ts
+++ b/packages/core/src/ide/ide-client.ts
@@ -33,34 +33,58 @@ export enum IDEConnectionStatus {
  * Manages the connection to and interaction with the IDE server.
  */
 export class IdeClient {
-  client: Client | undefined = undefined;
+  private static instance: IdeClient;
+  private client: Client | undefined = undefined;
   private state: IDEConnectionState = {
     status: IDEConnectionStatus.Disconnected,
+    details:
+      'IDE integration is currently disabled. To enable it, run /ide enable.',
   };
-  private static instance: IdeClient;
   private readonly currentIde: DetectedIde | undefined;
   private readonly currentIdeDisplayName: string | undefined;
 
-  constructor(ideMode: boolean) {
+  private constructor() {
     this.currentIde = detectIde();
     if (this.currentIde) {
       this.currentIdeDisplayName = getIdeDisplayName(this.currentIde);
     }
-    if (!ideMode) {
-      return;
-    }
-    this.init().catch((err) => {
-      logger.debug('Failed to initialize IdeClient:', err);
-    });
   }
 
-  static getInstance(ideMode: boolean): IdeClient {
+  static getInstance(): IdeClient {
     if (!IdeClient.instance) {
-      IdeClient.instance = new IdeClient(ideMode);
+      IdeClient.instance = new IdeClient();
     }
     return IdeClient.instance;
   }
 
+  async connect(): Promise<void> {
+    this.setState(IDEConnectionStatus.Connecting);
+
+    if (!this.currentIde || !this.currentIdeDisplayName) {
+      this.setState(IDEConnectionStatus.Disconnected);
+      return;
+    }
+
+    if (!this.validateWorkspacePath()) {
+      return;
+    }
+
+    const port = this.getPortFromEnv();
+    if (!port) {
+      return;
+    }
+
+    await this.establishConnection(port);
+  }
+
+  disconnect() {
+    this.setState(
+      IDEConnectionStatus.Disconnected,
+      'IDE integration disabled. To enable it again, run /ide enable.',
+    );
+    this.client?.close();
+  }
+
   getCurrentIde(): DetectedIde | undefined {
     return this.currentIde;
   }
@@ -69,46 +93,65 @@ export class IdeClient {
     return this.state;
   }
 
+  getDetectedIdeDisplayName(): string | undefined {
+    return this.currentIdeDisplayName;
+  }
+
   private setState(status: IDEConnectionStatus, details?: string) {
-    this.state = { status, details };
+    const isAlreadyDisconnected =
+      this.state.status === IDEConnectionStatus.Disconnected &&
+      status === IDEConnectionStatus.Disconnected;
+
+    // Only update details if the state wasn't already disconnected, so that
+    // the first detail message is preserved.
+    if (!isAlreadyDisconnected) {
+      this.state = { status, details };
+    }
 
     if (status === IDEConnectionStatus.Disconnected) {
-      logger.debug('IDE integration is disconnected. ', details);
+      logger.debug('IDE integration disconnected:', details);
       ideContext.clearIdeContext();
     }
   }
 
-  private getPortFromEnv(): string | undefined {
-    const port = process.env['GEMINI_CLI_IDE_SERVER_PORT'];
-    if (!port) {
+  private validateWorkspacePath(): boolean {
+    const ideWorkspacePath = process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'];
+    if (ideWorkspacePath === undefined) {
       this.setState(
         IDEConnectionStatus.Disconnected,
-        'Gemini CLI Companion extension not found. Install via /ide install and restart the CLI in a fresh terminal window.',
+        `Failed to connect to IDE companion extension for ${this.currentIdeDisplayName}. Please ensure the extension is running and try refreshing your terminal. To install the extension, run /ide install.`,
       );
-      return undefined;
+      return false;
     }
-    return port;
-  }
-
-  private validateWorkspacePath(): boolean {
-    const ideWorkspacePath = process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'];
-    if (!ideWorkspacePath) {
+    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.',
+        `To use this feature, please open a single workspace folder in ${this.currentIdeDisplayName} 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.`,
+        `Directory mismatch. Gemini CLI is running in a different location than the open workspace in ${this.currentIdeDisplayName}. Please run the CLI from the same directory as your project's root folder.`,
       );
       return false;
     }
     return true;
   }
 
+  private getPortFromEnv(): string | undefined {
+    const port = process.env['GEMINI_CLI_IDE_SERVER_PORT'];
+    if (!port) {
+      this.setState(
+        IDEConnectionStatus.Disconnected,
+        `Failed to connect to IDE companion extension for ${this.currentIdeDisplayName}. Please ensure the extension is running and try refreshing your terminal. To install the extension, run /ide install.`,
+      );
+      return undefined;
+    }
+    return port;
+  }
+
   private registerClientHandlers() {
     if (!this.client) {
       return;
@@ -120,20 +163,20 @@ export class IdeClient {
         ideContext.setIdeContext(notification.params);
       },
     );
-
     this.client.onerror = (_error) => {
-      this.setState(IDEConnectionStatus.Disconnected, 'Client error.');
+      this.setState(
+        IDEConnectionStatus.Disconnected,
+        `IDE connection error. The connection was lost unexpectedly. Please try reconnecting by running /ide enable`,
+      );
     };
-
     this.client.onclose = () => {
-      this.setState(IDEConnectionStatus.Disconnected, 'Connection closed.');
+      this.setState(
+        IDEConnectionStatus.Disconnected,
+        `IDE connection error. The connection was lost unexpectedly. Please try reconnecting by running /ide enable`,
+      );
     };
   }
 
-  async reconnect(ideMode: boolean) {
-    IdeClient.instance = new IdeClient(ideMode);
-  }
-
   private async establishConnection(port: string) {
     let transport: StreamableHTTPClientTransport | undefined;
     try {
@@ -142,20 +185,16 @@ export class IdeClient {
         // TODO(#3487): use the CLI version here.
         version: '1.0.0',
       });
-
       transport = new StreamableHTTPClientTransport(
         new URL(`http://localhost:${port}/mcp`),
       );
-
-      this.registerClientHandlers();
-
       await this.client.connect(transport);
-
+      this.registerClientHandlers();
       this.setState(IDEConnectionStatus.Connected);
-    } catch (error) {
+    } catch (_error) {
       this.setState(
         IDEConnectionStatus.Disconnected,
-        `Failed to connect to IDE server: ${error}`,
+        `Failed to connect to IDE companion extension for ${this.currentIdeDisplayName}. Please ensure the extension is running and try refreshing your terminal. To install the extension, run /ide install.`,
       );
       if (transport) {
         try {
@@ -166,42 +205,4 @@ export class IdeClient {
       }
     }
   }
-
-  async init(): Promise<void> {
-    if (this.state.status === IDEConnectionStatus.Connected) {
-      return;
-    }
-    if (!this.currentIde) {
-      this.setState(
-        IDEConnectionStatus.Disconnected,
-        'Not running in a supported IDE, skipping connection.',
-      );
-      return;
-    }
-
-    this.setState(IDEConnectionStatus.Connecting);
-
-    if (!this.validateWorkspacePath()) {
-      return;
-    }
-
-    const port = this.getPortFromEnv();
-    if (!port) {
-      return;
-    }
-
-    await this.establishConnection(port);
-  }
-
-  dispose() {
-    this.client?.close();
-  }
-
-  getDetectedIdeDisplayName(): string | undefined {
-    return this.currentIdeDisplayName;
-  }
-
-  setDisconnected() {
-    this.setState(IDEConnectionStatus.Disconnected);
-  }
 }