diff options
| -rw-r--r-- | packages/vscode-ide-companion/src/extension.ts | 28 | ||||
| -rw-r--r-- | packages/vscode-ide-companion/src/ide-server.ts | 258 |
2 files changed, 170 insertions, 116 deletions
diff --git a/packages/vscode-ide-companion/src/extension.ts b/packages/vscode-ide-companion/src/extension.ts index ecc63048..62901793 100644 --- a/packages/vscode-ide-companion/src/extension.ts +++ b/packages/vscode-ide-companion/src/extension.ts @@ -5,10 +5,32 @@ */ import * as vscode from 'vscode'; -import { startIDEServer } from './ide-server'; +import { IDEServer } from './ide-server'; + +let ideServer: IDEServer; +let logger: vscode.OutputChannel; export async function activate(context: vscode.ExtensionContext) { - startIDEServer(context); + logger = vscode.window.createOutputChannel('Gemini CLI IDE Companion'); + logger.show(); + logger.appendLine('Starting Gemini CLI IDE Companion server...'); + ideServer = new IDEServer(logger); + try { + await ideServer.start(context); + } catch (err) { + const message = err instanceof Error ? err.message : String(err); + logger.appendLine(`Failed to start IDE server: ${message}`); + } } -export function deactivate() {} +export function deactivate() { + if (ideServer) { + logger.appendLine('Deactivating Gemini CLI IDE Companion...'); + return ideServer.stop().finally(() => { + logger.dispose(); + }); + } + if (logger) { + logger.dispose(); + } +} diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts index 160cb54a..497929a4 100644 --- a/packages/vscode-ide-companion/src/ide-server.ts +++ b/packages/vscode-ide-companion/src/ide-server.ts @@ -13,14 +13,18 @@ import { isInitializeRequest, type JSONRPCNotification, } from '@modelcontextprotocol/sdk/types.js'; +import { Server as HTTPServer } from 'node:http'; -import { Server } from 'node:http'; +const MCP_SESSION_ID_HEADER = 'mcp-session-id'; +const IDE_SERVER_PORT_ENV_VAR = 'GEMINI_CLI_IDE_SERVER_PORT'; function sendActiveFileChangedNotification( transport: StreamableHTTPServerTransport, + logger: vscode.OutputChannel, ) { const editor = vscode.window.activeTextEditor; const filePath = editor ? editor.document.uri.fsPath : ''; + logger.appendLine(`Sending active file changed notification: ${filePath}`); const notification: JSONRPCNotification = { jsonrpc: '2.0', method: 'ide/activeFileChanged', @@ -29,116 +33,157 @@ function sendActiveFileChangedNotification( transport.send(notification); } -export async function startIDEServer(context: vscode.ExtensionContext) { - const app = express(); - app.use(express.json()); +export class IDEServer { + private server: HTTPServer | undefined; + private context: vscode.ExtensionContext | undefined; + private logger: vscode.OutputChannel; - const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {}; - const sessionsWithInitialNotification = new Set<string>(); + constructor(logger: vscode.OutputChannel) { + this.logger = logger; + } - const disposable = vscode.window.onDidChangeActiveTextEditor((_editor) => { - for (const transport of Object.values(transports)) { - sendActiveFileChangedNotification(transport); - } - }); - context.subscriptions.push(disposable); + async start(context: vscode.ExtensionContext) { + this.context = context; + const transports: { [sessionId: string]: StreamableHTTPServerTransport } = + {}; + const sessionsWithInitialNotification = new Set<string>(); - app.post('/mcp', async (req: Request, res: Response) => { - const sessionId = req.headers['mcp-session-id'] as string | undefined; - let transport: StreamableHTTPServerTransport; + const app = express(); + app.use(express.json()); + const mcpServer = createMcpServer(); - if (sessionId && transports[sessionId]) { - transport = transports[sessionId]; - } else if (!sessionId && isInitializeRequest(req.body)) { - transport = new StreamableHTTPServerTransport({ - sessionIdGenerator: () => randomUUID(), - onsessioninitialized: (newSessionId) => { - transports[newSessionId] = transport; - }, - }); + const disposable = vscode.window.onDidChangeActiveTextEditor((_editor) => { + for (const transport of Object.values(transports)) { + sendActiveFileChangedNotification(transport, this.logger); + } + }); + context.subscriptions.push(disposable); - transport.onclose = () => { - if (transport.sessionId) { - sessionsWithInitialNotification.delete(transport.sessionId); - delete transports[transport.sessionId]; - } - }; + app.post('/mcp', async (req: Request, res: Response) => { + const sessionId = req.headers[MCP_SESSION_ID_HEADER] as + | string + | undefined; + let transport: StreamableHTTPServerTransport; - 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' as const, + if (sessionId && transports[sessionId]) { + transport = transports[sessionId]; + } else if (!sessionId && isInitializeRequest(req.body)) { + transport = new StreamableHTTPServerTransport({ + sessionIdGenerator: () => randomUUID(), + onsessioninitialized: (newSessionId) => { + this.logger.appendLine(`New session initialized: ${newSessionId}`); + transports[newSessionId] = transport; + }, + }); + transport.onclose = () => { + if (transport.sessionId) { + this.logger.appendLine(`Session closed: ${transport.sessionId}`); + sessionsWithInitialNotification.delete(transport.sessionId); + delete transports[transport.sessionId]; + } + }; + mcpServer.connect(transport); + } else { + this.logger.appendLine( + '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'] as string | undefined; - if (!sessionId || !transports[sessionId]) { - 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.logger.appendLine(`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 handleSessionRequest = async (req: Request, res: Response) => { + const sessionId = req.headers[MCP_SESSION_ID_HEADER] as + | string + | undefined; + if (!sessionId || !transports[sessionId]) { + this.logger.appendLine('Invalid or missing session ID'); + res.status(400).send('Invalid or missing session ID'); + return; + } - const transport = transports[sessionId]; + const transport = transports[sessionId]; + try { + await transport.handleRequest(req, res); + } catch (error) { + const errorMessage = + error instanceof Error ? error.message : 'Unknown error'; + this.logger.appendLine( + `Error handling session request: ${errorMessage}`, + ); + if (!res.headersSent) { + res.status(400).send('Bad Request'); + } + } - 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'); + if (!sessionsWithInitialNotification.has(sessionId)) { + sendActiveFileChangedNotification(transport, this.logger); + sessionsWithInitialNotification.add(sessionId); } - } + }; - if (!sessionsWithInitialNotification.has(sessionId)) { - sendActiveFileChangedNotification(transport); - sessionsWithInitialNotification.add(sessionId); - } - }; + app.get('/mcp', handleSessionRequest); + + 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.logger.appendLine(`IDE server listening on port ${port}`); + } + }); + } - app.get('/mcp', handleSessionRequest); + async stop(): Promise<void> { + if (this.server) { + await new Promise<void>((resolve, reject) => { + this.server!.close((err?: Error) => { + if (err) { + this.logger.appendLine( + `Error shutting down IDE server: ${err.message}`, + ); + return reject(err); + } + this.logger.appendLine(`IDE server shut down`); + resolve(); + }); + }); + this.server = undefined; + } - const server = app.listen(0, () => { - const address = (server as Server).address(); - if (address && typeof address !== 'string') { - const port = address.port; - context.environmentVariableCollection.replace( - 'GEMINI_CLI_IDE_SERVER_PORT', - port.toString(), - ); - console.log(`MCP Streamable HTTP Server listening on port ${port}`); - } else { - const port = 0; - console.error('Failed to start server:', 'Unknown error'); - vscode.window.showErrorMessage( - `Companion server failed to start on port ${port}: Unknown error`, - ); + if (this.context) { + this.context.environmentVariableCollection.clear(); } - }); + } } const createMcpServer = () => { @@ -157,31 +202,18 @@ const createMcpServer = () => { inputSchema: {}, }, async () => { - try { - const activeEditor = vscode.window.activeTextEditor; - const filePath = activeEditor ? activeEditor.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.', - }, - ], - }; - } - } catch (error) { + const activeEditor = vscode.window.activeTextEditor; + const filePath = activeEditor ? activeEditor.document.uri.fsPath : ''; + if (filePath) { + return { + content: [{ type: 'text', text: `Active file: ${filePath}` }], + }; + } else { return { content: [ { type: 'text', - text: `Failed to get active file: ${ - (error as Error).message || 'Unknown error' - }`, + text: 'No file is currently active in the editor.', }, ], }; |