Minor refactoring of VS Code companion extension code (#4761)

1 files changed, 31 insertions, 26 deletions

diff --git a/packages/vscode-ide-companion/src/ide-server.ts b/packages/vscode-ide-companion/src/ide-server.ts
index 9111f349..75828485 100644
--- a/packages/vscode-ide-companion/src/ide-server.ts
+++ b/packages/vscode-ide-companion/src/ide-server.ts
@@ -21,30 +21,41 @@ const IDE_SERVER_PORT_ENV_VAR = 'GEMINI_CLI_IDE_SERVER_PORT';
 
 function sendOpenFilesChangedNotification(
   transport: StreamableHTTPServerTransport,
-  logger: vscode.OutputChannel,
+  log: (message: string) => void,
   recentFilesManager: RecentFilesManager,
 ) {
   const editor = vscode.window.activeTextEditor;
-  const filePath = editor ? editor.document.uri.fsPath : '';
-  logger.appendLine(`Sending active file changed notification: ${filePath}`);
+  const filePath =
+    editor && editor.document.uri.scheme === 'file'
+      ? editor.document.uri.fsPath
+      : '';
   const notification: JSONRPCNotification = {
     jsonrpc: '2.0',
     method: 'ide/openFilesChanged',
     params: {
       activeFile: filePath,
-      recentOpenFiles: recentFilesManager.recentFiles,
+      recentOpenFiles: recentFilesManager.recentFiles.filter(
+        (file) => file.filePath !== filePath,
+      ),
     },
   };
+  log(
+    `Sending active file changed notification: ${JSON.stringify(
+      notification,
+      null,
+      2,
+    )}`,
+  );
   transport.send(notification);
 }
 
 export class IDEServer {
   private server: HTTPServer | undefined;
   private context: vscode.ExtensionContext | undefined;
-  private logger: vscode.OutputChannel;
+  private log: (message: string) => void;
 
-  constructor(logger: vscode.OutputChannel) {
-    this.logger = logger;
+  constructor(log: (message: string) => void) {
+    this.log = log;
   }
 
   async start(context: vscode.ExtensionContext) {
@@ -62,7 +73,7 @@ export class IDEServer {
       for (const transport of Object.values(transports)) {
         sendOpenFilesChangedNotification(
           transport,
-          this.logger,
+          this.log.bind(this),
           recentFilesManager,
         );
       }
@@ -81,7 +92,7 @@ export class IDEServer {
         transport = new StreamableHTTPServerTransport({
           sessionIdGenerator: () => randomUUID(),
           onsessioninitialized: (newSessionId) => {
-            this.logger.appendLine(`New session initialized: ${newSessionId}`);
+            this.log(`New session initialized: ${newSessionId}`);
             transports[newSessionId] = transport;
           },
         });
@@ -90,26 +101,24 @@ export class IDEServer {
           try {
             transport.send({ jsonrpc: '2.0', method: 'ping' });
           } catch (e) {
-            // If sending a ping fails, the connection is likely broken.
-            // Log the error and clear the interval to prevent further attempts.
-            this.logger.append(
+            this.log(
               'Failed to send keep-alive ping, cleaning up interval.' + e,
             );
             clearInterval(keepAlive);
           }
-        }, 60000); // Send ping every 60 seconds
+        }, 60000); // 60 sec
 
         transport.onclose = () => {
           clearInterval(keepAlive);
           if (transport.sessionId) {
-            this.logger.appendLine(`Session closed: ${transport.sessionId}`);
+            this.log(`Session closed: ${transport.sessionId}`);
             sessionsWithInitialNotification.delete(transport.sessionId);
             delete transports[transport.sessionId];
           }
         };
         mcpServer.connect(transport);
       } else {
-        this.logger.appendLine(
+        this.log(
           'Bad Request: No valid session ID provided for non-initialize request.',
         );
         res.status(400).json({
@@ -129,7 +138,7 @@ export class IDEServer {
       } catch (error) {
         const errorMessage =
           error instanceof Error ? error.message : 'Unknown error';
-        this.logger.appendLine(`Error handling MCP request: ${errorMessage}`);
+        this.log(`Error handling MCP request: ${errorMessage}`);
         if (!res.headersSent) {
           res.status(500).json({
             jsonrpc: '2.0' as const,
@@ -148,7 +157,7 @@ export class IDEServer {
         | string
         | undefined;
       if (!sessionId || !transports[sessionId]) {
-        this.logger.appendLine('Invalid or missing session ID');
+        this.log('Invalid or missing session ID');
         res.status(400).send('Invalid or missing session ID');
         return;
       }
@@ -159,9 +168,7 @@ export class IDEServer {
       } catch (error) {
         const errorMessage =
           error instanceof Error ? error.message : 'Unknown error';
-        this.logger.appendLine(
-          `Error handling session request: ${errorMessage}`,
-        );
+        this.log(`Error handling session request: ${errorMessage}`);
         if (!res.headersSent) {
           res.status(400).send('Bad Request');
         }
@@ -170,7 +177,7 @@ export class IDEServer {
       if (!sessionsWithInitialNotification.has(sessionId)) {
         sendOpenFilesChangedNotification(
           transport,
-          this.logger,
+          this.log.bind(this),
           recentFilesManager,
         );
         sessionsWithInitialNotification.add(sessionId);
@@ -187,7 +194,7 @@ export class IDEServer {
           IDE_SERVER_PORT_ENV_VAR,
           port.toString(),
         );
-        this.logger.appendLine(`IDE server listening on port ${port}`);
+        this.log(`IDE server listening on port ${port}`);
       }
     });
   }
@@ -197,12 +204,10 @@ export class IDEServer {
       await new Promise<void>((resolve, reject) => {
         this.server!.close((err?: Error) => {
           if (err) {
-            this.logger.appendLine(
-              `Error shutting down IDE server: ${err.message}`,
-            );
+            this.log(`Error shutting down IDE server: ${err.message}`);
             return reject(err);
           }
-          this.logger.appendLine(`IDE server shut down`);
+          this.log(`IDE server shut down`);
           resolve();
         });
       });