Minor refactoring of IDE companion server (#4331)

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();
+  }
+}

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,22 +33,36 @@ function sendActiveFileChangedNotification(
   transport.send(notification);
 }
 
-export async function startIDEServer(context: vscode.ExtensionContext) {
+export class IDEServer {
+  private server: HTTPServer | undefined;
+  private context: vscode.ExtensionContext | undefined;
+  private logger: vscode.OutputChannel;
+
+  constructor(logger: vscode.OutputChannel) {
+    this.logger = logger;
+  }
+
+  async start(context: vscode.ExtensionContext) {
+    this.context = context;
+    const transports: { [sessionId: string]: StreamableHTTPServerTransport } =
+      {};
+    const sessionsWithInitialNotification = new Set<string>();
+
     const app = express();
     app.use(express.json());
-
+    const mcpServer = createMcpServer();
-  const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {};
-  const sessionsWithInitialNotification = new Set<string>();
 
     const disposable = vscode.window.onDidChangeActiveTextEditor((_editor) => {
       for (const transport of Object.values(transports)) {
-      sendActiveFileChangedNotification(transport);
+        sendActiveFileChangedNotification(transport, this.logger);
       }
     });
     context.subscriptions.push(disposable);
 
     app.post('/mcp', async (req: Request, res: Response) => {
-    const sessionId = req.headers['mcp-session-id'] as string | undefined;
+      const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
+        | string
+        | undefined;
       let transport: StreamableHTTPServerTransport;
 
       if (sessionId && transports[sessionId]) {
@@ -53,20 +71,22 @@ export async function startIDEServer(context: vscode.ExtensionContext) {
         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);
-      const server = createMcpServer();
-      server.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: {
@@ -82,7 +102,9 @@ export async function startIDEServer(context: vscode.ExtensionContext) {
       try {
         await transport.handleRequest(req, res, req.body);
       } catch (error) {
-      console.error('Error handling MCP request:', 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,
@@ -97,48 +119,71 @@ export async function startIDEServer(context: vscode.ExtensionContext) {
     });
 
     const handleSessionRequest = async (req: Request, res: Response) => {
-    const sessionId = req.headers['mcp-session-id'] as string | undefined;
+      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];
-
       try {
         await transport.handleRequest(req, res);
       } catch (error) {
-      console.error('Error handling MCP GET request:', 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');
         }
       }
 
       if (!sessionsWithInitialNotification.has(sessionId)) {
-      sendActiveFileChangedNotification(transport);
+        sendActiveFileChangedNotification(transport, this.logger);
         sessionsWithInitialNotification.add(sessionId);
       }
     };
 
     app.get('/mcp', handleSessionRequest);
 
-  const server = app.listen(0, () => {
+    this.server = app.listen(0, () => {
-    const address = (server as Server).address();
+      const address = (this.server as HTTPServer).address();
       if (address && typeof address !== 'string') {
         const port = address.port;
         context.environmentVariableCollection.replace(
-        'GEMINI_CLI_IDE_SERVER_PORT',
+          IDE_SERVER_PORT_ENV_VAR,
           port.toString(),
         );
-      console.log(`MCP Streamable HTTP Server listening on port ${port}`);
+        this.logger.appendLine(`IDE 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`,
-      );
       }
     });
+  }
+
+  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;
+    }
+
+    if (this.context) {
+      this.context.environmentVariableCollection.clear();
+    }
+  }
 }
 
 const createMcpServer = () => {
@@ -157,7 +202,6 @@ const createMcpServer = () => {
       inputSchema: {},
     },
     async () => {
-      try {
       const activeEditor = vscode.window.activeTextEditor;
       const filePath = activeEditor ? activeEditor.document.uri.fsPath : '';
       if (filePath) {
@@ -174,18 +218,6 @@ const createMcpServer = () => {
           ],
         };
       }
-      } catch (error) {
-        return {
-          content: [
-            {
-              type: 'text',
-              text: `Failed to get active file: ${
-                (error as Error).message || 'Unknown error'
-              }`,
-            },
-          ],
-        };
-      }
     },
   );
   return server;