fix: restart MCP server on vscode workspace folder changes

When workspace folders were added or removed, fireDidChange() notified
VS Code that the MCP server definitions changed, causing it to stop the
running server. However, because the returned McpStdioServerDefinition
had the same version as before, VS Code did not restart the server.

Add a monotonically increasing revision counter to McpProvider that
increments on each fireDidChange() call and is appended to the version
string (e.g. 2.25.1.1, 2.25.1.2). This signals VS Code that the
definition has genuinely changed and triggers a server restart with
the updated environment instead of only stopping the server.

- Add _revision counter and getEffectiveVersion() to McpProvider
- Add 4 unit tests for version-bumping behaviour
- Add integration test suite for workspace folder change scenarios

Author: data-douser

extensions/vscode/src/server/mcp-provider.ts

@@ -20,6 +20,14 @@ export class McpProvider
   private readonly _onDidChange = new vscode.EventEmitter<void>();
   readonly onDidChangeMcpServerDefinitions = this._onDidChange.event;
 
+  /**
+   * Monotonically increasing counter, bumped each time `fireDidChange()` is
+   * called. Appended to the version string returned by
+   * `provideMcpServerDefinitions()` so that VS Code sees a genuinely new
+   * server definition and restarts the MCP server instead of only stopping it.
+   */
+  private _revision = 0;
+
   constructor(
     private readonly serverManager: ServerManager,
     private readonly envBuilder: EnvironmentBuilder,
@@ -31,6 +39,7 @@ export class McpProvider
 
   /** Notify VS Code that the MCP server definitions have changed. */
   fireDidChange(): void {
+    this._revision++;
     this._onDidChange.fire();
   }
 
@@ -40,7 +49,7 @@ export class McpProvider
     const command = this.serverManager.getCommand();
     const args = this.serverManager.getArgs();
     const env = await this.envBuilder.build();
-    const version = this.serverManager.getVersion();
+    const version = this.getEffectiveVersion();
 
     this.logger.info(
       `Providing MCP server definition: ${command} ${args.join(' ')}`,
@@ -66,4 +75,26 @@ export class McpProvider
     server.env = env;
     return server;
   }
+
+  /**
+   * Computes the version string for the MCP server definition.
+   *
+   * Before any `fireDidChange()` call, returns the raw version from
+   * `serverManager.getVersion()` (preserving the original behaviour).
+   *
+   * After one or more `fireDidChange()` calls, appends a `.N` revision
+   * suffix so that the version is always different from the previous one.
+   * This is essential because VS Code uses the version to decide whether to
+   * restart a running server: if the version is unchanged, VS Code may stop
+   * the server without restarting it.
+   */
+  private getEffectiveVersion(): string | undefined {
+    if (this._revision === 0) {
+      return this.serverManager.getVersion();
+    }
+    const base =
+      this.serverManager.getVersion() ??
+      this.serverManager.getExtensionVersion();
+    return `${base}.${this._revision}`;
+  }
 }

extensions/vscode/test/server/mcp-provider.test.ts

