diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index 6d25367c..43c55e7d 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/packages/vscode-ide-companion/src/ide-server.ts @@ -8,28 +8,81 @@ import * as vscode from 'vscode'; import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js'; import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js'; import express, { Request, Response } from 'express'; +import { randomUUID } from 'node:crypto'; +import { + isInitializeRequest, + type JSONRPCNotification, +} from '@modelcontextprotocol/sdk/types.js'; -export async function startIDEServer(_context: vscode.ExtensionContext) { +export async function startIDEServer(context: vscode.ExtensionContext) { const app = express(); app.use(express.json()); - const mcpServer = createMcpServer(); - const transport = new StreamableHTTPServerTransport({ - sessionIdGenerator: undefined, - enableJsonResponse: true, - }); + const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; - mcpServer.connect(transport); + const disposable = vscode.window.onDidChangeActiveTextEditor((editor) => { + const filePath = editor ? editor.document.uri.fsPath : null; + const notification: JSONRPCNotification = { + jsonrpc: '2.0', + method: 'ide/activeFileChanged', + params: { filePath }, + }; + for (const transport of Object.values(transports)) { + transport.send(notification); + } + }); + context.subscriptions.push(disposable); app.post('/mcp', async (req: Request, res: Response) => { - console.log('Received MCP request:', req.body); + const sessionId = req.headers['mcp-session-id'] 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) => { + transports[newSessionId] = transport; + const editor = vscode.window.activeTextEditor; + const filePath = editor ? editor.document.uri.fsPath : null; + const notification: JSONRPCNotification = { + jsonrpc: '2.0', + method: 'ide/activeFileChanged', + params: { filePath }, + }; + transport.send(notification); + }, + }); + + transport.onclose = () => { + if (transport.sessionId) { + delete transports[transport.sessionId]; + } + }; + + const server = createMcpServer(); + server.connect(transport); + } else { + 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) { console.error('Error handling MCP request:', error); if (!res.headersSent) { res.status(500).json({ - jsonrpc: '2.0', + jsonrpc: '2.0' as const, error: { code: -32603, message: 'Internal server error', @@ -40,15 +93,29 @@ export async function startIDEServer(_context: vscode.ExtensionContext) { } }); - // Handle GET requests for SSE streams - app.get('/mcp', async (req: Request, res: Response) => { - res.status(405).set('Allow', 'POST').send('Method Not Allowed'); - }); + const handleSessionRequest = async (req: Request, res: Response) => { + const sessionId = req.headers['mcp-session-id'] as string | undefined; + if (!sessionId || !transports[sessionId]) { + res.status(400).send('Invalid or missing session ID'); + return; + } + + const transport = transports[sessionId]; + try { + await transport.handleRequest(req, res); + } catch (error) { + console.error('Error handling MCP GET request:', error); + if (!res.headersSent) { + res.status(400).send('Bad Request'); + } + } + }; + + app.get('/mcp', handleSessionRequest); - // Start the server // TODO(#3918): Generate dynamically and write to env variable const PORT = 3000; - app.listen(PORT, (error) => { + app.listen(PORT, (error?: Error) => { if (error) { console.error('Failed to start server:', error); vscode.window.showErrorMessage( @@ -60,10 +127,13 @@ export async function startIDEServer(_context: vscode.ExtensionContext) { } const createMcpServer = () => { - const server = new McpServer({ - name: 'vscode-ide-server', - version: '1.0.0', - }); + const server = new McpServer( + { + name: 'gemini-cli-companion-mcp-server', + version: '1.0.0', + }, + { capabilities: { logging: {} } }, + ); server.registerTool( 'getActiveFile', {