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

Co-authored-by: matt korwel <matt.korwel@gmail.com>
This commit is contained in:
Shreya Keshive 2025-08-05 18:52:58 -04:00 committed by GitHub
parent 6a72cd064b
commit 268627469b
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
16 changed files with 285 additions and 205 deletions

View File

@ -1007,16 +1007,6 @@ describe('loadCliConfig ideModeFeature', () => {
const config = await loadCliConfig(settings, [], 'test-session', argv);
expect(config.getIdeModeFeature()).toBe(false);
});
it('should be false when settings.ideModeFeature is true, but SANDBOX is set', async () => {
process.argv = ['node', 'script.js'];
const argv = await parseArguments();
process.env.TERM_PROGRAM = 'vscode';
process.env.SANDBOX = 'true';
const settings: Settings = { ideModeFeature: true };
const config = await loadCliConfig(settings, [], 'test-session', argv);
expect(config.getIdeModeFeature()).toBe(false);
});
});
vi.mock('fs', async () => {

View File

@ -22,7 +22,6 @@ import {
FileDiscoveryService,
TelemetryTarget,
FileFilteringOptions,
IdeClient,
} from '@google/gemini-cli-core';
import { Settings } from './settings.js';
@ -296,13 +295,10 @@ export async function loadCliConfig(
) ||
false;
const memoryImportFormat = settings.memoryImportFormat || 'tree';
const ideMode = settings.ideMode ?? false;
const ideModeFeature =
(argv.ideModeFeature ?? settings.ideModeFeature ?? false) &&
!process.env.SANDBOX;
const ideClient = IdeClient.getInstance(ideMode && ideModeFeature);
argv.ideModeFeature ?? settings.ideModeFeature ?? false;
const allExtensions = annotateActiveExtensions(
extensions,
@ -471,7 +467,6 @@ export async function loadCliConfig(
summarizeToolOutput: settings.summarizeToolOutput,
ideMode,
ideModeFeature,
ideClient,
});
}

View File

@ -42,9 +42,15 @@ describe('ideCommand', () => {
mockConfig = {
getIdeModeFeature: vi.fn(),
getIdeMode: vi.fn(),
getIdeClient: vi.fn(),
getIdeClient: vi.fn(() => ({
reconnect: vi.fn(),
disconnect: vi.fn(),
getCurrentIde: vi.fn(),
getDetectedIdeDisplayName: vi.fn(),
getConnectionStatus: vi.fn(),
})),
setIdeModeAndSyncConnection: vi.fn(),
setIdeMode: vi.fn(),
setIdeClientDisconnected: vi.fn(),
} as unknown as Config;
platformSpy = vi.spyOn(process, 'platform', 'get');

View File

@ -10,6 +10,7 @@ import {
IDEConnectionStatus,
getIdeDisplayName,
getIdeInstaller,
IdeClient,
} from '@google/gemini-cli-core';
import {
CommandContext,
@ -19,6 +20,35 @@ import {
} from './types.js';
import { SettingScope } from '../../config/settings.js';
function getIdeStatusMessage(ideClient: IdeClient): {
messageType: 'info' | 'error';
content: string;
} {
const connection = ideClient.getConnectionStatus();
switch (connection.status) {
case IDEConnectionStatus.Connected:
return {
messageType: 'info',
content: `🟢 Connected to ${ideClient.getDetectedIdeDisplayName()}`,
};
case IDEConnectionStatus.Connecting:
return {
messageType: 'info',
content: `🟡 Connecting...`,
};
default: {
let content = `🔴 Disconnected`;
if (connection?.details) {
content += `: ${connection.details}`;
}
return {
messageType: 'error',
content,
};
}
}
}
export const ideCommand = (config: Config | null): SlashCommand | null => {
if (!config || !config.getIdeModeFeature()) {
return null;
@ -54,33 +84,13 @@ export const ideCommand = (config: Config | null): SlashCommand | null => {
name: 'status',
description: 'check status of IDE integration',
kind: CommandKind.BUILT_IN,
action: (_context: CommandContext): SlashCommandActionReturn => {
const connection = ideClient.getConnectionStatus();
switch (connection.status) {
case IDEConnectionStatus.Connected:
return {
type: 'message',
messageType: 'info',
content: `🟢 Connected to ${ideClient.getDetectedIdeDisplayName()}`,
} as const;
case IDEConnectionStatus.Connecting:
return {
type: 'message',
messageType: 'info',
content: `🟡 Connecting...`,
} as const;
default: {
let content = `🔴 Disconnected`;
if (connection?.details) {
content += `: ${connection.details}`;
}
return {
type: 'message',
messageType: 'error',
content,
} as const;
}
}
action: (): SlashCommandActionReturn => {
const { messageType, content } = getIdeStatusMessage(ideClient);
return {
type: 'message',
messageType,
content,
} as const;
},
};
@ -110,6 +120,10 @@ export const ideCommand = (config: Config | null): SlashCommand | null => {
);
const result = await installer.install();
if (result.success) {
config.setIdeMode(true);
context.services.settings.setValue(SettingScope.User, 'ideMode', true);
}
context.ui.addItem(
{
type: result.success ? 'info' : 'error',
@ -126,8 +140,15 @@ export const ideCommand = (config: Config | null): SlashCommand | null => {
kind: CommandKind.BUILT_IN,
action: async (context: CommandContext) => {
context.services.settings.setValue(SettingScope.User, 'ideMode', true);
config.setIdeMode(true);
config.setIdeClientConnected();
await config.setIdeModeAndSyncConnection(true);
const { messageType, content } = getIdeStatusMessage(ideClient);
context.ui.addItem(
{
type: messageType,
text: content,
},
Date.now(),
);
},
};
@ -137,8 +158,15 @@ export const ideCommand = (config: Config | null): SlashCommand | null => {
kind: CommandKind.BUILT_IN,
action: async (context: CommandContext) => {
context.services.settings.setValue(SettingScope.User, 'ideMode', false);
config.setIdeMode(false);
config.setIdeClientDisconnected();
await config.setIdeModeAndSyncConnection(false);
const { messageType, content } = getIdeStatusMessage(ideClient);
context.ui.addItem(
{
type: messageType,
text: content,
},
Date.now(),
);
},
};

View File

@ -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(() => {

View File

@ -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> {

View File

@ -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');

View File

@ -33,154 +33,38 @@ 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;
}
getCurrentIde(): DetectedIde | undefined {
return this.currentIde;
}
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();
}
}
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.',
);
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.');
};
}
async reconnect(ideMode: boolean) {
IdeClient.instance = new IdeClient(ideMode);
}
private async establishConnection(port: string) {
let transport: StreamableHTTPClientTransport | undefined;
try {
this.client = new Client({
name: 'streamable-http-client',
// 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.setState(IDEConnectionStatus.Connected);
} catch (error) {
this.setState(
IDEConnectionStatus.Disconnected,
`Failed to connect to IDE server: ${error}`,
);
if (transport) {
try {
await transport.close();
} catch (closeError) {
logger.debug('Failed to close transport:', closeError);
}
}
}
}
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;
}
async connect(): Promise<void> {
this.setState(IDEConnectionStatus.Connecting);
if (!this.currentIde || !this.currentIdeDisplayName) {
this.setState(IDEConnectionStatus.Disconnected);
return;
}
if (!this.validateWorkspacePath()) {
return;
}
@ -193,15 +77,132 @@ export class IdeClient {
await this.establishConnection(port);
}
dispose() {
disconnect() {
this.setState(
IDEConnectionStatus.Disconnected,
'IDE integration disabled. To enable it again, run /ide enable.',
);
this.client?.close();
}
getCurrentIde(): DetectedIde | undefined {
return this.currentIde;
}
getConnectionStatus(): IDEConnectionState {
return this.state;
}
getDetectedIdeDisplayName(): string | undefined {
return this.currentIdeDisplayName;
}
setDisconnected() {
this.setState(IDEConnectionStatus.Disconnected);
private setState(status: IDEConnectionStatus, details?: string) {
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 disconnected:', details);
ideContext.clearIdeContext();
}
}
private validateWorkspacePath(): boolean {
const ideWorkspacePath = process.env['GEMINI_CLI_IDE_WORKSPACE_PATH'];
if (ideWorkspacePath === undefined) {
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 false;
}
if (ideWorkspacePath === '') {
this.setState(
IDEConnectionStatus.Disconnected,
`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,
`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;
}
this.client.setNotificationHandler(
IdeContextNotificationSchema,
(notification) => {
ideContext.setIdeContext(notification.params);
},
);
this.client.onerror = (_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,
`IDE connection error. The connection was lost unexpectedly. Please try reconnecting by running /ide enable`,
);
};
}
private async establishConnection(port: string) {
let transport: StreamableHTTPClientTransport | undefined;
try {
this.client = new Client({
name: 'streamable-http-client',
// TODO(#3487): use the CLI version here.
version: '1.0.0',
});
transport = new StreamableHTTPClientTransport(
new URL(`http://localhost:${port}/mcp`),
);
await this.client.connect(transport);
this.registerClientHandlers();
this.setState(IDEConnectionStatus.Connected);
} catch (_error) {
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.`,
);
if (transport) {
try {
await transport.close();
} catch (closeError) {
logger.debug('Failed to close transport:', closeError);
}
}
}
}
}

View File

@ -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 = [
{

View File

@ -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(

View File

@ -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';

View File

@ -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);
}

View File

@ -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(

View File

@ -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;

View File

@ -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', () => {

View File

@ -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