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

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

12 files changed, 163 insertions, 102 deletions

diff --git a/packages/core/src/config/config.test.ts b/packages/core/src/config/config.test.ts
index dd50fd41..64692139 100644
--- a/packages/core/src/config/config.test.ts
+++ b/packages/core/src/config/config.test.ts
@@ -18,7 +18,6 @@ import {
 } from '../core/contentGenerator.js';
 import { GeminiClient } from '../core/client.js';
 import { GitService } from '../services/gitService.js';
-import { IdeClient } from '../ide/ide-client.js';
 
 vi.mock('fs', async (importOriginal) => {
   const actual = await importOriginal<typeof import('fs')>();
@@ -120,7 +119,6 @@ describe('Server Config (config.ts)', () => {
     telemetry: TELEMETRY_SETTINGS,
     sessionId: SESSION_ID,
     model: MODEL,
-    ideClient: IdeClient.getInstance(false),
   };
 
   beforeEach(() => {
diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts
index 22996f3e..fa51a6af 100644
--- a/packages/core/src/config/config.ts
+++ b/packages/core/src/config/config.ts
@@ -48,6 +48,8 @@ import { shouldAttemptBrowserLaunch } from '../utils/browser.js';
 import { MCPOAuthConfig } from '../mcp/oauth-provider.js';
 import { IdeClient } from '../ide/ide-client.js';
 import type { Content } from '@google/genai';
+import { logIdeConnection } from '../telemetry/loggers.js';
+import { IdeConnectionEvent, IdeConnectionType } from '../telemetry/types.js';
 
 // Re-export OAuth config type
 export type { MCPOAuthConfig };
@@ -187,7 +189,6 @@ export interface ConfigParameters {
   summarizeToolOutput?: Record<string, SummarizeToolOutputSettings>;
   ideModeFeature?: boolean;
   ideMode?: boolean;
-  ideClient: IdeClient;
   loadMemoryFromIncludeDirectories?: boolean;
 }
 
@@ -305,7 +306,11 @@ export class Config {
     this.summarizeToolOutput = params.summarizeToolOutput;
     this.ideModeFeature = params.ideModeFeature ?? false;
     this.ideMode = params.ideMode ?? false;
-    this.ideClient = params.ideClient;
+    this.ideClient = IdeClient.getInstance();
+    if (this.ideMode && this.ideModeFeature) {
+      this.ideClient.connect();
+      logIdeConnection(this, new IdeConnectionEvent(IdeConnectionType.START));
+    }
     this.loadMemoryFromIncludeDirectories =
       params.loadMemoryFromIncludeDirectories ?? false;
 
@@ -633,10 +638,6 @@ export class Config {
     return this.ideModeFeature;
   }
 
-  getIdeClient(): IdeClient {
-    return this.ideClient;
-  }
-
   getIdeMode(): boolean {
     return this.ideMode;
   }
@@ -645,12 +646,18 @@ export class Config {
     this.ideMode = value;
   }
 
-  setIdeClientDisconnected(): void {
-    this.ideClient.setDisconnected();
+  async setIdeModeAndSyncConnection(value: boolean): Promise<void> {
+    this.ideMode = value;
+    if (value) {
+      await this.ideClient.connect();
+      logIdeConnection(this, new IdeConnectionEvent(IdeConnectionType.SESSION));
+    } else {
+      this.ideClient.disconnect();
+    }
   }
 
-  setIdeClientConnected(): void {
-    this.ideClient.reconnect(this.ideMode && this.ideModeFeature);
+  getIdeClient(): IdeClient {
+    return this.ideClient;
   }
 
   async getGitService(): Promise<GitService> {
diff --git a/packages/core/src/config/flashFallback.test.ts b/packages/core/src/config/flashFallback.test.ts
index 0b68f993..5665a7e0 100644
--- a/packages/core/src/config/flashFallback.test.ts
+++ b/packages/core/src/config/flashFallback.test.ts
@@ -7,7 +7,7 @@
 import { describe, it, expect, beforeEach, vi } from 'vitest';
 import { Config } from './config.js';
 import { DEFAULT_GEMINI_MODEL, DEFAULT_GEMINI_FLASH_MODEL } from './models.js';
-import { IdeClient } from '../ide/ide-client.js';
+
 import fs from 'node:fs';
 
 vi.mock('node:fs');
@@ -26,7 +26,6 @@ describe('Flash Model Fallback Configuration', () => {
       debugMode: false,
       cwd: '/test',
       model: DEFAULT_GEMINI_MODEL,
-      ideClient: IdeClient.getInstance(false),
     });
 
     // Initialize contentGeneratorConfig for testing
@@ -51,7 +50,6 @@ describe('Flash Model Fallback Configuration', () => {
         debugMode: false,
         cwd: '/test',
         model: DEFAULT_GEMINI_MODEL,
-        ideClient: IdeClient.getInstance(false),
       });
 
       // Should not crash when contentGeneratorConfig is undefined
@@ -75,7 +73,6 @@ describe('Flash Model Fallback Configuration', () => {
         debugMode: false,
         cwd: '/test',
         model: 'custom-model',
-        ideClient: IdeClient.getInstance(false),
       });
 
       expect(newConfig.getModel()).toBe('custom-model');
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);
-  }
 }
diff --git a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts
index 6b85a664..649d82b6 100644
--- a/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts
+++ b/packages/core/src/telemetry/clearcut-logger/clearcut-logger.ts
@@ -21,6 +21,7 @@ import {
   NextSpeakerCheckEvent,
   SlashCommandEvent,
   MalformedJsonResponseEvent,
+  IdeConnectionEvent,
 } from '../types.js';
 import { EventMetadataKey } from './event-metadata-key.js';
 import { Config } from '../../config/config.js';
@@ -44,6 +45,7 @@ const loop_detected_event_name = 'loop_detected';
 const next_speaker_check_event_name = 'next_speaker_check';
 const slash_command_event_name = 'slash_command';
 const malformed_json_response_event_name = 'malformed_json_response';
+const ide_connection_event_name = 'ide_connection';
 
 export interface LogResponse {
   nextRequestWaitMs?: number;
@@ -578,6 +580,18 @@ export class ClearcutLogger {
     this.flushIfNeeded();
   }
 
+  logIdeConnectionEvent(event: IdeConnectionEvent): void {
+    const data = [
+      {
+        gemini_cli_key: EventMetadataKey.GEMINI_CLI_IDE_CONNECTION_TYPE,
+        value: JSON.stringify(event.connection_type),
+      },
+    ];
+
+    this.enqueueLogEvent(this.createLogEvent(ide_connection_event_name, data));
+    this.flushIfNeeded();
+  }
+
   logEndSessionEvent(event: EndSessionEvent): void {
     const data = [
       {
diff --git a/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts b/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts
index 0fc35894..54f570f1 100644
--- a/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts
+++ b/packages/core/src/telemetry/clearcut-logger/event-metadata-key.ts
@@ -190,6 +190,13 @@ export enum EventMetadataKey {
 
   // Logs the model that produced the malformed JSON response.
   GEMINI_CLI_MALFORMED_JSON_RESPONSE_MODEL = 45,
+
+  // ==========================================================================
+  // IDE Connection Event Keys
+  // ===========================================================================
+
+  // Logs the type of the IDE connection.
+  GEMINI_CLI_IDE_CONNECTION_TYPE = 46,
 }
 
 export function getEventMetadataKey(
diff --git a/packages/core/src/telemetry/constants.ts b/packages/core/src/telemetry/constants.ts
index 7dd5c8d1..7d840815 100644
--- a/packages/core/src/telemetry/constants.ts
+++ b/packages/core/src/telemetry/constants.ts
@@ -15,6 +15,7 @@ export const EVENT_CLI_CONFIG = 'gemini_cli.config';
 export const EVENT_FLASH_FALLBACK = 'gemini_cli.flash_fallback';
 export const EVENT_NEXT_SPEAKER_CHECK = 'gemini_cli.next_speaker_check';
 export const EVENT_SLASH_COMMAND = 'gemini_cli.slash_command';
+export const EVENT_IDE_CONNECTION = 'gemini_cli.ide_connection';
 
 export const METRIC_TOOL_CALL_COUNT = 'gemini_cli.tool.call.count';
 export const METRIC_TOOL_CALL_LATENCY = 'gemini_cli.tool.call.latency';
diff --git a/packages/core/src/telemetry/loggers.ts b/packages/core/src/telemetry/loggers.ts
index 2aa0d86a..e3726ccb 100644
--- a/packages/core/src/telemetry/loggers.ts
+++ b/packages/core/src/telemetry/loggers.ts
@@ -12,6 +12,7 @@ import {
   EVENT_API_REQUEST,
   EVENT_API_RESPONSE,
   EVENT_CLI_CONFIG,
+  EVENT_IDE_CONNECTION,
   EVENT_TOOL_CALL,
   EVENT_USER_PROMPT,
   EVENT_FLASH_FALLBACK,
@@ -23,6 +24,7 @@ import {
   ApiErrorEvent,
   ApiRequestEvent,
   ApiResponseEvent,
+  IdeConnectionEvent,
   StartSessionEvent,
   ToolCallEvent,
   UserPromptEvent,
@@ -355,3 +357,24 @@ export function logSlashCommand(
   };
   logger.emit(logRecord);
 }
+
+export function logIdeConnection(
+  config: Config,
+  event: IdeConnectionEvent,
+): void {
+  ClearcutLogger.getInstance(config)?.logIdeConnectionEvent(event);
+  if (!isTelemetrySdkInitialized()) return;
+
+  const attributes: LogAttributes = {
+    ...getCommonAttributes(config),
+    ...event,
+    'event.name': EVENT_IDE_CONNECTION,
+  };
+
+  const logger = logs.getLogger(SERVICE_NAME);
+  const logRecord: LogRecord = {
+    body: `Ide connection. Type: ${event.connection_type}.`,
+    attributes,
+  };
+  logger.emit(logRecord);
+}
diff --git a/packages/core/src/telemetry/telemetry.test.ts b/packages/core/src/telemetry/telemetry.test.ts
index 8ebb3d9a..9734e382 100644
--- a/packages/core/src/telemetry/telemetry.test.ts
+++ b/packages/core/src/telemetry/telemetry.test.ts
@@ -12,7 +12,6 @@ import {
 } from './sdk.js';
 import { Config } from '../config/config.js';
 import { NodeSDK } from '@opentelemetry/sdk-node';
-import { IdeClient } from '../ide/ide-client.js';
 
 vi.mock('@opentelemetry/sdk-node');
 vi.mock('../config/config.js');
@@ -30,7 +29,6 @@ describe('telemetry', () => {
       targetDir: '/test/dir',
       debugMode: false,
       cwd: '/test/dir',
-      ideClient: IdeClient.getInstance(false),
     });
     vi.spyOn(mockConfig, 'getTelemetryEnabled').mockReturnValue(true);
     vi.spyOn(mockConfig, 'getTelemetryOtlpEndpoint').mockReturnValue(
diff --git a/packages/core/src/telemetry/types.ts b/packages/core/src/telemetry/types.ts
index 9d1fd77a..668421f0 100644
--- a/packages/core/src/telemetry/types.ts
+++ b/packages/core/src/telemetry/types.ts
@@ -308,6 +308,23 @@ export class MalformedJsonResponseEvent {
   }
 }
 
+export enum IdeConnectionType {
+  START = 'start',
+  SESSION = 'session',
+}
+
+export class IdeConnectionEvent {
+  'event.name': 'ide_connection';
+  'event.timestamp': string; // ISO 8601
+  connection_type: IdeConnectionType;
+
+  constructor(connection_type: IdeConnectionType) {
+    this['event.name'] = 'ide_connection';
+    this['event.timestamp'] = new Date().toISOString();
+    this.connection_type = connection_type;
+  }
+}
+
 export type TelemetryEvent =
   | StartSessionEvent
   | EndSessionEvent
@@ -320,4 +337,5 @@ export type TelemetryEvent =
   | LoopDetectedEvent
   | NextSpeakerCheckEvent
   | SlashCommandEvent
-  | MalformedJsonResponseEvent;
+  | MalformedJsonResponseEvent
+  | IdeConnectionEvent;
diff --git a/packages/core/src/tools/tool-registry.test.ts b/packages/core/src/tools/tool-registry.test.ts
index 88b23d84..24b6ca5f 100644
--- a/packages/core/src/tools/tool-registry.test.ts
+++ b/packages/core/src/tools/tool-registry.test.ts
@@ -30,7 +30,7 @@ import {
   Schema,
 } from '@google/genai';
 import { spawn } from 'node:child_process';
-import { IdeClient } from '../ide/ide-client.js';
+
 import fs from 'node:fs';
 
 vi.mock('node:fs');
@@ -140,7 +140,6 @@ const baseConfigParams: ConfigParameters = {
   geminiMdFileCount: 0,
   approvalMode: ApprovalMode.DEFAULT,
   sessionId: 'test-session-id',
-  ideClient: IdeClient.getInstance(false),
 };
 
 describe('ToolRegistry', () => {
diff --git a/packages/core/src/utils/flashFallback.integration.test.ts b/packages/core/src/utils/flashFallback.integration.test.ts
index 7f18b24f..9211ad2f 100644
--- a/packages/core/src/utils/flashFallback.integration.test.ts
+++ b/packages/core/src/utils/flashFallback.integration.test.ts
@@ -17,7 +17,6 @@ import {
 import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js';
 import { retryWithBackoff } from './retry.js';
 import { AuthType } from '../core/contentGenerator.js';
-import { IdeClient } from '../ide/ide-client.js';
 
 vi.mock('node:fs');
 
@@ -35,7 +34,6 @@ describe('Flash Fallback Integration', () => {
       debugMode: false,
       cwd: '/test',
       model: 'gemini-2.5-pro',
-      ideClient: IdeClient.getInstance(false),
     });
 
     // Reset simulation state for each test