diff --git a/extensions/vscode/esbuild.config.js b/extensions/vscode/esbuild.config.js index 3a0682a9..ee55e4c1 100644 --- a/extensions/vscode/esbuild.config.js +++ b/extensions/vscode/esbuild.config.js @@ -41,6 +41,7 @@ const testSuiteConfig = { 'test/suite/mcp-resource-e2e.integration.test.ts', 'test/suite/mcp-server.integration.test.ts', 'test/suite/mcp-tool-e2e.integration.test.ts', + 'test/suite/workspace-folder-change.integration.test.ts', 'test/suite/workspace-scenario.integration.test.ts', ], outdir: 'dist/test/suite', diff --git a/extensions/vscode/src/extension.ts b/extensions/vscode/src/extension.ts index 7f2febbe..92af9897 100644 --- a/extensions/vscode/src/extension.ts +++ b/extensions/vscode/src/extension.ts @@ -81,19 +81,20 @@ export async function activate( context.subscriptions.push( vscode.workspace.onDidChangeConfiguration((e) => { if (e.affectsConfiguration('codeql-mcp')) { - logger.info('Configuration changed — rebuilding environment'); - envBuilder.invalidate(); - mcpProvider.fireDidChange(); + logger.info('Configuration changed — requesting MCP server restart'); + mcpProvider.requestRestart(); } }), ); - // Re-scan when workspace folders change + // Invalidate cached environment when workspace folders change. + // VS Code itself manages MCP server lifecycle when roots change + // (stopping and restarting the server as needed). We just clear + // the cached env so the next server start picks up updated folders. context.subscriptions.push( vscode.workspace.onDidChangeWorkspaceFolders(() => { - logger.info('Workspace folders changed — refreshing watchers'); + logger.info('Workspace folders changed — invalidating environment cache'); envBuilder.invalidate(); - mcpProvider.fireDidChange(); }), ); diff --git a/extensions/vscode/src/server/mcp-provider.ts b/extensions/vscode/src/server/mcp-provider.ts index 75086655..de30ea29 100644 --- a/extensions/vscode/src/server/mcp-provider.ts +++ b/extensions/vscode/src/server/mcp-provider.ts @@ -20,27 +20,69 @@ export class McpProvider private readonly _onDidChange = new vscode.EventEmitter(); readonly onDidChangeMcpServerDefinitions = this._onDidChange.event; + /** + * Monotonically increasing counter, bumped by `requestRestart()`. + * Appended to the version string so that VS Code sees a genuinely new + * server definition and triggers a stop → restart cycle. + */ + private _revision = 0; + + /** + * Cached extension version to avoid repeated synchronous `readFileSync` + * calls on every `getEffectiveVersion()` invocation. + */ + private readonly _extensionVersion: string; + constructor( private readonly serverManager: ServerManager, private readonly envBuilder: EnvironmentBuilder, private readonly logger: Logger, ) { super(); + this._extensionVersion = serverManager.getExtensionVersion(); this.push(this._onDidChange); } - /** Notify VS Code that the MCP server definitions have changed. */ + /** + * Soft notification: tell VS Code that definitions may have changed. + * + * Does NOT bump the version, so VS Code will re-query + * `provideMcpServerDefinitions()` / `resolveMcpServerDefinition()` but + * will NOT restart the server. Use for lightweight updates (file watcher + * events, extension changes, background install completion) where the + * running server can continue with its current environment. + */ fireDidChange(): void { this._onDidChange.fire(); } + /** + * Request that VS Code restart the MCP server with a fresh environment. + * + * Invalidates the cached environment and bumps the internal revision counter + * so that the next call to `provideMcpServerDefinitions()` returns a + * definition with a different `version` string. VS Code compares the new + * version to the running server's version and, seeing a change, triggers a + * stop → start cycle. + * + * Use for changes that require a server restart (configuration changes). + */ + requestRestart(): void { + this.envBuilder.invalidate(); + this._revision++; + this.logger.info( + `Requesting ql-mcp restart (revision ${this._revision})...`, + ); + this._onDidChange.fire(); + } + async provideMcpServerDefinitions( _token: vscode.CancellationToken, ): Promise { 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 +108,26 @@ export class McpProvider server.env = env; return server; } + + /** + * Computes the version string for the MCP server definition. + * + * Always returns a concrete string so that VS Code has a reliable + * baseline for version comparison. When `serverManager.getVersion()` + * returns `undefined` (the "latest" / unpinned case), the extension + * version is used as the base instead. + * + * After one or more `requestRestart()` calls, a `+rN` revision suffix + * is appended so that the version is always different from the + * previous one. VS Code uses the version to decide whether to + * restart a running server: a changed version triggers a stop → start + * cycle. + */ + private getEffectiveVersion(): string { + const base = this.serverManager.getVersion() ?? this._extensionVersion; + if (this._revision === 0) { + return base; + } + return `${base}+r${this._revision}`; + } } diff --git a/extensions/vscode/test/extension.test.ts b/extensions/vscode/test/extension.test.ts index b0f3b400..e3d53ab2 100644 --- a/extensions/vscode/test/extension.test.ts +++ b/extensions/vscode/test/extension.test.ts @@ -50,6 +50,7 @@ vi.mock('../src/server/mcp-provider', () => ({ resolveMcpServerDefinition: vi.fn().mockResolvedValue(undefined), onDidChangeMcpServerDefinitions: vi.fn(), fireDidChange: vi.fn(), + requestRestart: vi.fn(), dispose: vi.fn(), }; }), @@ -102,6 +103,8 @@ vi.mock('../src/bridge/environment-builder', () => ({ import { activate, deactivate } from '../src/extension'; import * as vscode from 'vscode'; +import { McpProvider } from '../src/server/mcp-provider'; +import { EnvironmentBuilder } from '../src/bridge/environment-builder'; function createMockContext(): vscode.ExtensionContext { return { @@ -168,4 +171,39 @@ describe('Extension', () => { it('should deactivate even if activate was never called', () => { expect(() => deactivate()).not.toThrow(); }); + + it('should invalidate env cache on workspace folder changes', async () => { + // Capture the workspace folder change callback + let workspaceFolderChangeCallback: Function | undefined; + const spy = vi.spyOn(vscode.workspace, 'onDidChangeWorkspaceFolders').mockImplementation( + (cb: Function) => { workspaceFolderChangeCallback = cb; return { dispose: vi.fn() }; }, + ); + + await activate(ctx); + + // The extension should have registered a workspace folder change handler + expect(workspaceFolderChangeCallback).toBeDefined(); + + // Get the mock instances created during activation + const mcpProviderInstance = vi.mocked(McpProvider).mock.results[0]?.value; + const envBuilderInstance = vi.mocked(EnvironmentBuilder).mock.results[0]?.value; + + // Reset call counts from activation + mcpProviderInstance.fireDidChange.mockClear(); + mcpProviderInstance.requestRestart.mockClear(); + envBuilderInstance.invalidate.mockClear(); + + try { + // Simulate workspace folder change + workspaceFolderChangeCallback!(); + + // Cache invalidated immediately; no VS Code notification. + // VS Code manages MCP server lifecycle (stop/restart) when roots change. + expect(envBuilderInstance.invalidate).toHaveBeenCalledTimes(1); + expect(mcpProviderInstance.fireDidChange).not.toHaveBeenCalled(); + expect(mcpProviderInstance.requestRestart).not.toHaveBeenCalled(); + } finally { + spy.mockRestore(); + } + }); }); diff --git a/extensions/vscode/test/fixtures/multi-root-workspace/folder-c/.gitkeep b/extensions/vscode/test/fixtures/multi-root-workspace/folder-c/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/extensions/vscode/test/fixtures/multi-root-workspace/folder-d/.gitkeep b/extensions/vscode/test/fixtures/multi-root-workspace/folder-d/.gitkeep new file mode 100644 index 00000000..e69de29b diff --git a/extensions/vscode/test/fixtures/multi-root-workspace/test.code-workspace b/extensions/vscode/test/fixtures/multi-root-workspace/test.code-workspace index ba155da6..fcbbbe8b 100644 --- a/extensions/vscode/test/fixtures/multi-root-workspace/test.code-workspace +++ b/extensions/vscode/test/fixtures/multi-root-workspace/test.code-workspace @@ -1,6 +1,16 @@ { "folders": [ - { "path": "folder-a" }, - { "path": "folder-b" } + { + "path": "folder-a" + }, + { + "path": "folder-b" + }, + { + "path": "folder-c" + }, + { + "path": "folder-d" + } ] } diff --git a/extensions/vscode/test/server/mcp-provider.test.ts b/extensions/vscode/test/server/mcp-provider.test.ts index 77c09f69..a66c5aa3 100644 --- a/extensions/vscode/test/server/mcp-provider.test.ts +++ b/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'), @@ -111,6 +112,26 @@ describe('McpProvider', () => { expect(definitions).toHaveLength(1); }); + it('should always provide a defined version string even when serverVersion is latest', async () => { + // When serverManager.getVersion() returns undefined ("latest" mode), + // the definition must still carry a concrete version string so that + // VS Code has a baseline for version comparison. An undefined initial + // version prevents VS Code from detecting changes after requestRestart(). + // + // NOTE: McpProvider caches getExtensionVersion() in its constructor, so the + // mock must be configured before constructing the provider. + serverManager.getVersion.mockReturnValue(undefined); + serverManager.getExtensionVersion.mockReturnValue('2.25.1'); + provider = new McpProvider(serverManager, envBuilder, logger); + + const token = { isCancellationRequested: false, onCancellationRequested: vi.fn() }; + const definitions = await provider.provideMcpServerDefinitions(token as any); + + expect(definitions![0].version).toBeDefined(); + expect(typeof definitions![0].version).toBe('string'); + expect(definitions![0].version!.length).toBeGreaterThan(0); + }); + it('should pass version from serverManager when pinned', async () => { serverManager.getVersion.mockReturnValue('2.20.0'); @@ -120,6 +141,78 @@ describe('McpProvider', () => { expect(definitions![0].version).toBe('2.20.0'); }); + it('should not change version on fireDidChange (soft signal)', 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); + + expect(after![0].version).toBe(versionBefore); + }); + + it('should produce a different version after requestRestart', async () => { + const token = { isCancellationRequested: false, onCancellationRequested: vi.fn() }; + const before = await provider.provideMcpServerDefinitions(token as any); + const versionBefore = before![0].version; + + provider.requestRestart(); + 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 requestRestart calls', async () => { + const token = { isCancellationRequested: false, onCancellationRequested: vi.fn() }; + + provider.requestRestart(); + const defs1 = await provider.provideMcpServerDefinitions(token as any); + + provider.requestRestart(); + 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 requestRestart', async () => { + serverManager.getVersion.mockReturnValue('2.20.0'); + const token = { isCancellationRequested: false, onCancellationRequested: vi.fn() }; + + provider.requestRestart(); + const definitions = await provider.provideMcpServerDefinitions(token as any); + + expect(definitions![0].version).toMatch(/^2\.20\.0\+r\d+$/); + }); + + it('should append revision to extension version after requestRestart when version is latest', async () => { + serverManager.getVersion.mockReturnValue(undefined); + serverManager.getExtensionVersion.mockReturnValue('2.25.1'); + const token = { isCancellationRequested: false, onCancellationRequested: vi.fn() }; + + provider.requestRestart(); + const definitions = await provider.provideMcpServerDefinitions(token as any); + + expect(definitions![0].version).toMatch(/^2\.25\.1\+r\d+$/); + }); + + it('should fire change event when requestRestart is called', () => { + const listener = vi.fn(); + provider.onDidChangeMcpServerDefinitions(listener); + + provider.requestRestart(); + + expect(listener).toHaveBeenCalledTimes(1); + }); + + it('should invalidate env cache when requestRestart is called', () => { + provider.requestRestart(); + + expect(envBuilder.invalidate).toHaveBeenCalledTimes(1); + }); + it('should resolve definition by refreshing environment', async () => { const updatedEnv = { CODEQL_PATH: '/new/path', TRANSPORT_MODE: 'stdio' }; envBuilder.build.mockResolvedValue(updatedEnv); diff --git a/extensions/vscode/test/suite/workspace-folder-change.integration.test.ts b/extensions/vscode/test/suite/workspace-folder-change.integration.test.ts new file mode 100644 index 00000000..1e9afa57 --- /dev/null +++ b/extensions/vscode/test/suite/workspace-folder-change.integration.test.ts @@ -0,0 +1,178 @@ +/** + * Integration tests for workspace folder change handling. + * + * These run inside the Extension Development Host with the REAL VS Code API + * under the `multiRoot` profile (4 workspace folders: folder-a … folder-d). + * + * They verify that: + * 1. The cached environment is rebuilt to include newly added folders. + * 2. The cached environment is rebuilt to exclude removed folders. + * 3. Original fixture folders survive after workspace mutations. + * + * Note: VS Code itself manages the MCP server lifecycle (stop/start) when + * workspace roots change. Our extension invalidates its env cache so the + * next server start picks up the updated workspace folder list. + */ + +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'; + +/** + * Create a temporary directory and resolve its real path (macOS `/tmp` is a + * symlink to `/private/tmp`; VS Code resolves symlinks in `uri.fsPath`). + */ +function createTempDir(prefix: string): string { + return fs.realpathSync(fs.mkdtempSync(path.join(os.tmpdir(), prefix))); +} + +/** Wait for `onDidChangeWorkspaceFolders` to fire once. */ +function waitForWorkspaceFolderChange(): Promise { + return new Promise((resolve) => { + const d = vscode.workspace.onDidChangeWorkspaceFolders(() => { + d.dispose(); + resolve(); + }); + }); +} + +/** + * Add a folder to the workspace and wait for the change event to settle. + * Throws if `updateWorkspaceFolders` returns `false`. + */ +async function addWorkspaceFolder(uri: vscode.Uri): Promise { + const folders = vscode.workspace.workspaceFolders ?? []; + const changePromise = waitForWorkspaceFolderChange(); + const ok = vscode.workspace.updateWorkspaceFolders(folders.length, 0, { uri }); + assert.ok(ok, 'updateWorkspaceFolders (add) returned false'); + await changePromise; +} + +/** + * Remove a workspace folder by its `fsPath` and wait for the change event. + * No-op if the folder is not currently in the workspace. + */ +async function removeWorkspaceFolder(fsPath: string): Promise { + const folders = vscode.workspace.workspaceFolders ?? []; + const idx = folders.findIndex((f) => f.uri.fsPath === fsPath); + if (idx < 0) return; + const changePromise = waitForWorkspaceFolderChange(); + const ok = vscode.workspace.updateWorkspaceFolders(idx, 1); + assert.ok(ok, 'updateWorkspaceFolders (remove) returned false'); + await changePromise; +} + +suite('Workspace Folder Change Tests', () => { + let api: any; + + suiteSetup(async function () { + const ext = vscode.extensions.getExtension(EXTENSION_ID); + assert.ok(ext, `Extension ${EXTENSION_ID} not found`); + api = ext.isActive ? ext.exports : await ext.activate(); + assert.ok(api.environmentBuilder, 'API missing environmentBuilder'); + + const folders = vscode.workspace.workspaceFolders; + console.log( + `[workspace-folder-change] Starting with ${folders?.length ?? 0} workspace folders`, + ); + if (!folders || folders.length < 2) { + console.log('[workspace-folder-change] Skipping — requires multiRoot workspace (>= 2 folders)'); + this.skip(); + } + }); + + // --------------------------------------------------------------- + // Test 1: Environment includes the newly added folder + // --------------------------------------------------------------- + test('CODEQL_MCP_WORKSPACE_FOLDERS should include newly added folder', async () => { + const envBuilder = api.environmentBuilder; + const tempDir = createTempDir('ql-mcp-env-'); + + // Populate the environment cache before mutating workspace folders. + const envBefore = await envBuilder.build(); + const beforeFoldersEnv = + envBefore.CODEQL_MCP_WORKSPACE_FOLDERS?.split(path.delimiter) ?? []; + assert.ok( + !beforeFoldersEnv.includes(tempDir), + `Precondition failed: tempDir ${tempDir} unexpectedly present in CODEQL_MCP_WORKSPACE_FOLDERS before add`, + ); + + await addWorkspaceFolder(vscode.Uri.file(tempDir)); + + try { + const envAfter = await envBuilder.build(); + assert.ok( + envAfter.CODEQL_MCP_WORKSPACE_FOLDERS, + 'CODEQL_MCP_WORKSPACE_FOLDERS should be set', + ); + const workspaceFoldersEnv = envAfter.CODEQL_MCP_WORKSPACE_FOLDERS.split(path.delimiter); + assert.ok( + workspaceFoldersEnv.includes(tempDir), + `CODEQL_MCP_WORKSPACE_FOLDERS should include ${tempDir}: ${envAfter.CODEQL_MCP_WORKSPACE_FOLDERS}`, + ); + } finally { + await removeWorkspaceFolder(tempDir); + try { fs.rmSync(tempDir, { recursive: true, force: true }); } catch { /* best-effort */ } + } + }); + + // --------------------------------------------------------------- + // Test 2: Environment excludes a removed folder + // --------------------------------------------------------------- + test('CODEQL_MCP_WORKSPACE_FOLDERS should exclude removed folder', async () => { + const envBuilder = api.environmentBuilder; + const tempDir = createTempDir('ql-mcp-env-rm-'); + + await addWorkspaceFolder(vscode.Uri.file(tempDir)); + + // Verify it's there + let env = await envBuilder.build(); + assert.ok( + env.CODEQL_MCP_WORKSPACE_FOLDERS.split(path.delimiter).includes(tempDir), + 'Folder should be present after adding', + ); + + // Now remove it + await removeWorkspaceFolder(tempDir); + + try { + env = await envBuilder.build(); + assert.ok( + !env.CODEQL_MCP_WORKSPACE_FOLDERS.split(path.delimiter).includes(tempDir), + `Folder should be absent after removal: ${env.CODEQL_MCP_WORKSPACE_FOLDERS}`, + ); + } finally { + await removeWorkspaceFolder(tempDir); + try { fs.rmSync(tempDir, { recursive: true, force: true }); } catch { /* best-effort */ } + } + }); + + // --------------------------------------------------------------- + // Test 3: Original fixture folders survive after add + remove + // --------------------------------------------------------------- + test('Fixture workspace folders should remain after add and remove cycle', async () => { + const foldersBefore = (vscode.workspace.workspaceFolders ?? []).map( + (f) => f.uri.fsPath, + ); + + const tempDir = createTempDir('ql-mcp-preserve-'); + await addWorkspaceFolder(vscode.Uri.file(tempDir)); + await removeWorkspaceFolder(tempDir); + + const foldersAfter = (vscode.workspace.workspaceFolders ?? []).map( + (f) => f.uri.fsPath, + ); + + assert.deepStrictEqual( + foldersAfter, + foldersBefore, + 'Original workspace folders must be preserved after an add+remove cycle', + ); + + try { fs.rmSync(tempDir, { recursive: true, force: true }); } catch { /* best-effort */ } + }); +});