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/cli/src/config/config.test.ts | 96 ---------------- packages/cli/src/config/config.ts | 40 ++----- packages/cli/src/ui/commands/ideCommand.test.ts | 64 +++++------ packages/cli/src/ui/commands/ideCommand.ts | 48 +++----- packages/core/package.json | 4 +- packages/core/src/config/config.ts | 8 ++ packages/core/src/core/client.test.ts | 4 +- packages/core/src/core/client.ts | 2 +- packages/core/src/ide/ide-client.ts | 100 +++++++++++++++++ packages/core/src/ide/ideContext.test.ts | 140 ++++++++++++++++++++++++ packages/core/src/ide/ideContext.ts | 118 ++++++++++++++++++++ packages/core/src/index.ts | 5 +- packages/core/src/services/ideContext.test.ts | 140 ------------------------ packages/core/src/services/ideContext.ts | 122 --------------------- packages/core/src/tools/mcp-client.ts | 18 --- packages/vscode-ide-companion/src/ide-server.ts | 29 ----- 16 files changed, 430 insertions(+), 508 deletions(-) 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 delete mode 100644 packages/core/src/services/ideContext.test.ts delete mode 100644 packages/core/src/services/ideContext.ts (limited to 'packages') diff --git a/packages/cli/src/config/config.test.ts b/packages/cli/src/config/config.test.ts index c0e9c215..55780320 100644 --- a/packages/cli/src/config/config.test.ts +++ b/packages/cli/src/config/config.test.ts @@ -1011,100 +1011,4 @@ describe('loadCliConfig ideMode', () => { const config = await loadCliConfig(settings, [], 'test-session', argv); expect(config.getIdeMode()).toBe(false); }); - - it('should add _ide_server when ideMode is true', async () => { - process.argv = ['node', 'script.js', '--ide-mode']; - const argv = await parseArguments(); - process.env.TERM_PROGRAM = 'vscode'; - process.env.GEMINI_CLI_IDE_SERVER_PORT = '3000'; - const settings: Settings = {}; - const config = await loadCliConfig(settings, [], 'test-session', argv); - expect(config.getIdeMode()).toBe(true); - const mcpServers = config.getMcpServers(); - expect(mcpServers['_ide_server']).toBeDefined(); - expect(mcpServers['_ide_server'].httpUrl).toBe('http://localhost:3000/mcp'); - expect(mcpServers['_ide_server'].description).toBe('IDE connection'); - expect(mcpServers['_ide_server'].trust).toBe(false); - }); - - it('should warn if ideMode is true and no port is set', async () => { - const consoleWarnSpy = vi - .spyOn(console, 'warn') - .mockImplementation(() => {}); - process.argv = ['node', 'script.js', '--ide-mode']; - const argv = await parseArguments(); - process.env.TERM_PROGRAM = 'vscode'; - const settings: Settings = {}; - await loadCliConfig(settings, [], 'test-session', argv); - expect(consoleWarnSpy).toHaveBeenCalledWith( - '[WARN]', - 'Could not connect to IDE. Make sure you have the companion VS Code extension installed from the marketplace or via /ide install.', - ); - consoleWarnSpy.mockRestore(); - }); - - it('should warn and overwrite if settings contain the reserved _ide_server name and ideMode is active', async () => { - const consoleWarnSpy = vi - .spyOn(console, 'warn') - .mockImplementation(() => {}); - - process.argv = ['node', 'script.js', '--ide-mode']; - const argv = await parseArguments(); - process.env.TERM_PROGRAM = 'vscode'; - process.env.GEMINI_CLI_IDE_SERVER_PORT = '3000'; - const settings: Settings = { - mcpServers: { - _ide_server: new ServerConfig.MCPServerConfig( - undefined, - undefined, - undefined, - undefined, - 'http://malicious:1234', - ), - }, - }; - - const config = await loadCliConfig(settings, [], 'test-session', argv); - - expect(consoleWarnSpy).toHaveBeenCalledWith( - '[WARN]', - 'Ignoring user-defined MCP server config for "_ide_server" as it is a reserved name.', - ); - - const mcpServers = config.getMcpServers(); - expect(mcpServers['_ide_server']).toBeDefined(); - expect(mcpServers['_ide_server'].httpUrl).toBe('http://localhost:3000/mcp'); - - consoleWarnSpy.mockRestore(); - }); - - it('should NOT warn if settings contain the reserved _ide_server name and ideMode is NOT active', async () => { - const consoleWarnSpy = vi - .spyOn(console, 'warn') - .mockImplementation(() => {}); - - process.argv = ['node', 'script.js']; - const argv = await parseArguments(); - const settings: Settings = { - mcpServers: { - _ide_server: new ServerConfig.MCPServerConfig( - undefined, - undefined, - undefined, - undefined, - 'http://malicious:1234', - ), - }, - }; - - const config = await loadCliConfig(settings, [], 'test-session', argv); - - expect(consoleWarnSpy).not.toHaveBeenCalled(); - - const mcpServers = config.getMcpServers(); - expect(mcpServers['_ide_server']).toBeDefined(); - expect(mcpServers['_ide_server'].url).toBe('http://malicious:1234'); - - consoleWarnSpy.mockRestore(); - }); }); diff --git a/packages/cli/src/config/config.ts b/packages/cli/src/config/config.ts index 650f3aa2..27e3ec09 100644 --- a/packages/cli/src/config/config.ts +++ b/packages/cli/src/config/config.ts @@ -19,8 +19,7 @@ import { FileDiscoveryService, TelemetryTarget, FileFilteringOptions, - MCPServerConfig, - IDE_SERVER_NAME, + IdeClient, } from '@google/gemini-cli-core'; import { Settings } from './settings.js'; @@ -264,6 +263,11 @@ export async function loadCliConfig( process.env.TERM_PROGRAM === 'vscode' && !process.env.SANDBOX; + let ideClient: IdeClient | undefined; + if (ideMode) { + ideClient = new IdeClient(); + } + const allExtensions = annotateActiveExtensions( extensions, argv.extensions || [], @@ -355,37 +359,6 @@ export async function loadCliConfig( } } - if (ideMode) { - if (mcpServers[IDE_SERVER_NAME]) { - logger.warn( - `Ignoring user-defined MCP server config for "${IDE_SERVER_NAME}" as it is a reserved name.`, - ); - } - const companionPort = process.env.GEMINI_CLI_IDE_SERVER_PORT; - if (companionPort) { - const httpUrl = `http://localhost:${companionPort}/mcp`; - mcpServers[IDE_SERVER_NAME] = new MCPServerConfig( - undefined, // command - undefined, // args - undefined, // env - undefined, // cwd - undefined, // url - httpUrl, // httpUrl - undefined, // headers - undefined, // tcp - undefined, // timeout - false, // trust - 'IDE connection', // description - undefined, // includeTools - undefined, // excludeTools - ); - } else { - logger.warn( - 'Could not connect to IDE. Make sure you have the companion VS Code extension installed from the marketplace or via /ide install.', - ); - } - } - const sandboxConfig = await loadSandboxConfig(settings, argv); return new Config({ @@ -450,6 +423,7 @@ export async function loadCliConfig( noBrowser: !!process.env.NO_BROWSER, summarizeToolOutput: settings.summarizeToolOutput, ideMode, + ideClient, }); } diff --git a/packages/cli/src/ui/commands/ideCommand.test.ts b/packages/cli/src/ui/commands/ideCommand.test.ts index 51322843..d1d72466 100644 --- a/packages/cli/src/ui/commands/ideCommand.test.ts +++ b/packages/cli/src/ui/commands/ideCommand.test.ts @@ -19,25 +19,10 @@ import { type Config } from '@google/gemini-cli-core'; import * as child_process from 'child_process'; import { glob } from 'glob'; -import { - getMCPDiscoveryState, - getMCPServerStatus, - IDE_SERVER_NAME, - MCPDiscoveryState, - MCPServerStatus, -} from '@google/gemini-cli-core'; +import { IDEConnectionStatus } from '@google/gemini-cli-core/index.js'; vi.mock('child_process'); vi.mock('glob'); -vi.mock('@google/gemini-cli-core', async (importOriginal) => { - const original = - await importOriginal(); - return { - ...original, - getMCPServerStatus: vi.fn(), - getMCPDiscoveryState: vi.fn(), - }; -}); function regexEscape(value: string) { return value.replace(/[.*+?^${}()|[\]\\]/g, '\\$&'); @@ -49,8 +34,6 @@ describe('ideCommand', () => { let execSyncSpy: MockInstance; let globSyncSpy: MockInstance; let platformSpy: MockInstance; - let getMCPServerStatusSpy: MockInstance; - let getMCPDiscoveryStateSpy: MockInstance; beforeEach(() => { mockContext = { @@ -61,13 +44,12 @@ describe('ideCommand', () => { mockConfig = { getIdeMode: vi.fn(), + getIdeClient: vi.fn(), } as unknown as Config; execSyncSpy = vi.spyOn(child_process, 'execSync'); globSyncSpy = vi.spyOn(glob, 'sync'); platformSpy = vi.spyOn(process, 'platform', 'get'); - getMCPServerStatusSpy = vi.mocked(getMCPServerStatus); - getMCPDiscoveryStateSpy = vi.mocked(getMCPDiscoveryState); }); afterEach(() => { @@ -91,15 +73,21 @@ describe('ideCommand', () => { }); describe('status subcommand', () => { + const mockGetConnectionStatus = vi.fn(); beforeEach(() => { vi.mocked(mockConfig.getIdeMode).mockReturnValue(true); + vi.mocked(mockConfig.getIdeClient).mockReturnValue({ + getConnectionStatus: mockGetConnectionStatus, + } as ReturnType); }); it('should show connected status', () => { - getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.CONNECTED); + mockGetConnectionStatus.mockReturnValue({ + status: IDEConnectionStatus.Connected, + }); const command = ideCommand(mockConfig); const result = command!.subCommands![0].action!(mockContext, ''); - expect(getMCPServerStatusSpy).toHaveBeenCalledWith(IDE_SERVER_NAME); + expect(mockGetConnectionStatus).toHaveBeenCalled(); expect(result).toEqual({ type: 'message', messageType: 'info', @@ -108,37 +96,45 @@ describe('ideCommand', () => { }); it('should show connecting status', () => { - getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.CONNECTING); + mockGetConnectionStatus.mockReturnValue({ + status: IDEConnectionStatus.Connecting, + }); const command = ideCommand(mockConfig); const result = command!.subCommands![0].action!(mockContext, ''); + expect(mockGetConnectionStatus).toHaveBeenCalled(); expect(result).toEqual({ type: 'message', messageType: 'info', - content: '🔄 Initializing...', + content: `🟡 Connecting...`, }); }); - - it('should show discovery in progress status', () => { - getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.DISCONNECTED); - getMCPDiscoveryStateSpy.mockReturnValue(MCPDiscoveryState.IN_PROGRESS); + it('should show disconnected status', () => { + mockGetConnectionStatus.mockReturnValue({ + status: IDEConnectionStatus.Disconnected, + }); const command = ideCommand(mockConfig); const result = command!.subCommands![0].action!(mockContext, ''); + expect(mockGetConnectionStatus).toHaveBeenCalled(); expect(result).toEqual({ type: 'message', - messageType: 'info', - content: '🔄 Initializing...', + messageType: 'error', + content: `🔴 Disconnected`, }); }); - it('should show disconnected status', () => { - getMCPServerStatusSpy.mockReturnValue(MCPServerStatus.DISCONNECTED); - getMCPDiscoveryStateSpy.mockReturnValue(MCPDiscoveryState.COMPLETED); + it('should show disconnected status with details', () => { + const details = 'Something went wrong'; + mockGetConnectionStatus.mockReturnValue({ + status: IDEConnectionStatus.Disconnected, + details, + }); const command = ideCommand(mockConfig); const result = command!.subCommands![0].action!(mockContext, ''); + expect(mockGetConnectionStatus).toHaveBeenCalled(); expect(result).toEqual({ type: 'message', messageType: 'error', - content: '🔴 Disconnected', + content: `🔴 Disconnected: ${details}`, }); }); }); diff --git a/packages/cli/src/ui/commands/ideCommand.ts b/packages/cli/src/ui/commands/ideCommand.ts index 6fc4f50b..31f2371f 100644 --- a/packages/cli/src/ui/commands/ideCommand.ts +++ b/packages/cli/src/ui/commands/ideCommand.ts @@ -5,14 +5,7 @@ */ import { fileURLToPath } from 'url'; -import { - Config, - getMCPDiscoveryState, - getMCPServerStatus, - IDE_SERVER_NAME, - MCPDiscoveryState, - MCPServerStatus, -} from '@google/gemini-cli-core'; +import { Config, IDEConnectionStatus } from '@google/gemini-cli-core'; import { CommandContext, SlashCommand, @@ -56,36 +49,31 @@ export const ideCommand = (config: Config | null): SlashCommand | null => { description: 'check status of IDE integration', kind: CommandKind.BUILT_IN, action: (_context: CommandContext): SlashCommandActionReturn => { - const status = getMCPServerStatus(IDE_SERVER_NAME); - const discoveryState = getMCPDiscoveryState(); - switch (status) { - case MCPServerStatus.CONNECTED: + const connection = config.getIdeClient()?.getConnectionStatus(); + switch (connection?.status) { + case IDEConnectionStatus.Connected: return { type: 'message', messageType: 'info', content: `🟢 Connected`, - }; - case MCPServerStatus.CONNECTING: + } as const; + case IDEConnectionStatus.Connecting: return { type: 'message', messageType: 'info', - content: `🔄 Initializing...`, - }; - case MCPServerStatus.DISCONNECTED: - default: - if (discoveryState === MCPDiscoveryState.IN_PROGRESS) { - return { - type: 'message', - messageType: 'info', - content: `🔄 Initializing...`, - }; - } else { - return { - type: 'message', - messageType: 'error', - content: `🔴 Disconnected`, - }; + content: `🟡 Connecting...`, + } as const; + default: { + let content = `🔴 Disconnected`; + if (connection?.details) { + content += `: ${connection.details}`; } + return { + type: 'message', + messageType: 'error', + content, + } as const; + } } }, }, diff --git a/packages/core/package.json b/packages/core/package.json index 40f10aa0..ba4735ea 100644 --- a/packages/core/package.json +++ b/packages/core/package.json @@ -31,6 +31,7 @@ "@types/glob": "^8.1.0", "@types/html-to-text": "^9.0.4", "ajv": "^8.17.1", + "chardet": "^2.1.0", "diff": "^7.0.0", "dotenv": "^17.1.0", "glob": "^10.4.5", @@ -44,8 +45,7 @@ "simple-git": "^3.28.0", "strip-ansi": "^7.1.0", "undici": "^7.10.0", - "ws": "^8.18.0", - "chardet": "^2.1.0" + "ws": "^8.18.0" }, "devDependencies": { "@types/diff": "^7.0.2", diff --git a/packages/core/src/config/config.ts b/packages/core/src/config/config.ts index 485a56c4..96b6f2cb 100644 --- a/packages/core/src/config/config.ts +++ b/packages/core/src/config/config.ts @@ -45,6 +45,7 @@ import { import { ClearcutLogger } from '../telemetry/clearcut-logger/clearcut-logger.js'; import { shouldAttemptBrowserLaunch } from '../utils/browser.js'; import { MCPOAuthConfig } from '../mcp/oauth-provider.js'; +import { IdeClient } from '../ide/ide-client.js'; // Re-export OAuth config type export type { MCPOAuthConfig }; @@ -180,6 +181,7 @@ export interface ConfigParameters { noBrowser?: boolean; summarizeToolOutput?: Record; ideMode?: boolean; + ideClient?: IdeClient; } export class Config { @@ -221,6 +223,7 @@ export class Config { private readonly extensionContextFilePaths: string[]; private readonly noBrowser: boolean; private readonly ideMode: boolean; + private readonly ideClient: IdeClient | undefined; private modelSwitchedDuringSession: boolean = false; private readonly maxSessionTurns: number; private readonly listExtensions: boolean; @@ -286,6 +289,7 @@ export class Config { this.noBrowser = params.noBrowser ?? false; this.summarizeToolOutput = params.summarizeToolOutput; this.ideMode = params.ideMode ?? false; + this.ideClient = params.ideClient; if (params.contextFileName) { setGeminiMdFilename(params.contextFileName); @@ -574,6 +578,10 @@ export class Config { return this.ideMode; } + getIdeClient(): IdeClient | undefined { + return this.ideClient; + } + async getGitService(): Promise { if (!this.gitService) { this.gitService = new GitService(this.targetDir); diff --git a/packages/core/src/core/client.test.ts b/packages/core/src/core/client.test.ts index 44b19f56..25ea9bc1 100644 --- a/packages/core/src/core/client.test.ts +++ b/packages/core/src/core/client.test.ts @@ -23,7 +23,7 @@ import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js'; import { FileDiscoveryService } from '../services/fileDiscoveryService.js'; import { setSimulate429 } from '../utils/testUtils.js'; import { tokenLimit } from './tokenLimits.js'; -import { ideContext } from '../services/ideContext.js'; +import { ideContext } from '../ide/ideContext.js'; // --- Mocks --- const mockChatCreateFn = vi.fn(); @@ -72,7 +72,7 @@ vi.mock('../telemetry/index.js', () => ({ logApiResponse: vi.fn(), logApiError: vi.fn(), })); -vi.mock('../services/ideContext.js'); +vi.mock('../ide/ideContext.js'); describe('findIndexAfterFraction', () => { const history: Content[] = [ diff --git a/packages/core/src/core/client.ts b/packages/core/src/core/client.ts index 6f482307..77683a45 100644 --- a/packages/core/src/core/client.ts +++ b/packages/core/src/core/client.ts @@ -42,7 +42,7 @@ import { import { ProxyAgent, setGlobalDispatcher } from 'undici'; import { DEFAULT_GEMINI_FLASH_MODEL } from '../config/models.js'; import { LoopDetectionService } from '../services/loopDetectionService.js'; -import { ideContext } from '../services/ideContext.js'; +import { ideContext } from '../ide/ideContext.js'; import { logFlashDecidedToContinue } from '../telemetry/loggers.js'; import { FlashDecidedToContinueEvent } from '../telemetry/types.js'; 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(); diff --git a/packages/core/src/index.ts b/packages/core/src/index.ts index f560afb4..9d87ce32 100644 --- a/packages/core/src/index.ts +++ b/packages/core/src/index.ts @@ -40,7 +40,10 @@ export * from './utils/systemEncoding.js'; // Export services export * from './services/fileDiscoveryService.js'; export * from './services/gitService.js'; -export * from './services/ideContext.js'; + +// Export IDE specific logic +export * from './ide/ide-client.js'; +export * from './ide/ideContext.js'; // Export base tool definitions export * from './tools/tools.js'; diff --git a/packages/core/src/services/ideContext.test.ts b/packages/core/src/services/ideContext.test.ts deleted file mode 100644 index 1cb09c53..00000000 --- a/packages/core/src/services/ideContext.test.ts +++ /dev/null @@ -1,140 +0,0 @@ -/** - * @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/services/ideContext.ts b/packages/core/src/services/ideContext.ts deleted file mode 100644 index f8a50f12..00000000 --- a/packages/core/src/services/ideContext.ts +++ /dev/null @@ -1,122 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { z } from 'zod'; - -/** - * The reserved server name for the IDE's MCP server. - */ -export const IDE_SERVER_NAME = '_ide_server'; -/** - * 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(); diff --git a/packages/core/src/tools/mcp-client.ts b/packages/core/src/tools/mcp-client.ts index 3c482100..c59b1592 100644 --- a/packages/core/src/tools/mcp-client.ts +++ b/packages/core/src/tools/mcp-client.ts @@ -24,11 +24,6 @@ import { ToolRegistry } from './tool-registry.js'; import { MCPOAuthProvider } from '../mcp/oauth-provider.js'; import { OAuthUtils } from '../mcp/oauth-utils.js'; import { MCPOAuthTokenStorage } from '../mcp/oauth-token-storage.js'; -import { - OpenFilesNotificationSchema, - IDE_SERVER_NAME, - ideContext, -} from '../services/ideContext.js'; import { getErrorMessage } from '../utils/errors.js'; export const MCP_DEFAULT_TIMEOUT_MSEC = 10 * 60 * 1000; // default to 10 minutes @@ -379,24 +374,11 @@ export async function connectAndDiscover( ); try { updateMCPServerStatus(mcpServerName, MCPServerStatus.CONNECTED); - mcpClient.onerror = (error) => { console.error(`MCP ERROR (${mcpServerName}):`, error.toString()); updateMCPServerStatus(mcpServerName, MCPServerStatus.DISCONNECTED); - if (mcpServerName === IDE_SERVER_NAME) { - ideContext.clearOpenFilesContext(); - } }; - if (mcpServerName === IDE_SERVER_NAME) { - mcpClient.setNotificationHandler( - OpenFilesNotificationSchema, - (notification) => { - ideContext.setOpenFilesContext(notification.params); - }, - ); - } - const tools = await discoverTools( mcpServerName, mcpServerConfig, diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index 3072029b..f47463ba 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/packages/vscode-ide-companion/src/ide-server.ts @@ -244,34 +244,5 @@ const createMcpServer = () => { }, { capabilities: { logging: {} } }, ); - server.registerTool( - 'getOpenFiles', - { - description: - '(IDE Tool) Get the path of the file currently active in VS Code.', - inputSchema: {}, - }, - async () => { - const editor = vscode.window.activeTextEditor; - const filePath = - editor && editor.document.uri.scheme === 'file' - ? editor.document.uri.fsPath - : ''; - if (filePath) { - return { - content: [{ type: 'text', text: `Active file: ${filePath}` }], - }; - } else { - return { - content: [ - { - type: 'text', - text: 'No file is currently active in the editor.', - }, - ], - }; - } - }, - ); return server; }; -- cgit v1.2.3