From 1b8ba5ca6bf739e4100a1d313721988f953acb49 Mon Sep 17 00:00:00 2001 From: christine betts Date: Fri, 25 Jul 2025 17:46:55 +0000 Subject: [ide-mode] Create an IDE manager class to handle connecting to and exposing methods from the IDE server (#4797) --- packages/core/src/ide/ide-client.ts | 100 ++++++++++++++++++++++ packages/core/src/ide/ideContext.test.ts | 140 +++++++++++++++++++++++++++++++ packages/core/src/ide/ideContext.ts | 118 ++++++++++++++++++++++++++ 3 files changed, 358 insertions(+) create mode 100644 packages/core/src/ide/ide-client.ts create mode 100644 packages/core/src/ide/ideContext.test.ts create mode 100644 packages/core/src/ide/ideContext.ts (limited to 'packages/core/src/ide') diff --git a/packages/core/src/ide/ide-client.ts b/packages/core/src/ide/ide-client.ts new file mode 100644 index 00000000..eeed60b2 --- /dev/null +++ b/packages/core/src/ide/ide-client.ts @@ -0,0 +1,100 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { ideContext, OpenFilesNotificationSchema } from '../ide/ideContext.js'; +import { Client } from '@modelcontextprotocol/sdk/client/index.js'; +import { StreamableHTTPClientTransport } from '@modelcontextprotocol/sdk/client/streamableHttp.js'; + +const logger = { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + debug: (...args: any[]) => + console.debug('[DEBUG] [ImportProcessor]', ...args), +}; + +export type IDEConnectionState = { + status: IDEConnectionStatus; + details?: string; +}; + +export enum IDEConnectionStatus { + Connected = 'connected', + Disconnected = 'disconnected', + Connecting = 'connecting', +} + +/** + * Manages the connection to and interaction with the IDE server. + */ +export class IdeClient { + client: Client | undefined = undefined; + connectionStatus: IDEConnectionStatus = IDEConnectionStatus.Disconnected; + + constructor() { + this.connectToMcpServer().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.'; + } + } + return { + status: this.connectionStatus, + details, + }; + } + + async connectToMcpServer(): Promise { + 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.', + ); + this.connectionStatus = IDEConnectionStatus.Disconnected; + return; + } + + try { + this.client = new Client({ + name: 'streamable-http-client', + // TODO(#3487): use the CLI version here. + version: '1.0.0', + }); + const transport = new StreamableHTTPClientTransport( + new URL(`http://localhost:${idePort}/mcp`), + ); + await this.client.connect(transport); + this.client.setNotificationHandler( + OpenFilesNotificationSchema, + (notification) => { + ideContext.setOpenFilesContext(notification.params); + }, + ); + this.client.onerror = (error) => { + logger.debug('IDE MCP client error:', error); + this.connectionStatus = IDEConnectionStatus.Disconnected; + ideContext.clearOpenFilesContext(); + }; + this.client.onclose = () => { + logger.debug('IDE MCP client connection closed.'); + this.connectionStatus = IDEConnectionStatus.Disconnected; + ideContext.clearOpenFilesContext(); + }; + + this.connectionStatus = IDEConnectionStatus.Connected; + } catch (error) { + this.connectionStatus = IDEConnectionStatus.Disconnected; + logger.debug('Failed to connect to MCP server:', error); + } + } +} diff --git a/packages/core/src/ide/ideContext.test.ts b/packages/core/src/ide/ideContext.test.ts new file mode 100644 index 00000000..1cb09c53 --- /dev/null +++ b/packages/core/src/ide/ideContext.test.ts @@ -0,0 +1,140 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { describe, it, expect, beforeEach, vi } from 'vitest'; +import { createIdeContextStore } from './ideContext.js'; + +describe('ideContext - Active File', () => { + let ideContext: ReturnType; + + beforeEach(() => { + // Create a fresh, isolated instance for each test + ideContext = createIdeContextStore(); + }); + + it('should return undefined initially for active file context', () => { + expect(ideContext.getOpenFilesContext()).toBeUndefined(); + }); + + it('should set and retrieve the active file context', () => { + const testFile = { + activeFile: '/path/to/test/file.ts', + selectedText: '1234', + }; + + ideContext.setOpenFilesContext(testFile); + + const activeFile = ideContext.getOpenFilesContext(); + expect(activeFile).toEqual(testFile); + }); + + it('should update the active file context when called multiple times', () => { + const firstFile = { + activeFile: '/path/to/first.js', + selectedText: '1234', + }; + ideContext.setOpenFilesContext(firstFile); + + const secondFile = { + activeFile: '/path/to/second.py', + cursor: { line: 20, character: 30 }, + }; + ideContext.setOpenFilesContext(secondFile); + + const activeFile = ideContext.getOpenFilesContext(); + expect(activeFile).toEqual(secondFile); + }); + + it('should handle empty string for file path', () => { + const testFile = { + activeFile: '', + selectedText: '1234', + }; + ideContext.setOpenFilesContext(testFile); + expect(ideContext.getOpenFilesContext()).toEqual(testFile); + }); + + it('should notify subscribers when active file context changes', () => { + const subscriber1 = vi.fn(); + const subscriber2 = vi.fn(); + + ideContext.subscribeToOpenFiles(subscriber1); + ideContext.subscribeToOpenFiles(subscriber2); + + const testFile = { + activeFile: '/path/to/subscribed.ts', + cursor: { line: 15, character: 25 }, + }; + ideContext.setOpenFilesContext(testFile); + + expect(subscriber1).toHaveBeenCalledTimes(1); + expect(subscriber1).toHaveBeenCalledWith(testFile); + expect(subscriber2).toHaveBeenCalledTimes(1); + expect(subscriber2).toHaveBeenCalledWith(testFile); + + // Test with another update + const newFile = { + activeFile: '/path/to/new.js', + selectedText: '1234', + }; + ideContext.setOpenFilesContext(newFile); + + expect(subscriber1).toHaveBeenCalledTimes(2); + expect(subscriber1).toHaveBeenCalledWith(newFile); + expect(subscriber2).toHaveBeenCalledTimes(2); + expect(subscriber2).toHaveBeenCalledWith(newFile); + }); + + it('should stop notifying a subscriber after unsubscribe', () => { + const subscriber1 = vi.fn(); + const subscriber2 = vi.fn(); + + const unsubscribe1 = ideContext.subscribeToOpenFiles(subscriber1); + ideContext.subscribeToOpenFiles(subscriber2); + + ideContext.setOpenFilesContext({ + activeFile: '/path/to/file1.txt', + selectedText: '1234', + }); + expect(subscriber1).toHaveBeenCalledTimes(1); + expect(subscriber2).toHaveBeenCalledTimes(1); + + unsubscribe1(); + + ideContext.setOpenFilesContext({ + activeFile: '/path/to/file2.txt', + selectedText: '1234', + }); + expect(subscriber1).toHaveBeenCalledTimes(1); // Should not be called again + expect(subscriber2).toHaveBeenCalledTimes(2); + }); + + it('should allow the cursor to be optional', () => { + const testFile = { + activeFile: '/path/to/test/file.ts', + }; + + ideContext.setOpenFilesContext(testFile); + + const activeFile = ideContext.getOpenFilesContext(); + expect(activeFile).toEqual(testFile); + }); + + it('should clear the active file context', () => { + const testFile = { + activeFile: '/path/to/test/file.ts', + selectedText: '1234', + }; + + ideContext.setOpenFilesContext(testFile); + + expect(ideContext.getOpenFilesContext()).toEqual(testFile); + + ideContext.clearOpenFilesContext(); + + expect(ideContext.getOpenFilesContext()).toBeUndefined(); + }); +}); diff --git a/packages/core/src/ide/ideContext.ts b/packages/core/src/ide/ideContext.ts new file mode 100644 index 00000000..bc7383a1 --- /dev/null +++ b/packages/core/src/ide/ideContext.ts @@ -0,0 +1,118 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { z } from 'zod'; + +/** + * Zod schema for validating a cursor position. + */ +export const CursorSchema = z.object({ + line: z.number(), + character: z.number(), +}); +export type Cursor = z.infer; + +/** + * Zod schema for validating an active file context from the IDE. + */ +export const OpenFilesSchema = z.object({ + activeFile: z.string(), + selectedText: z.string().optional(), + cursor: CursorSchema.optional(), + recentOpenFiles: z + .array( + z.object({ + filePath: z.string(), + timestamp: z.number(), + }), + ) + .optional(), +}); +export type OpenFiles = z.infer; + +/** + * Zod schema for validating the 'ide/openFilesChanged' notification from the IDE. + */ +export const OpenFilesNotificationSchema = z.object({ + method: z.literal('ide/openFilesChanged'), + params: OpenFilesSchema, +}); + +type OpenFilesSubscriber = (openFiles: OpenFiles | undefined) => void; + +/** + * Creates a new store for managing the IDE's active file context. + * This factory function encapsulates the state and logic, allowing for the creation + * of isolated instances, which is particularly useful for testing. + * + * @returns An object with methods to interact with the active file context. + */ +export function createIdeContextStore() { + let openFilesContext: OpenFiles | undefined = undefined; + const subscribers = new Set(); + + /** + * Notifies all registered subscribers about the current active file context. + */ + function notifySubscribers(): void { + for (const subscriber of subscribers) { + subscriber(openFilesContext); + } + } + + /** + * Sets the active file context and notifies all registered subscribers of the change. + * @param newOpenFiles The new active file context from the IDE. + */ + function setOpenFilesContext(newOpenFiles: OpenFiles): void { + openFilesContext = newOpenFiles; + notifySubscribers(); + } + + /** + * Clears the active file context and notifies all registered subscribers of the change. + */ + function clearOpenFilesContext(): void { + openFilesContext = undefined; + notifySubscribers(); + } + + /** + * Retrieves the current active file context. + * @returns The `OpenFiles` object if a file is active; otherwise, `undefined`. + */ + function getOpenFilesContext(): OpenFiles | undefined { + return openFilesContext; + } + + /** + * Subscribes to changes in the active file context. + * + * When the active file context changes, the provided `subscriber` function will be called. + * Note: The subscriber is not called with the current value upon subscription. + * + * @param subscriber The function to be called when the active file context changes. + * @returns A function that, when called, will unsubscribe the provided subscriber. + */ + function subscribeToOpenFiles(subscriber: OpenFilesSubscriber): () => void { + subscribers.add(subscriber); + return () => { + subscribers.delete(subscriber); + }; + } + + return { + setOpenFilesContext, + getOpenFilesContext, + subscribeToOpenFiles, + clearOpenFilesContext, + }; +} + +/** + * The default, shared instance of the IDE context store for the application. + */ +export const ideContext = createIdeContextStore(); -- cgit v1.2.3