diff --git a/packages/vscode-ide-companion/src/extension-multi-folder.test.ts b/packages/vscode-ide-companion/src/extension-multi-folder.test.ts deleted file mode 100644 index c8fff810..00000000 --- a/packages/vscode-ide-companion/src/extension-multi-folder.test.ts +++ /dev/null @@ -1,211 +0,0 @@ -/** - * @license - * Copyright 2025 Google LLC - * SPDX-License-Identifier: Apache-2.0 - */ - -import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; -import * as vscode from 'vscode'; -import * as path from 'path'; -import { activate } from './extension.js'; - -vi.mock('vscode', () => ({ - window: { - createOutputChannel: vi.fn(() => ({ - appendLine: vi.fn(), - })), - showInformationMessage: vi.fn(), - createTerminal: vi.fn(() => ({ - show: vi.fn(), - sendText: vi.fn(), - })), - onDidChangeActiveTextEditor: vi.fn(), - activeTextEditor: undefined, - tabGroups: { - all: [], - close: vi.fn(), - }, - showTextDocument: vi.fn(), - }, - workspace: { - workspaceFolders: [], - onDidCloseTextDocument: vi.fn(), - registerTextDocumentContentProvider: vi.fn(), - onDidChangeWorkspaceFolders: vi.fn(), - }, - commands: { - registerCommand: vi.fn(), - executeCommand: vi.fn(), - }, - Uri: { - joinPath: vi.fn(), - file: (path: string) => ({ fsPath: path }), - }, - ExtensionMode: { - Development: 1, - Production: 2, - }, - EventEmitter: vi.fn(() => ({ - event: vi.fn(), - fire: vi.fn(), - dispose: vi.fn(), - })), -})); - -describe('activate with multiple folders', () => { - let context: vscode.ExtensionContext; - let onDidChangeWorkspaceFoldersCallback: ( - e: vscode.WorkspaceFoldersChangeEvent, - ) => void; - - beforeEach(() => { - context = { - subscriptions: [], - environmentVariableCollection: { - replace: vi.fn(), - }, - globalState: { - get: vi.fn().mockReturnValue(true), - update: vi.fn(), - }, - extensionUri: { - fsPath: '/path/to/extension', - }, - } as unknown as vscode.ExtensionContext; - - vi.mocked(vscode.workspace.onDidChangeWorkspaceFolders).mockImplementation( - (callback) => { - onDidChangeWorkspaceFoldersCallback = callback; - return { dispose: vi.fn() }; - }, - ); - }); - - afterEach(() => { - vi.restoreAllMocks(); - }); - - it('should set a single folder path', async () => { - const workspaceFoldersSpy = vi.spyOn( - vscode.workspace, - 'workspaceFolders', - 'get', - ); - workspaceFoldersSpy.mockReturnValue([ - { uri: { fsPath: '/foo/bar' } }, - ] as vscode.WorkspaceFolder[]); - - await activate(context); - - expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith( - 'GEMINI_CLI_IDE_WORKSPACE_PATH', - '/foo/bar', - ); - }); - - it('should set multiple folder paths, separated by OS-specific path delimiter', async () => { - const workspaceFoldersSpy = vi.spyOn( - vscode.workspace, - 'workspaceFolders', - 'get', - ); - workspaceFoldersSpy.mockReturnValue([ - { uri: { fsPath: '/foo/bar' } }, - { uri: { fsPath: '/baz/qux' } }, - ] as vscode.WorkspaceFolder[]); - - await activate(context); - - expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith( - 'GEMINI_CLI_IDE_WORKSPACE_PATH', - ['/foo/bar', '/baz/qux'].join(path.delimiter), - ); - }); - - it('should set an empty string if no folders are open', async () => { - const workspaceFoldersSpy = vi.spyOn( - vscode.workspace, - 'workspaceFolders', - 'get', - ); - workspaceFoldersSpy.mockReturnValue([]); - - await activate(context); - - expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith( - 'GEMINI_CLI_IDE_WORKSPACE_PATH', - '', - ); - }); - - it('should update the path when workspace folders change', async () => { - const workspaceFoldersSpy = vi.spyOn( - vscode.workspace, - 'workspaceFolders', - 'get', - ); - workspaceFoldersSpy.mockReturnValue([ - { uri: { fsPath: '/foo/bar' } }, - ] as vscode.WorkspaceFolder[]); - - await activate(context); - - expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith( - 'GEMINI_CLI_IDE_WORKSPACE_PATH', - '/foo/bar', - ); - - // Simulate adding a folder - workspaceFoldersSpy.mockReturnValue([ - { uri: { fsPath: '/foo/bar' } }, - { uri: { fsPath: '/baz/qux' } }, - ] as vscode.WorkspaceFolder[]); - onDidChangeWorkspaceFoldersCallback({ - added: [{ uri: { fsPath: '/baz/qux' } } as vscode.WorkspaceFolder], - removed: [], - }); - - expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith( - 'GEMINI_CLI_IDE_WORKSPACE_PATH', - ['/foo/bar', '/baz/qux'].join(path.delimiter), - ); - - // Simulate removing a folder - workspaceFoldersSpy.mockReturnValue([ - { uri: { fsPath: '/baz/qux' } }, - ] as vscode.WorkspaceFolder[]); - onDidChangeWorkspaceFoldersCallback({ - added: [], - removed: [{ uri: { fsPath: '/foo/bar' } } as vscode.WorkspaceFolder], - }); - - expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith( - 'GEMINI_CLI_IDE_WORKSPACE_PATH', - '/baz/qux', - ); - }); - - it.skipIf(process.platform !== 'win32')( - 'should handle windows paths', - async () => { - const workspaceFoldersSpy = vi.spyOn( - vscode.workspace, - 'workspaceFolders', - 'get', - ); - workspaceFoldersSpy.mockReturnValue([ - { uri: { fsPath: 'c:/foo/bar' } }, - { uri: { fsPath: 'd:/baz/qux' } }, - ] as vscode.WorkspaceFolder[]); - - await activate(context); - - expect( - context.environmentVariableCollection.replace, - ).toHaveBeenCalledWith( - 'GEMINI_CLI_IDE_WORKSPACE_PATH', - 'c:/foo/bar;d:/baz/qux', - ); - }, - ); -}); diff --git a/packages/vscode-ide-companion/src/extension.ts b/packages/vscode-ide-companion/src/extension.ts index 30bab801..4e03a290 100644 --- a/packages/vscode-ide-companion/src/extension.ts +++ b/packages/vscode-ide-companion/src/extension.ts @@ -5,13 +5,11 @@ */ import * as vscode from 'vscode'; -import * as path from 'path'; import { IDEServer } from './ide-server.js'; import { DiffContentProvider, DiffManager } from './diff-manager.js'; import { createLogger } from './utils/logger.js'; const INFO_MESSAGE_SHOWN_KEY = 'geminiCliInfoMessageShown'; -const IDE_WORKSPACE_PATH_ENV_VAR = 'GEMINI_CLI_IDE_WORKSPACE_PATH'; export const DIFF_SCHEME = 'gemini-diff'; let ideServer: IDEServer; @@ -19,31 +17,11 @@ let logger: vscode.OutputChannel; let log: (message: string) => void = () => {}; -function updateWorkspacePath(context: vscode.ExtensionContext) { - const workspaceFolders = vscode.workspace.workspaceFolders; - if (workspaceFolders && workspaceFolders.length > 0) { - const workspacePaths = workspaceFolders - .map((folder) => folder.uri.fsPath) - .join(path.delimiter); - context.environmentVariableCollection.replace( - IDE_WORKSPACE_PATH_ENV_VAR, - workspacePaths, - ); - } else { - context.environmentVariableCollection.replace( - IDE_WORKSPACE_PATH_ENV_VAR, - '', - ); - } -} - export async function activate(context: vscode.ExtensionContext) { logger = vscode.window.createOutputChannel('Gemini CLI IDE Companion'); log = createLogger(context, logger); log('Extension activated'); - updateWorkspacePath(context); - const diffContentProvider = new DiffContentProvider(); const diffManager = new DiffManager(log, diffContentProvider); @@ -94,7 +72,7 @@ export async function activate(context: vscode.ExtensionContext) { context.subscriptions.push( vscode.workspace.onDidChangeWorkspaceFolders(() => { - updateWorkspacePath(context); + ideServer.updateWorkspacePath(); }), vscode.commands.registerCommand('gemini-cli.runGeminiCLI', async () => { const workspaceFolders = vscode.workspace.workspaceFolders; diff --git a/packages/vscode-ide-companion/src/ide-server.test.ts b/packages/vscode-ide-companion/src/ide-server.test.ts new file mode 100644 index 00000000..726259e4 --- /dev/null +++ b/packages/vscode-ide-companion/src/ide-server.test.ts @@ -0,0 +1,287 @@ +/** + * @license + * Copyright 2025 Google LLC + * SPDX-License-Identifier: Apache-2.0 + */ + +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import type * as vscode from 'vscode'; +import * as fs from 'node:fs/promises'; +import * as os from 'node:os'; +import * as path from 'node:path'; +import { IDEServer } from './ide-server.js'; +import { DiffManager } from './diff-manager.js'; + +const mocks = vi.hoisted(() => ({ + diffManager: { + onDidChange: vi.fn(() => ({ dispose: vi.fn() })), + } as unknown as DiffManager, +})); + +vi.mock('node:fs/promises', () => ({ + writeFile: vi.fn(() => Promise.resolve(undefined)), + unlink: vi.fn(() => Promise.resolve(undefined)), +})); + +vi.mock('node:os', async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + tmpdir: vi.fn(() => '/tmp'), + }; +}); + +const vscodeMock = vi.hoisted(() => ({ + workspace: { + workspaceFolders: [ + { + uri: { + fsPath: '/test/workspace1', + }, + }, + { + uri: { + fsPath: '/test/workspace2', + }, + }, + ], + }, +})); + +vi.mock('vscode', () => vscodeMock); + +vi.mock('./open-files-manager', () => { + const OpenFilesManager = vi.fn(); + OpenFilesManager.prototype.onDidChange = vi.fn(() => ({ dispose: vi.fn() })); + return { OpenFilesManager }; +}); + +describe('IDEServer', () => { + let ideServer: IDEServer; + let mockContext: vscode.ExtensionContext; + let mockLog: (message: string) => void; + + const getPortFromMock = ( + replaceMock: ReturnType< + () => vscode.ExtensionContext['environmentVariableCollection']['replace'] + >, + ) => { + const port = vi + .mocked(replaceMock) + .mock.calls.find((call) => call[0] === 'GEMINI_CLI_IDE_SERVER_PORT')?.[1]; + + if (port === undefined) { + expect.fail('Port was not set'); + } + return port; + }; + + beforeEach(() => { + mockLog = vi.fn(); + ideServer = new IDEServer(mockLog, mocks.diffManager); + mockContext = { + subscriptions: [], + environmentVariableCollection: { + replace: vi.fn(), + clear: vi.fn(), + }, + } as unknown as vscode.ExtensionContext; + }); + + afterEach(async () => { + await ideServer.stop(); + vi.restoreAllMocks(); + vscodeMock.workspace.workspaceFolders = [ + { uri: { fsPath: '/test/workspace1' } }, + { uri: { fsPath: '/test/workspace2' } }, + ]; + }); + + it('should set environment variables and workspace path on start with multiple folders', async () => { + await ideServer.start(mockContext); + + const replaceMock = mockContext.environmentVariableCollection.replace; + expect(replaceMock).toHaveBeenCalledTimes(2); + + expect(replaceMock).toHaveBeenNthCalledWith( + 1, + 'GEMINI_CLI_IDE_SERVER_PORT', + expect.any(String), // port is a number as a string + ); + + const expectedWorkspacePaths = [ + '/test/workspace1', + '/test/workspace2', + ].join(path.delimiter); + + expect(replaceMock).toHaveBeenNthCalledWith( + 2, + 'GEMINI_CLI_IDE_WORKSPACE_PATH', + expectedWorkspacePaths, + ); + + const port = getPortFromMock(replaceMock); + const expectedPortFile = path.join( + '/tmp', + `gemini-ide-server-${process.ppid}.json`, + ); + expect(fs.writeFile).toHaveBeenCalledWith( + expectedPortFile, + JSON.stringify({ + port: parseInt(port, 10), + workspacePath: expectedWorkspacePaths, + }), + ); + }); + + it('should set a single folder path', async () => { + vscodeMock.workspace.workspaceFolders = [{ uri: { fsPath: '/foo/bar' } }]; + + await ideServer.start(mockContext); + const replaceMock = mockContext.environmentVariableCollection.replace; + + expect(replaceMock).toHaveBeenCalledWith( + 'GEMINI_CLI_IDE_WORKSPACE_PATH', + '/foo/bar', + ); + + const port = getPortFromMock(replaceMock); + const expectedPortFile = path.join( + '/tmp', + `gemini-ide-server-${process.ppid}.json`, + ); + expect(fs.writeFile).toHaveBeenCalledWith( + expectedPortFile, + JSON.stringify({ + port: parseInt(port, 10), + workspacePath: '/foo/bar', + }), + ); + }); + + it('should set an empty string if no folders are open', async () => { + vscodeMock.workspace.workspaceFolders = []; + + await ideServer.start(mockContext); + const replaceMock = mockContext.environmentVariableCollection.replace; + + expect(replaceMock).toHaveBeenCalledWith( + 'GEMINI_CLI_IDE_WORKSPACE_PATH', + '', + ); + + const port = getPortFromMock(replaceMock); + const expectedPortFile = path.join( + '/tmp', + `gemini-ide-server-${process.ppid}.json`, + ); + expect(fs.writeFile).toHaveBeenCalledWith( + expectedPortFile, + JSON.stringify({ + port: parseInt(port, 10), + workspacePath: '', + }), + ); + }); + + it('should update the path when workspace folders change', async () => { + vscodeMock.workspace.workspaceFolders = [{ uri: { fsPath: '/foo/bar' } }]; + await ideServer.start(mockContext); + const replaceMock = mockContext.environmentVariableCollection.replace; + + expect(replaceMock).toHaveBeenCalledWith( + 'GEMINI_CLI_IDE_WORKSPACE_PATH', + '/foo/bar', + ); + + // Simulate adding a folder + vscodeMock.workspace.workspaceFolders = [ + { uri: { fsPath: '/foo/bar' } }, + { uri: { fsPath: '/baz/qux' } }, + ]; + await ideServer.updateWorkspacePath(); + + const expectedWorkspacePaths = ['/foo/bar', '/baz/qux'].join( + path.delimiter, + ); + expect(replaceMock).toHaveBeenCalledWith( + 'GEMINI_CLI_IDE_WORKSPACE_PATH', + expectedWorkspacePaths, + ); + + const port = getPortFromMock(replaceMock); + const expectedPortFile = path.join( + '/tmp', + `gemini-ide-server-${process.ppid}.json`, + ); + expect(fs.writeFile).toHaveBeenCalledWith( + expectedPortFile, + JSON.stringify({ + port: parseInt(port, 10), + workspacePath: expectedWorkspacePaths, + }), + ); + + // Simulate removing a folder + vscodeMock.workspace.workspaceFolders = [{ uri: { fsPath: '/baz/qux' } }]; + await ideServer.updateWorkspacePath(); + + expect(replaceMock).toHaveBeenCalledWith( + 'GEMINI_CLI_IDE_WORKSPACE_PATH', + '/baz/qux', + ); + expect(fs.writeFile).toHaveBeenCalledWith( + expectedPortFile, + JSON.stringify({ + port: parseInt(port, 10), + workspacePath: '/baz/qux', + }), + ); + }); + + it('should clear env vars and delete port file on stop', async () => { + await ideServer.start(mockContext); + const portFile = path.join( + '/tmp', + `gemini-ide-server-${process.ppid}.json`, + ); + expect(fs.writeFile).toHaveBeenCalledWith(portFile, expect.any(String)); + + await ideServer.stop(); + + expect(mockContext.environmentVariableCollection.clear).toHaveBeenCalled(); + expect(fs.unlink).toHaveBeenCalledWith(portFile); + }); + + it.skipIf(process.platform !== 'win32')( + 'should handle windows paths', + async () => { + vscodeMock.workspace.workspaceFolders = [ + { uri: { fsPath: 'c:\\foo\\bar' } }, + { uri: { fsPath: 'd:\\baz\\qux' } }, + ]; + + await ideServer.start(mockContext); + const replaceMock = mockContext.environmentVariableCollection.replace; + const expectedWorkspacePaths = 'c:\\foo\\bar;d:\\baz\\qux'; + + expect(replaceMock).toHaveBeenCalledWith( + 'GEMINI_CLI_IDE_WORKSPACE_PATH', + expectedWorkspacePaths, + ); + + const port = getPortFromMock(replaceMock); + const expectedPortFile = path.join( + '/tmp', + `gemini-ide-server-${process.ppid}.json`, + ); + expect(fs.writeFile).toHaveBeenCalledWith( + expectedPortFile, + JSON.stringify({ + port: parseInt(port, 10), + workspacePath: expectedWorkspacePaths, + }), + ); + }, + ); +}); diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index ee77bdb8..2a49e74c 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/packages/vscode-ide-companion/src/ide-server.ts @@ -21,6 +21,37 @@ import { OpenFilesManager } from './open-files-manager.js'; const MCP_SESSION_ID_HEADER = 'mcp-session-id'; const IDE_SERVER_PORT_ENV_VAR = 'GEMINI_CLI_IDE_SERVER_PORT'; +const IDE_WORKSPACE_PATH_ENV_VAR = 'GEMINI_CLI_IDE_WORKSPACE_PATH'; + +function writePortAndWorkspace( + context: vscode.ExtensionContext, + port: number, + portFile: string, + log: (message: string) => void, +): Promise { + const workspaceFolders = vscode.workspace.workspaceFolders; + const workspacePath = + workspaceFolders && workspaceFolders.length > 0 + ? workspaceFolders.map((folder) => folder.uri.fsPath).join(path.delimiter) + : ''; + + context.environmentVariableCollection.replace( + IDE_SERVER_PORT_ENV_VAR, + port.toString(), + ); + context.environmentVariableCollection.replace( + IDE_WORKSPACE_PATH_ENV_VAR, + workspacePath, + ); + + log(`Writing port file to: ${portFile}`); + return fs + .writeFile(portFile, JSON.stringify({ port, workspacePath })) + .catch((err) => { + const message = err instanceof Error ? err.message : String(err); + log(`Failed to write port to file: ${message}`); + }); +} function sendIdeContextUpdateNotification( transport: StreamableHTTPServerTransport, @@ -50,6 +81,7 @@ export class IDEServer { private context: vscode.ExtensionContext | undefined; private log: (message: string) => void; private portFile: string; + private port: number | undefined; diffManager: DiffManager; constructor(log: (message: string) => void, diffManager: DiffManager) { @@ -61,158 +93,170 @@ export class IDEServer { ); } - async start(context: vscode.ExtensionContext) { - this.context = context; - const sessionsWithInitialNotification = new Set(); - const transports: { [sessionId: string]: StreamableHTTPServerTransport } = - {}; + start(context: vscode.ExtensionContext): Promise { + return new Promise((resolve) => { + this.context = context; + const sessionsWithInitialNotification = new Set(); + const transports: { [sessionId: string]: StreamableHTTPServerTransport } = + {}; - const app = express(); - app.use(express.json()); - const mcpServer = createMcpServer(this.diffManager); + const app = express(); + app.use(express.json()); + const mcpServer = createMcpServer(this.diffManager); - const openFilesManager = new OpenFilesManager(context); - const onDidChangeSubscription = openFilesManager.onDidChange(() => { - for (const transport of Object.values(transports)) { - sendIdeContextUpdateNotification( - transport, - this.log.bind(this), - openFilesManager, - ); - } - }); - context.subscriptions.push(onDidChangeSubscription); - const onDidChangeDiffSubscription = this.diffManager.onDidChange( - (notification) => { + const openFilesManager = new OpenFilesManager(context); + const onDidChangeSubscription = openFilesManager.onDidChange(() => { for (const transport of Object.values(transports)) { - transport.send(notification); + sendIdeContextUpdateNotification( + transport, + this.log.bind(this), + openFilesManager, + ); } - }, - ); - context.subscriptions.push(onDidChangeDiffSubscription); + }); + context.subscriptions.push(onDidChangeSubscription); + const onDidChangeDiffSubscription = this.diffManager.onDidChange( + (notification) => { + for (const transport of Object.values(transports)) { + transport.send(notification); + } + }, + ); + context.subscriptions.push(onDidChangeDiffSubscription); - app.post('/mcp', async (req: Request, res: Response) => { - const sessionId = req.headers[MCP_SESSION_ID_HEADER] as - | string - | undefined; - let transport: StreamableHTTPServerTransport; + app.post('/mcp', async (req: Request, res: Response) => { + const sessionId = req.headers[MCP_SESSION_ID_HEADER] as + | string + | undefined; + let transport: StreamableHTTPServerTransport; - if (sessionId && transports[sessionId]) { - transport = transports[sessionId]; - } else if (!sessionId && isInitializeRequest(req.body)) { - transport = new StreamableHTTPServerTransport({ - sessionIdGenerator: () => randomUUID(), - onsessioninitialized: (newSessionId) => { - this.log(`New session initialized: ${newSessionId}`); - transports[newSessionId] = transport; - }, - }); - const keepAlive = setInterval(() => { - try { - transport.send({ jsonrpc: '2.0', method: 'ping' }); - } catch (e) { - this.log( - 'Failed to send keep-alive ping, cleaning up interval.' + e, - ); + if (sessionId && transports[sessionId]) { + transport = transports[sessionId]; + } else if (!sessionId && isInitializeRequest(req.body)) { + transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + onsessioninitialized: (newSessionId) => { + this.log(`New session initialized: ${newSessionId}`); + transports[newSessionId] = transport; + }, + }); + const keepAlive = setInterval(() => { + try { + transport.send({ jsonrpc: '2.0', method: 'ping' }); + } catch (e) { + this.log( + 'Failed to send keep-alive ping, cleaning up interval.' + e, + ); + clearInterval(keepAlive); + } + }, 60000); // 60 sec + + transport.onclose = () => { clearInterval(keepAlive); - } - }, 60000); // 60 sec - - transport.onclose = () => { - clearInterval(keepAlive); - if (transport.sessionId) { - this.log(`Session closed: ${transport.sessionId}`); - sessionsWithInitialNotification.delete(transport.sessionId); - delete transports[transport.sessionId]; - } - }; - mcpServer.connect(transport); - } else { - this.log( - 'Bad Request: No valid session ID provided for non-initialize request.', - ); - res.status(400).json({ - jsonrpc: '2.0', - error: { - code: -32000, - message: - 'Bad Request: No valid session ID provided for non-initialize request.', - }, - id: null, - }); - return; - } - - try { - await transport.handleRequest(req, res, req.body); - } catch (error) { - const errorMessage = - error instanceof Error ? error.message : 'Unknown error'; - this.log(`Error handling MCP request: ${errorMessage}`); - if (!res.headersSent) { - res.status(500).json({ - jsonrpc: '2.0' as const, + if (transport.sessionId) { + this.log(`Session closed: ${transport.sessionId}`); + sessionsWithInitialNotification.delete(transport.sessionId); + delete transports[transport.sessionId]; + } + }; + mcpServer.connect(transport); + } else { + this.log( + 'Bad Request: No valid session ID provided for non-initialize request.', + ); + res.status(400).json({ + jsonrpc: '2.0', error: { - code: -32603, - message: 'Internal server error', + code: -32000, + message: + 'Bad Request: No valid session ID provided for non-initialize request.', }, id: null, }); + return; } - } - }); - const handleSessionRequest = async (req: Request, res: Response) => { - const sessionId = req.headers[MCP_SESSION_ID_HEADER] as - | string - | undefined; - if (!sessionId || !transports[sessionId]) { - this.log('Invalid or missing session ID'); - res.status(400).send('Invalid or missing session ID'); - return; - } - - const transport = transports[sessionId]; - try { - await transport.handleRequest(req, res); - } catch (error) { - const errorMessage = - error instanceof Error ? error.message : 'Unknown error'; - this.log(`Error handling session request: ${errorMessage}`); - if (!res.headersSent) { - res.status(400).send('Bad Request'); + try { + await transport.handleRequest(req, res, req.body); + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : 'Unknown error'; + this.log(`Error handling MCP request: ${errorMessage}`); + if (!res.headersSent) { + res.status(500).json({ + jsonrpc: '2.0' as const, + error: { + code: -32603, + message: 'Internal server error', + }, + id: null, + }); + } } - } + }); - if (!sessionsWithInitialNotification.has(sessionId)) { - sendIdeContextUpdateNotification( - transport, - this.log.bind(this), - openFilesManager, - ); - sessionsWithInitialNotification.add(sessionId); - } - }; + const handleSessionRequest = async (req: Request, res: Response) => { + const sessionId = req.headers[MCP_SESSION_ID_HEADER] as + | string + | undefined; + if (!sessionId || !transports[sessionId]) { + this.log('Invalid or missing session ID'); + res.status(400).send('Invalid or missing session ID'); + return; + } - app.get('/mcp', handleSessionRequest); + const transport = transports[sessionId]; + try { + await transport.handleRequest(req, res); + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : 'Unknown error'; + this.log(`Error handling session request: ${errorMessage}`); + if (!res.headersSent) { + res.status(400).send('Bad Request'); + } + } - this.server = app.listen(0, () => { - const address = (this.server as HTTPServer).address(); - if (address && typeof address !== 'string') { - const port = address.port; - context.environmentVariableCollection.replace( - IDE_SERVER_PORT_ENV_VAR, - port.toString(), - ); - this.log(`IDE server listening on port ${port}`); - fs.writeFile(this.portFile, JSON.stringify({ port })).catch((err) => { - this.log(`Failed to write port to file: ${err}`); - }); - this.log(this.portFile); - } + if (!sessionsWithInitialNotification.has(sessionId)) { + sendIdeContextUpdateNotification( + transport, + this.log.bind(this), + openFilesManager, + ); + sessionsWithInitialNotification.add(sessionId); + } + }; + + app.get('/mcp', handleSessionRequest); + + this.server = app.listen(0, async () => { + const address = (this.server as HTTPServer).address(); + if (address && typeof address !== 'string') { + this.port = address.port; + this.log(`IDE server listening on port ${this.port}`); + await writePortAndWorkspace( + context, + this.port, + this.portFile, + this.log, + ); + } + resolve(); + }); }); } + async updateWorkspacePath(): Promise { + if (this.context && this.port) { + await writePortAndWorkspace( + this.context, + this.port, + this.portFile, + this.log, + ); + } + } + async stop(): Promise { if (this.server) { await new Promise((resolve, reject) => {