From 7881184f6e7660c9f798feebecf91b8da5978363 Mon Sep 17 00:00:00 2001 From: Dmitriy Vasyura Date: Sun, 20 Sep 2026 12:32:43 -0700 Subject: [PATCH] mcp: skip startup for removed server definitions Treat a collection or definition removed during discovery reconciliation as a cancelled start instead of throwing. Log the skipped startup and preserve genuine lazy-load, delegate, and enterprise policy errors. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../contrib/mcp/common/mcpRegistry.ts | 3 +- .../contrib/mcp/common/mcpRegistryTypes.ts | 2 +- .../mcp/test/common/mcpRegistry.test.ts | 94 ++++++++++++++++++- 3 files changed, 96 insertions(+), 3 deletions(-) diff --git a/src/vs/workbench/contrib/mcp/common/mcpRegistry.ts b/src/vs/workbench/contrib/mcp/common/mcpRegistry.ts index fef5ad9ec09941..c1de3ae5a8d5be 100644 --- a/src/vs/workbench/contrib/mcp/common/mcpRegistry.ts +++ b/src/vs/workbench/contrib/mcp/common/mcpRegistry.ts @@ -516,7 +516,8 @@ export class McpRegistry extends Disposable implements IMcpRegistry { const definition = collection?.serverDefinitions.get().find(s => s.id === definitionRef.id); if (!collection || !definition) { - throw new Error(`Collection or definition not found for ${collectionRef.id} and ${definitionRef.id}`); + logger.debug(`Skipping MCP server ${definitionRef.id}: collection ${collectionRef.id} or server definition is no longer registered.`); + return undefined; } const delegate = this._delegates.get().find(d => d.canStart(collection, definition)); diff --git a/src/vs/workbench/contrib/mcp/common/mcpRegistryTypes.ts b/src/vs/workbench/contrib/mcp/common/mcpRegistryTypes.ts index 6b589320fa67fd..a69b01ff37046e 100644 --- a/src/vs/workbench/contrib/mcp/common/mcpRegistryTypes.ts +++ b/src/vs/workbench/contrib/mcp/common/mcpRegistryTypes.ts @@ -95,6 +95,6 @@ export interface IMcpRegistry { setSavedInput(inputId: string, target: ConfigurationTarget, value: string): Promise; /** Gets saved inputs from storage. */ getSavedInputs(scope: StorageScope): Promise<{ [id: string]: IResolvedValue }>; - /** Creates a connection for the collection and definition. */ + /** Creates a connection, or returns undefined if startup is cancelled or the collection or definition is no longer registered. */ resolveConnection(options: IMcpResolveConnectionOptions): Promise; } diff --git a/src/vs/workbench/contrib/mcp/test/common/mcpRegistry.test.ts b/src/vs/workbench/contrib/mcp/test/common/mcpRegistry.test.ts index c80059ef3edd89..224cc3064a8ca5 100644 --- a/src/vs/workbench/contrib/mcp/test/common/mcpRegistry.test.ts +++ b/src/vs/workbench/contrib/mcp/test/common/mcpRegistry.test.ts @@ -5,7 +5,7 @@ import * as assert from 'assert'; import * as sinon from 'sinon'; -import { timeout } from '../../../../../base/common/async.js'; +import { DeferredPromise, timeout } from '../../../../../base/common/async.js'; import { Disposable, toDisposable } from '../../../../../base/common/lifecycle.js'; import { ISettableObservable, observableValue } from '../../../../../base/common/observable.js'; import { URI } from '../../../../../base/common/uri.js'; @@ -257,6 +257,10 @@ suite('Workbench - MCP - Registry', () => { }; }); + teardown(() => { + sinon.restore(); + }); + test('registerCollection adds collection to registry', () => { const disposable = registry.registerCollection(testCollection); store.add(disposable); @@ -330,6 +334,94 @@ suite('Workbench - MCP - Registry', () => { assert.strictEqual(registry.delegates.get().length, 0); }); + for (const removed of ['collection', 'definition'] as const) { + for (const duringLoad of [false, true]) { + test(`resolveConnection skips a removed ${removed} ${duringLoad ? 'during lazy loading' : 'before resolution'}`, async () => { + const loaded = new DeferredPromise(); + const resolveLaunch = sinon.stub().resolves(baseDefinition.launch); + const resolveInputs = sinon.spy(testConfigResolverService, 'resolveWithInteraction'); + const delegate = new TestMcpHostDelegate(); + const substituteVariables = sinon.spy(delegate, 'substituteVariables'); + store.add(registry.registerDelegate(delegate)); + testCollection.serverDefinitions.set([baseDefinition], undefined); + const collection: McpCollectionDefinition = { + ...testCollection, + resolveServerLanch: resolveLaunch, + lazy: duringLoad ? { isCached: true, load: () => loaded.p } : undefined, + }; + const registration = store.add(registry.registerCollection(collection)); + const resolve = () => registry.resolveConnection({ collectionRef: collection, definitionRef: baseDefinition, logger, trustNonceBearer, taskManager }); + const pending = duringLoad ? resolve() : undefined; + + if (removed === 'collection') { + registration.dispose(); + } else { + testCollection.serverDefinitions.set([], undefined); + } + await loaded.complete(); + + const connection = await (pending ?? resolve()); + if (connection) { + store.add(connection); + } + assert.deepStrictEqual({ + connection, + launchResolutions: resolveLaunch.callCount, + variableSubstitutions: substituteVariables.callCount, + inputResolutions: resolveInputs.callCount, + sandboxLaunches: testMcpSandboxService.callCount, + }, { + connection: undefined, + launchResolutions: 0, + variableSubstitutions: 0, + inputResolutions: 0, + sandboxLaunches: 0, + }); + }); + } + } + + test('resolveConnection still rejects a lazy loading error', async () => { + const error = new Error('Failed to load the MCP collection'); + const collection: McpCollectionDefinition = { + ...testCollection, + lazy: { isCached: true, load: async () => { throw error; } }, + }; + store.add(registry.registerCollection(collection)); + + await assert.rejects( + registry.resolveConnection({ collectionRef: collection, definitionRef: baseDefinition, logger, trustNonceBearer, taskManager }), + error, + ); + }); + + test('resolveConnection still rejects a missing delegate', async () => { + testCollection.serverDefinitions.set([baseDefinition], undefined); + store.add(registry.registerCollection(testCollection)); + + await assert.rejects( + registry.resolveConnection({ collectionRef: testCollection, definitionRef: baseDefinition, logger, trustNonceBearer, taskManager }), + /No delegate found that can handle the connection/, + ); + }); + + test('resolveConnection still rejects an enterprise customization restriction', async () => { + testCollection.serverDefinitions.set([baseDefinition], undefined); + store.add(registry.registerCollection(testCollection)); + await configurationService.setUserConfiguration(COPILOT_STRICT_PLUGIN_ONLY_CUSTOMIZATION_CONFIG, true); + configurationService.onDidChangeConfigurationEmitter.fire({ + source: ConfigurationTarget.USER, + affectedKeys: new Set([COPILOT_STRICT_PLUGIN_ONLY_CUSTOMIZATION_CONFIG]), + change: { keys: [COPILOT_STRICT_PLUGIN_ONLY_CUSTOMIZATION_CONFIG], overrides: [] }, + affectsConfiguration: key => key === COPILOT_STRICT_PLUGIN_ONLY_CUSTOMIZATION_CONFIG, + }); + + await assert.rejects( + registry.resolveConnection({ collectionRef: testCollection, definitionRef: baseDefinition, logger, trustNonceBearer, taskManager }), + /MCP collection test-collection is blocked by enterprise customization policy/, + ); + }); + test('resolveConnection creates connection with resolved variables and memorizes them until cleared', async () => { const definition: McpServerDefinition = { ...baseDefinition,