fix: ql-mcp server must handle vscode workspace folder changes

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