diff options
Diffstat (limited to 'packages/vscode-ide-companion/src/ide-server.ts')
| -rw-r--r-- | packages/vscode-ide-companion/src/ide-server.ts | 306 |
1 files changed, 175 insertions, 131 deletions
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<void> { + 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<string>(); - const transports: { [sessionId: string]: StreamableHTTPServerTransport } = - {}; + start(context: vscode.ExtensionContext): Promise<void> { + return new Promise((resolve) => { + this.context = context; + const sessionsWithInitialNotification = new Set<string>(); + 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); - - 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, - ); - clearInterval(keepAlive); + }); + context.subscriptions.push(onDidChangeSubscription); + const onDidChangeDiffSubscription = this.diffManager.onDidChange( + (notification) => { + for (const transport of Object.values(transports)) { + transport.send(notification); } - }, 60000); // 60 sec + }, + ); + context.subscriptions.push(onDidChangeDiffSubscription); - 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; - } + app.post('/mcp', async (req: Request, res: Response) => { + const sessionId = req.headers[MCP_SESSION_ID_HEADER] as + | string + | undefined; + let transport: StreamableHTTPServerTransport; - 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 (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); + 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; - } + 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, + }); + } + } + }); - 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'); + 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; } - } - if (!sessionsWithInitialNotification.has(sessionId)) { - sendIdeContextUpdateNotification( - transport, - this.log.bind(this), - openFilesManager, - ); - sessionsWithInitialNotification.add(sessionId); - } - }; + 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'); + } + } - app.get('/mcp', handleSessionRequest); + if (!sessionsWithInitialNotification.has(sessionId)) { + sendIdeContextUpdateNotification( + transport, + this.log.bind(this), + openFilesManager, + ); + sessionsWithInitialNotification.add(sessionId); + } + }; - 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); - } + 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<void> { + if (this.context && this.port) { + await writePortAndWorkspace( + this.context, + this.port, + this.portFile, + this.log, + ); + } + } + async stop(): Promise<void> { if (this.server) { await new Promise<void>((resolve, reject) => { |