@@ -7,6 +7,7 @@ function createMockServerManager() {
   return {
     getCommand: vi.fn().mockReturnValue('npx'),
     getArgs: vi.fn().mockReturnValue(['-y', 'codeql-development-mcp-server']),
+    getExtensionVersion: vi.fn().mockReturnValue('2.25.1'),
     getVersion: vi.fn().mockReturnValue(undefined),
     getDescription: vi.fn().mockReturnValue('npx -y codeql-development-mcp-server'),
     getInstallDir: vi.fn().mockReturnValue('/mock/install'),
@@ -120,6 +121,52 @@ describe('McpProvider', () => {
     expect(definitions![0].version).toBe('2.20.0');
   });
 
+  it('should produce a different version after fireDidChange to trigger server restart', async () => {
+    const token = { isCancellationRequested: false, onCancellationRequested: vi.fn() };
+    const before = await provider.provideMcpServerDefinitions(token as any);
+    const versionBefore = before![0].version;
+
+    provider.fireDidChange();
+    const after = await provider.provideMcpServerDefinitions(token as any);
+    const versionAfter = after![0].version;
+
+    expect(versionAfter).toBeDefined();
+    expect(versionAfter).not.toBe(versionBefore);
+  });
+
+  it('should produce distinct versions after successive fireDidChange calls', async () => {
+    const token = { isCancellationRequested: false, onCancellationRequested: vi.fn() };
+
+    provider.fireDidChange();
+    const defs1 = await provider.provideMcpServerDefinitions(token as any);
+
+    provider.fireDidChange();
+    const defs2 = await provider.provideMcpServerDefinitions(token as any);
+
+    expect(defs1![0].version).not.toBe(defs2![0].version);
+  });
+
+  it('should append revision to pinned version after fireDidChange', async () => {
+    serverManager.getVersion.mockReturnValue('2.20.0');
+    const token = { isCancellationRequested: false, onCancellationRequested: vi.fn() };
+
+    provider.fireDidChange();
+    const definitions = await provider.provideMcpServerDefinitions(token as any);
+
+    expect(definitions![0].version).toMatch(/^2\.20\.0\.\d+$/);
+  });
+
+  it('should append revision to extension version after fireDidChange when version is latest', async () => {
+    serverManager.getVersion.mockReturnValue(undefined);
+    serverManager.getExtensionVersion.mockReturnValue('2.25.1');
+    const token = { isCancellationRequested: false, onCancellationRequested: vi.fn() };
+
+    provider.fireDidChange();
+    const definitions = await provider.provideMcpServerDefinitions(token as any);
+
+    expect(definitions![0].version).toMatch(/^2\.25\.1\.\d+$/);
+  });
+
   it('should resolve definition by refreshing environment', async () => {
     const updatedEnv = { CODEQL_PATH: '/new/path', TRANSPORT_MODE: 'stdio' };
     envBuilder.build.mockResolvedValue(updatedEnv);

extensions/vscode/test/suite/workspace-folder-change.integration.test.ts

@@ -0,0 +1,229 @@
+/**
+ * Integration tests for workspace folder change handling.
+ *
+ * These run inside the Extension Development Host with the REAL VS Code API.
+ * They verify that the MCP server definition version changes when workspace
+ * folders are added or removed, which signals VS Code to restart the server
+ * with the updated environment rather than simply stopping it.
+ *
+ * Bug: Prior to the fix, fireDidChange() notified VS Code that the MCP server
+ * definitions changed, causing it to stop the running server. However, because
+ * the returned McpStdioServerDefinition had the same version as before, VS Code
+ * did not restart the server — it only stopped it.
+ */
+
+import * as assert from 'assert';
+import * as fs from 'fs';
+import * as os from 'os';
+import * as path from 'path';
+import * as vscode from 'vscode';
+
+const EXTENSION_ID = 'advanced-security.vscode-codeql-development-mcp-server';
+
+suite('Workspace Folder Change Tests', () => {
+  let api: any;
+
+  suiteSetup(async () => {
+    const ext = vscode.extensions.getExtension(EXTENSION_ID);
+    assert.ok(ext, `Extension ${EXTENSION_ID} not found`);
+    api = ext.isActive ? ext.exports : await ext.activate();
+  });
+
+  test('MCP definition version should change after workspace folder is added', async function () {
+    const provider = api.mcpProvider;
+    if (!provider) {
+      this.skip();
+      return;
+    }
+
+    const token: vscode.CancellationToken = {
+      isCancellationRequested: false,
+      onCancellationRequested: () => ({ dispose: () => {} }),
+    };
+
+    // Get initial definitions
+    const beforeDefs = await provider.provideMcpServerDefinitions(token);
+    assert.ok(beforeDefs && beforeDefs.length >= 1, 'Should provide at least one definition');
+    const versionBefore = beforeDefs[0].version;
+
+    // Create a real temporary directory to add as a workspace folder
+    const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ql-mcp-test-'));
+
+    // Listen for the onDidChangeMcpServerDefinitions event.
+    // Mocha's test timeout (60 s) handles the failure case.
+    const changePromise = new Promise<void>((resolve) => {
+      const disposable = provider.onDidChangeMcpServerDefinitions(() => {
+        disposable.dispose();
+        resolve();
+      });
+    });
+
+    // Add the temporary folder to the workspace
+    const folders = vscode.workspace.workspaceFolders ?? [];
+    const added = vscode.workspace.updateWorkspaceFolders(
+      folders.length,
+      0,
+      { uri: vscode.Uri.file(tempDir) },
+    );
+
+    if (!added) {
+      // Clean up and skip if updateWorkspaceFolders is not supported in this profile
+      fs.rmdirSync(tempDir);
+      this.skip();
+      return;
+    }
+
+    try {
+      // Wait for the MCP definition change event
+      await changePromise;
+
+      // Get definitions after the workspace folder change
+      const afterDefs = await provider.provideMcpServerDefinitions(token);
+      assert.ok(afterDefs && afterDefs.length >= 1, 'Should still provide definitions after folder change');
+      const versionAfter = afterDefs[0].version;
+
+      // The version MUST be different so VS Code knows to restart the server
+      assert.notStrictEqual(
+        versionAfter,
+        versionBefore,
+        'MCP server definition version must change after workspace folder update ' +
+        'to signal VS Code to restart the server instead of only stopping it',
+      );
+    } finally {
+      // Cleanup: remove the added folder and temp directory
+      const updatedFolders = vscode.workspace.workspaceFolders ?? [];
+      const idx = updatedFolders.findIndex((f) => f.uri.fsPath === tempDir);
+      if (idx >= 0) {
+        vscode.workspace.updateWorkspaceFolders(idx, 1);
+      }
+      try {
+        fs.rmdirSync(tempDir);
+      } catch {
+        // Best-effort cleanup
+      }
+    }
+  });
+
+  test('MCP definition version should change after workspace folder is removed', async function () {
+    const provider = api.mcpProvider;
+    const folders = vscode.workspace.workspaceFolders;
+    if (!provider || !folders || folders.length < 2) {
+      // Need at least 2 folders to remove one without closing the workspace
+      this.skip();
+      return;
+    }
+
+    const token: vscode.CancellationToken = {
+      isCancellationRequested: false,
+      onCancellationRequested: () => ({ dispose: () => {} }),
+    };
+
+    // Get initial definitions
+    const beforeDefs = await provider.provideMcpServerDefinitions(token);
+    assert.ok(beforeDefs && beforeDefs.length >= 1, 'Should provide at least one definition');
+    const versionBefore = beforeDefs[0].version;
+
+    // Remember the last folder so we can re-add it after removal
+    const lastFolder = folders[folders.length - 1];
+    const removedUri = lastFolder.uri;
+
+    // Listen for the onDidChangeMcpServerDefinitions event.
+    // Mocha's test timeout (60 s) handles the failure case.
+    const changePromise = new Promise<void>((resolve) => {
+      const disposable = provider.onDidChangeMcpServerDefinitions(() => {
+        disposable.dispose();
+        resolve();
+      });
+    });
+
+    // Remove the last workspace folder
+    const removed = vscode.workspace.updateWorkspaceFolders(folders.length - 1, 1);
+    if (!removed) {
+      this.skip();
+      return;
+    }
+
+    try {
+      // Wait for the MCP definition change event
+      await changePromise;
+
+      // Get definitions after the workspace folder change
+      const afterDefs = await provider.provideMcpServerDefinitions(token);
+      assert.ok(afterDefs && afterDefs.length >= 1, 'Should still provide definitions after folder removal');
+      const versionAfter = afterDefs[0].version;
+
+      // The version MUST be different so VS Code knows to restart the server
+      assert.notStrictEqual(
+        versionAfter,
+        versionBefore,
+        'MCP server definition version must change after workspace folder removal ' +
+        'to signal VS Code to restart the server instead of only stopping it',
+      );
+    } finally {
+      // Cleanup: re-add the removed folder
+      const currentFolders = vscode.workspace.workspaceFolders ?? [];
+      vscode.workspace.updateWorkspaceFolders(
+        currentFolders.length,
+        0,
+        { uri: removedUri },
+      );
+    }
+  });
+
+  test('Environment should reflect updated workspace folders after change', async function () {
+    const envBuilder = api.environmentBuilder;
+    const folders = vscode.workspace.workspaceFolders;
+    if (!envBuilder || !folders || folders.length === 0) {
+      this.skip();
+      return;
+    }
+
+    // Create a real temporary directory
+    const tempDir = fs.mkdtempSync(path.join(os.tmpdir(), 'ql-mcp-env-test-'));
+
+    // Listen for workspace folder changes through the VS Code API
+    const changePromise = new Promise<void>((resolve) => {
+      const disposable = vscode.workspace.onDidChangeWorkspaceFolders(() => {
+        disposable.dispose();
+        resolve();
+      });
+    });
+
+    // Add the temporary folder
+    const added = vscode.workspace.updateWorkspaceFolders(
+      folders.length,
+      0,
+      { uri: vscode.Uri.file(tempDir) },
+    );
+
+    if (!added) {
+      fs.rmdirSync(tempDir);
+      this.skip();
+      return;
+    }
+
+    try {
+      await changePromise;
+
+      // invalidate() was called by the event handler, so build() returns fresh env
+      const env = await envBuilder.build();
+      assert.ok(env.CODEQL_MCP_WORKSPACE_FOLDERS, 'CODEQL_MCP_WORKSPACE_FOLDERS should be set');
+      assert.ok(
+        env.CODEQL_MCP_WORKSPACE_FOLDERS.includes(tempDir),
+        `CODEQL_MCP_WORKSPACE_FOLDERS should include the newly added folder: ${env.CODEQL_MCP_WORKSPACE_FOLDERS}`,
+      );
+    } finally {
+      // Cleanup
+      const updatedFolders = vscode.workspace.workspaceFolders ?? [];
+      const idx = updatedFolders.findIndex((f) => f.uri.fsPath === tempDir);
+      if (idx >= 0) {
+        vscode.workspace.updateWorkspaceFolders(idx, 1);
+      }
+      try {
+        fs.rmdirSync(tempDir);
+      } catch {
+        // Best-effort cleanup
+      }
+    }
+  });
+});