feat(ide ext): Write workspace path to port file (#6659)

packages/vscode-ide-companion/src/extension-multi-folder.test.ts

@@ -1,211 +0,0 @@
-/**
- * @license
- * Copyright 2025 Google LLC
- * SPDX-License-Identifier: Apache-2.0
- */
-
-import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
-import * as vscode from 'vscode';
-import * as path from 'path';
-import { activate } from './extension.js';
-
-vi.mock('vscode', () => ({
-  window: {
-    createOutputChannel: vi.fn(() => ({
-      appendLine: vi.fn(),
-    })),
-    showInformationMessage: vi.fn(),
-    createTerminal: vi.fn(() => ({
-      show: vi.fn(),
-      sendText: vi.fn(),
-    })),
-    onDidChangeActiveTextEditor: vi.fn(),
-    activeTextEditor: undefined,
-    tabGroups: {
-      all: [],
-      close: vi.fn(),
-    },
-    showTextDocument: vi.fn(),
-  },
-  workspace: {
-    workspaceFolders: [],
-    onDidCloseTextDocument: vi.fn(),
-    registerTextDocumentContentProvider: vi.fn(),
-    onDidChangeWorkspaceFolders: vi.fn(),
-  },
-  commands: {
-    registerCommand: vi.fn(),
-    executeCommand: vi.fn(),
-  },
-  Uri: {
-    joinPath: vi.fn(),
-    file: (path: string) => ({ fsPath: path }),
-  },
-  ExtensionMode: {
-    Development: 1,
-    Production: 2,
-  },
-  EventEmitter: vi.fn(() => ({
-    event: vi.fn(),
-    fire: vi.fn(),
-    dispose: vi.fn(),
-  })),
-}));
-
-describe('activate with multiple folders', () => {
-  let context: vscode.ExtensionContext;
-  let onDidChangeWorkspaceFoldersCallback: (
-    e: vscode.WorkspaceFoldersChangeEvent,
-  ) => void;
-
-  beforeEach(() => {
-    context = {
-      subscriptions: [],
-      environmentVariableCollection: {
-        replace: vi.fn(),
-      },
-      globalState: {
-        get: vi.fn().mockReturnValue(true),
-        update: vi.fn(),
-      },
-      extensionUri: {
-        fsPath: '/path/to/extension',
-      },
-    } as unknown as vscode.ExtensionContext;
-
-    vi.mocked(vscode.workspace.onDidChangeWorkspaceFolders).mockImplementation(
-      (callback) => {
-        onDidChangeWorkspaceFoldersCallback = callback;
-        return { dispose: vi.fn() };
-      },
-    );
-  });
-
-  afterEach(() => {
-    vi.restoreAllMocks();
-  });
-
-  it('should set a single folder path', async () => {
-    const workspaceFoldersSpy = vi.spyOn(
-      vscode.workspace,
-      'workspaceFolders',
-      'get',
-    );
-    workspaceFoldersSpy.mockReturnValue([
-      { uri: { fsPath: '/foo/bar' } },
-    ] as vscode.WorkspaceFolder[]);
-
-    await activate(context);
-
-    expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith(
-      'GEMINI_CLI_IDE_WORKSPACE_PATH',
-      '/foo/bar',
-    );
-  });
-
-  it('should set multiple folder paths, separated by OS-specific path delimiter', async () => {
-    const workspaceFoldersSpy = vi.spyOn(
-      vscode.workspace,
-      'workspaceFolders',
-      'get',
-    );
-    workspaceFoldersSpy.mockReturnValue([
-      { uri: { fsPath: '/foo/bar' } },
-      { uri: { fsPath: '/baz/qux' } },
-    ] as vscode.WorkspaceFolder[]);
-
-    await activate(context);
-
-    expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith(
-      'GEMINI_CLI_IDE_WORKSPACE_PATH',
-      ['/foo/bar', '/baz/qux'].join(path.delimiter),
-    );
-  });
-
-  it('should set an empty string if no folders are open', async () => {
-    const workspaceFoldersSpy = vi.spyOn(
-      vscode.workspace,
-      'workspaceFolders',
-      'get',
-    );
-    workspaceFoldersSpy.mockReturnValue([]);
-
-    await activate(context);
-
-    expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith(
-      'GEMINI_CLI_IDE_WORKSPACE_PATH',
-      '',
-    );
-  });
-
-  it('should update the path when workspace folders change', async () => {
-    const workspaceFoldersSpy = vi.spyOn(
-      vscode.workspace,
-      'workspaceFolders',
-      'get',
-    );
-    workspaceFoldersSpy.mockReturnValue([
-      { uri: { fsPath: '/foo/bar' } },
-    ] as vscode.WorkspaceFolder[]);
-
-    await activate(context);
-
-    expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith(
-      'GEMINI_CLI_IDE_WORKSPACE_PATH',
-      '/foo/bar',
-    );
-
-    // Simulate adding a folder
-    workspaceFoldersSpy.mockReturnValue([
-      { uri: { fsPath: '/foo/bar' } },
-      { uri: { fsPath: '/baz/qux' } },
-    ] as vscode.WorkspaceFolder[]);
-    onDidChangeWorkspaceFoldersCallback({
-      added: [{ uri: { fsPath: '/baz/qux' } } as vscode.WorkspaceFolder],
-      removed: [],
-    });
-
-    expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith(
-      'GEMINI_CLI_IDE_WORKSPACE_PATH',
-      ['/foo/bar', '/baz/qux'].join(path.delimiter),
-    );
-
-    // Simulate removing a folder
-    workspaceFoldersSpy.mockReturnValue([
-      { uri: { fsPath: '/baz/qux' } },
-    ] as vscode.WorkspaceFolder[]);
-    onDidChangeWorkspaceFoldersCallback({
-      added: [],
-      removed: [{ uri: { fsPath: '/foo/bar' } } as vscode.WorkspaceFolder],
-    });
-
-    expect(context.environmentVariableCollection.replace).toHaveBeenCalledWith(
-      'GEMINI_CLI_IDE_WORKSPACE_PATH',
-      '/baz/qux',
-    );
-  });
-
-  it.skipIf(process.platform !== 'win32')(
-    'should handle windows paths',
-    async () => {
-      const workspaceFoldersSpy = vi.spyOn(
-        vscode.workspace,
-        'workspaceFolders',
-        'get',
-      );
-      workspaceFoldersSpy.mockReturnValue([
-        { uri: { fsPath: 'c:/foo/bar' } },
-        { uri: { fsPath: 'd:/baz/qux' } },
-      ] as vscode.WorkspaceFolder[]);
-
-      await activate(context);
-
-      expect(
-        context.environmentVariableCollection.replace,
-      ).toHaveBeenCalledWith(
-        'GEMINI_CLI_IDE_WORKSPACE_PATH',
-        'c:/foo/bar;d:/baz/qux',
-      );
-    },
-  );
-});

packages/vscode-ide-companion/src/extension.ts

@@ -5,13 +5,11 @@
  */
 
 import * as vscode from 'vscode';
-import * as path from 'path';
 import { IDEServer } from './ide-server.js';
 import { DiffContentProvider, DiffManager } from './diff-manager.js';
 import { createLogger } from './utils/logger.js';
 
 const INFO_MESSAGE_SHOWN_KEY = 'geminiCliInfoMessageShown';
-const IDE_WORKSPACE_PATH_ENV_VAR = 'GEMINI_CLI_IDE_WORKSPACE_PATH';
 export const DIFF_SCHEME = 'gemini-diff';
 
 let ideServer: IDEServer;
@@ -19,31 +17,11 @@ let logger: vscode.OutputChannel;
 
 let log: (message: string) => void = () => {};
 
-function updateWorkspacePath(context: vscode.ExtensionContext) {
-  const workspaceFolders = vscode.workspace.workspaceFolders;
-  if (workspaceFolders && workspaceFolders.length > 0) {
-    const workspacePaths = workspaceFolders
-      .map((folder) => folder.uri.fsPath)
-      .join(path.delimiter);
-    context.environmentVariableCollection.replace(
-      IDE_WORKSPACE_PATH_ENV_VAR,
-      workspacePaths,
-    );
-  } else {
-    context.environmentVariableCollection.replace(
-      IDE_WORKSPACE_PATH_ENV_VAR,
-      '',
-    );
-  }
-}
-
 export async function activate(context: vscode.ExtensionContext) {
   logger = vscode.window.createOutputChannel('Gemini CLI IDE Companion');
   log = createLogger(context, logger);
   log('Extension activated');
 
-  updateWorkspacePath(context);
-
   const diffContentProvider = new DiffContentProvider();
   const diffManager = new DiffManager(log, diffContentProvider);
 
@@ -94,7 +72,7 @@ export async function activate(context: vscode.ExtensionContext) {
 
   context.subscriptions.push(
     vscode.workspace.onDidChangeWorkspaceFolders(() => {
-      updateWorkspacePath(context);
+      ideServer.updateWorkspacePath();
     }),
     vscode.commands.registerCommand('gemini-cli.runGeminiCLI', async () => {
       const workspaceFolders = vscode.workspace.workspaceFolders;

packages/vscode-ide-companion/src/ide-server.test.ts

@@ -0,0 +1,287 @@
+/**
+ * @license
+ * Copyright 2025 Google LLC
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import type * as vscode from 'vscode';
+import * as fs from 'node:fs/promises';
+import * as os from 'node:os';
+import * as path from 'node:path';
+import { IDEServer } from './ide-server.js';
+import { DiffManager } from './diff-manager.js';
+
+const mocks = vi.hoisted(() => ({
+  diffManager: {
+    onDidChange: vi.fn(() => ({ dispose: vi.fn() })),
+  } as unknown as DiffManager,
+}));
+
+vi.mock('node:fs/promises', () => ({
+  writeFile: vi.fn(() => Promise.resolve(undefined)),
+  unlink: vi.fn(() => Promise.resolve(undefined)),
+}));
+
+vi.mock('node:os', async (importOriginal) => {
+  const actual = await importOriginal<typeof os>();
+  return {
+    ...actual,
+    tmpdir: vi.fn(() => '/tmp'),
+  };
+});
+
+const vscodeMock = vi.hoisted(() => ({
+  workspace: {
+    workspaceFolders: [
+      {
+        uri: {
+          fsPath: '/test/workspace1',
+        },
+      },
+      {
+        uri: {
+          fsPath: '/test/workspace2',
+        },
+      },
+    ],
+  },
+}));
+
+vi.mock('vscode', () => vscodeMock);
+
+vi.mock('./open-files-manager', () => {
+  const OpenFilesManager = vi.fn();
+  OpenFilesManager.prototype.onDidChange = vi.fn(() => ({ dispose: vi.fn() }));
+  return { OpenFilesManager };
+});
+
+describe('IDEServer', () => {
+  let ideServer: IDEServer;
+  let mockContext: vscode.ExtensionContext;
+  let mockLog: (message: string) => void;
+
+  const getPortFromMock = (
+    replaceMock: ReturnType<
+      () => vscode.ExtensionContext['environmentVariableCollection']['replace']
+    >,
+  ) => {
+    const port = vi
+      .mocked(replaceMock)
+      .mock.calls.find((call) => call[0] === 'GEMINI_CLI_IDE_SERVER_PORT')?.[1];
+
+    if (port === undefined) {
+      expect.fail('Port was not set');
+    }
+    return port;
+  };
+
+  beforeEach(() => {
+    mockLog = vi.fn();
+    ideServer = new IDEServer(mockLog, mocks.diffManager);
+    mockContext = {
+      subscriptions: [],
+      environmentVariableCollection: {
+        replace: vi.fn(),
+        clear: vi.fn(),
+      },
+    } as unknown as vscode.ExtensionContext;
+  });
+
+  afterEach(async () => {
+    await ideServer.stop();
+    vi.restoreAllMocks();
+    vscodeMock.workspace.workspaceFolders = [
+      { uri: { fsPath: '/test/workspace1' } },
+      { uri: { fsPath: '/test/workspace2' } },
+    ];
+  });
+
+  it('should set environment variables and workspace path on start with multiple folders', async () => {
+    await ideServer.start(mockContext);
+
+    const replaceMock = mockContext.environmentVariableCollection.replace;
+    expect(replaceMock).toHaveBeenCalledTimes(2);
+
+    expect(replaceMock).toHaveBeenNthCalledWith(
+      1,
+      'GEMINI_CLI_IDE_SERVER_PORT',
+      expect.any(String), // port is a number as a string
+    );
+
+    const expectedWorkspacePaths = [
+      '/test/workspace1',
+      '/test/workspace2',
+    ].join(path.delimiter);
+
+    expect(replaceMock).toHaveBeenNthCalledWith(
+      2,
+      'GEMINI_CLI_IDE_WORKSPACE_PATH',
+      expectedWorkspacePaths,
+    );
+
+    const port = getPortFromMock(replaceMock);
+    const expectedPortFile = path.join(
+      '/tmp',
+      `gemini-ide-server-${process.ppid}.json`,
+    );
+    expect(fs.writeFile).toHaveBeenCalledWith(
+      expectedPortFile,
+      JSON.stringify({
+        port: parseInt(port, 10),
+        workspacePath: expectedWorkspacePaths,
+      }),
+    );
+  });
+
+  it('should set a single folder path', async () => {
+    vscodeMock.workspace.workspaceFolders = [{ uri: { fsPath: '/foo/bar' } }];
+
+    await ideServer.start(mockContext);
+    const replaceMock = mockContext.environmentVariableCollection.replace;
+
+    expect(replaceMock).toHaveBeenCalledWith(
+      'GEMINI_CLI_IDE_WORKSPACE_PATH',
+      '/foo/bar',
+    );
+
+    const port = getPortFromMock(replaceMock);
+    const expectedPortFile = path.join(
+      '/tmp',
+      `gemini-ide-server-${process.ppid}.json`,
+    );
+    expect(fs.writeFile).toHaveBeenCalledWith(
+      expectedPortFile,
+      JSON.stringify({
+        port: parseInt(port, 10),
+        workspacePath: '/foo/bar',
+      }),
+    );
+  });
+
+  it('should set an empty string if no folders are open', async () => {
+    vscodeMock.workspace.workspaceFolders = [];
+
+    await ideServer.start(mockContext);
+    const replaceMock = mockContext.environmentVariableCollection.replace;
+
+    expect(replaceMock).toHaveBeenCalledWith(
+      'GEMINI_CLI_IDE_WORKSPACE_PATH',
+      '',
+    );
+
+    const port = getPortFromMock(replaceMock);
+    const expectedPortFile = path.join(
+      '/tmp',
+      `gemini-ide-server-${process.ppid}.json`,
+    );
+    expect(fs.writeFile).toHaveBeenCalledWith(
+      expectedPortFile,
+      JSON.stringify({
+        port: parseInt(port, 10),
+        workspacePath: '',
+      }),
+    );
+  });
+
+  it('should update the path when workspace folders change', async () => {
+    vscodeMock.workspace.workspaceFolders = [{ uri: { fsPath: '/foo/bar' } }];
+    await ideServer.start(mockContext);
+    const replaceMock = mockContext.environmentVariableCollection.replace;
+
+    expect(replaceMock).toHaveBeenCalledWith(
+      'GEMINI_CLI_IDE_WORKSPACE_PATH',
+      '/foo/bar',
+    );
+
+    // Simulate adding a folder
+    vscodeMock.workspace.workspaceFolders = [
+      { uri: { fsPath: '/foo/bar' } },
+      { uri: { fsPath: '/baz/qux' } },
+    ];
+    await ideServer.updateWorkspacePath();
+
+    const expectedWorkspacePaths = ['/foo/bar', '/baz/qux'].join(
+      path.delimiter,
+    );
+    expect(replaceMock).toHaveBeenCalledWith(
+      'GEMINI_CLI_IDE_WORKSPACE_PATH',
+      expectedWorkspacePaths,
+    );
+
+    const port = getPortFromMock(replaceMock);
+    const expectedPortFile = path.join(
+      '/tmp',
+      `gemini-ide-server-${process.ppid}.json`,
+    );
+    expect(fs.writeFile).toHaveBeenCalledWith(
+      expectedPortFile,
+      JSON.stringify({
+        port: parseInt(port, 10),
+        workspacePath: expectedWorkspacePaths,
+      }),
+    );
+
+    // Simulate removing a folder
+    vscodeMock.workspace.workspaceFolders = [{ uri: { fsPath: '/baz/qux' } }];
+    await ideServer.updateWorkspacePath();
+
+    expect(replaceMock).toHaveBeenCalledWith(
+      'GEMINI_CLI_IDE_WORKSPACE_PATH',
+      '/baz/qux',
+    );
+    expect(fs.writeFile).toHaveBeenCalledWith(
+      expectedPortFile,
+      JSON.stringify({
+        port: parseInt(port, 10),
+        workspacePath: '/baz/qux',
+      }),
+    );
+  });
+
+  it('should clear env vars and delete port file on stop', async () => {
+    await ideServer.start(mockContext);
+    const portFile = path.join(
+      '/tmp',
+      `gemini-ide-server-${process.ppid}.json`,
+    );
+    expect(fs.writeFile).toHaveBeenCalledWith(portFile, expect.any(String));
+
+    await ideServer.stop();
+
+    expect(mockContext.environmentVariableCollection.clear).toHaveBeenCalled();
+    expect(fs.unlink).toHaveBeenCalledWith(portFile);
+  });
+
+  it.skipIf(process.platform !== 'win32')(
+    'should handle windows paths',
+    async () => {
+      vscodeMock.workspace.workspaceFolders = [
+        { uri: { fsPath: 'c:\\foo\\bar' } },
+        { uri: { fsPath: 'd:\\baz\\qux' } },
+      ];
+
+      await ideServer.start(mockContext);
+      const replaceMock = mockContext.environmentVariableCollection.replace;
+      const expectedWorkspacePaths = 'c:\\foo\\bar;d:\\baz\\qux';
+
+      expect(replaceMock).toHaveBeenCalledWith(
+        'GEMINI_CLI_IDE_WORKSPACE_PATH',
+        expectedWorkspacePaths,
+      );
+
+      const port = getPortFromMock(replaceMock);
+      const expectedPortFile = path.join(
+        '/tmp',
+        `gemini-ide-server-${process.ppid}.json`,
+      );
+      expect(fs.writeFile).toHaveBeenCalledWith(
+        expectedPortFile,
+        JSON.stringify({
+          port: parseInt(port, 10),
+          workspacePath: expectedWorkspacePaths,
+        }),
+      );
+    },
+  );
+});

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) {
+  start(context: vscode.ExtensionContext): Promise<void> {
-    this.context = context;
+    return new Promise((resolve) => {
-    const sessionsWithInitialNotification = new Set<string>();
+      this.context = context;
-    const transports: { [sessionId: string]: StreamableHTTPServerTransport } =
+      const sessionsWithInitialNotification = new Set<string>();
-      {};
+      const transports: { [sessionId: string]: StreamableHTTPServerTransport } =
+        {};
 
-    const app = express();
+      const app = express();
-    app.use(express.json());
+      app.use(express.json());
-    const mcpServer = createMcpServer(this.diffManager);
+      const mcpServer = createMcpServer(this.diffManager);
 
-    const openFilesManager = new OpenFilesManager(context);
+      const openFilesManager = new OpenFilesManager(context);
-    const onDidChangeSubscription = openFilesManager.onDidChange(() => {
+      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) => {
         for (const transport of Object.values(transports)) {
-          transport.send(notification);
+          sendIdeContextUpdateNotification(
+            transport,
+            this.log.bind(this),
+            openFilesManager,
+          );
         }
-      },
+      });
-    );
+      context.subscriptions.push(onDidChangeSubscription);
-    context.subscriptions.push(onDidChangeDiffSubscription);
+      const onDidChangeDiffSubscription = this.diffManager.onDidChange(
+        (notification) => {
+          for (const transport of Object.values(transports)) {
+            transport.send(notification);
+          }
+        },
+      );
+      context.subscriptions.push(onDidChangeDiffSubscription);
 
-    app.post('/mcp', async (req: Request, res: Response) => {
+      app.post('/mcp', async (req: Request, res: Response) => {
-      const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
+        const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
-        | string
+          | string
-        | undefined;
+          | undefined;
-      let transport: StreamableHTTPServerTransport;
+        let transport: StreamableHTTPServerTransport;
 
-      if (sessionId && transports[sessionId]) {
+        if (sessionId && transports[sessionId]) {
-        transport = transports[sessionId];
+          transport = transports[sessionId];
-      } else if (!sessionId && isInitializeRequest(req.body)) {
+        } else if (!sessionId && isInitializeRequest(req.body)) {
-        transport = new StreamableHTTPServerTransport({
+          transport = new StreamableHTTPServerTransport({
-          sessionIdGenerator: () => randomUUID(),
+            sessionIdGenerator: () => randomUUID(),
-          onsessioninitialized: (newSessionId) => {
+            onsessioninitialized: (newSessionId) => {
-            this.log(`New session initialized: ${newSessionId}`);
+              this.log(`New session initialized: ${newSessionId}`);
-            transports[newSessionId] = transport;
+              transports[newSessionId] = transport;
-          },
+            },
-        });
+          });
-        const keepAlive = setInterval(() => {
+          const keepAlive = setInterval(() => {
-          try {
+            try {
-            transport.send({ jsonrpc: '2.0', method: 'ping' });
+              transport.send({ jsonrpc: '2.0', method: 'ping' });
-          } catch (e) {
+            } catch (e) {
-            this.log(
+              this.log(
-              'Failed to send keep-alive ping, cleaning up interval.' + e,
+                'Failed to send keep-alive ping, cleaning up interval.' + e,
-            );
+              );
+              clearInterval(keepAlive);
+            }
+          }, 60000); // 60 sec
+
+          transport.onclose = () => {
             clearInterval(keepAlive);
-          }
+            if (transport.sessionId) {
-        }, 60000); // 60 sec
+              this.log(`Session closed: ${transport.sessionId}`);
-
+              sessionsWithInitialNotification.delete(transport.sessionId);
-        transport.onclose = () => {
+              delete transports[transport.sessionId];
-          clearInterval(keepAlive);
+            }
-          if (transport.sessionId) {
+          };
-            this.log(`Session closed: ${transport.sessionId}`);
+          mcpServer.connect(transport);
-            sessionsWithInitialNotification.delete(transport.sessionId);
+        } else {
-            delete transports[transport.sessionId];
+          this.log(
-          }
+            'Bad Request: No valid session ID provided for non-initialize request.',
-        };
+          );
-        mcpServer.connect(transport);
+          res.status(400).json({
-      } else {
+            jsonrpc: '2.0',
-        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;
-      }
-
-      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,
+              code: -32000,
-              message: 'Internal server error',
+              message:
+                'Bad Request: No valid session ID provided for non-initialize request.',
             },
             id: null,
           });
+          return;
         }
-      }
-    });
 
-    const handleSessionRequest = async (req: Request, res: Response) => {
+        try {
-      const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
+          await transport.handleRequest(req, res, req.body);
-        | string
+        } catch (error) {
-        | undefined;
+          const errorMessage =
-      if (!sessionId || !transports[sessionId]) {
+            error instanceof Error ? error.message : 'Unknown error';
-        this.log('Invalid or missing session ID');
+          this.log(`Error handling MCP request: ${errorMessage}`);
-        res.status(400).send('Invalid or missing session ID');
+          if (!res.headersSent) {
-        return;
+            res.status(500).json({
-      }
+              jsonrpc: '2.0' as const,
-
+              error: {
-      const transport = transports[sessionId];
+                code: -32603,
-      try {
+                message: 'Internal server error',
-        await transport.handleRequest(req, res);
+              },
-      } catch (error) {
+              id: null,
-        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');
         }
-      }
+      });
 
-      if (!sessionsWithInitialNotification.has(sessionId)) {
+      const handleSessionRequest = async (req: Request, res: Response) => {
-        sendIdeContextUpdateNotification(
+        const sessionId = req.headers[MCP_SESSION_ID_HEADER] as
-          transport,
+          | string
-          this.log.bind(this),
+          | undefined;
-          openFilesManager,
+        if (!sessionId || !transports[sessionId]) {
-        );
+          this.log('Invalid or missing session ID');
-        sessionsWithInitialNotification.add(sessionId);
+          res.status(400).send('Invalid or missing session ID');
-      }
+          return;
-    };
+        }
 
-    app.get('/mcp', handleSessionRequest);
+        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');
+          }
+        }
 
-    this.server = app.listen(0, () => {
+        if (!sessionsWithInitialNotification.has(sessionId)) {
-      const address = (this.server as HTTPServer).address();
+          sendIdeContextUpdateNotification(
-      if (address && typeof address !== 'string') {
+            transport,
-        const port = address.port;
+            this.log.bind(this),
-        context.environmentVariableCollection.replace(
+            openFilesManager,
-          IDE_SERVER_PORT_ENV_VAR,
+          );
-          port.toString(),
+          sessionsWithInitialNotification.add(sessionId);
-        );
+        }
-        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}`);
+      app.get('/mcp', handleSessionRequest);
-        });
+
-        this.log(this.portFile);
+      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) => {