mcp: improve tool prefixing behavior (#249929)

Closes #249281

Uses `mcp_servername_toolname` to form the tool.

Author: connor4312

src/vs/workbench/contrib/mcp/common/mcpRegistry.ts

@@ -5,13 +5,11 @@
 
 import { Codicon } from '../../../../base/common/codicons.js';
 import { Emitter } from '../../../../base/common/event.js';
-import { StringSHA1 } from '../../../../base/common/hash.js';
 import { MarkdownString } from '../../../../base/common/htmlContent.js';
 import { Lazy } from '../../../../base/common/lazy.js';
 import { Disposable, IDisposable } from '../../../../base/common/lifecycle.js';
 import { derived, IObservable, observableValue } from '../../../../base/common/observable.js';
 import { basename } from '../../../../base/common/resources.js';
-import { indexOfPattern } from '../../../../base/common/strings.js';
 import { localize } from '../../../../nls.js';
 import { ConfigurationTarget, IConfigurationService } from '../../../../platform/configuration/common/configuration.js';
 import { IDialogService } from '../../../../platform/dialogs/common/dialogs.js';
@@ -37,8 +35,6 @@ const createTrustMemento = observableMemento<Readonly<Record<string, boolean>>>(
 	key: 'mcp.trustedCollections'
 });
 
-const collectionPrefixLen = 3;
-
 export class McpRegistry extends Disposable implements IMcpRegistry {
 	declare public readonly _serviceBrand: undefined;
 
@@ -54,41 +50,6 @@ export class McpRegistry extends Disposable implements IMcpRegistry {
 		return this._collections.read(reader);
 	});
 
-	private readonly _collectionToPrefixes = this._collections.map(c => {
-		// This creates tool prefixes based on a hash of the collection ID. This is
-		// a short prefix because tool names that are too long can cause errors (#243602).
-		// So we take a hash (in order for tools to be stable, because randomized
-		// names can cause hallicinations if present in history) and then adjust
-		// them if there are any collisions.
-		type CollectionHash = { view: number; hash: string; collection: McpCollectionDefinition };
-
-		const hashes = c.map((collection): CollectionHash => {
-			const sha = new StringSHA1();
-			sha.update(collection.id);
-			const hash = sha.digest();
-			// Gemini errors if the name starts with a number (microsoft/vscode-copilot-release#7152)
-			return { view: indexOfPattern(hash, /[a-z]/i), hash, collection };
-		});
-
-		const view = (h: CollectionHash) => h.hash.slice(h.view, h.view + collectionPrefixLen);
-
-		let collided = false;
-		do {
-			hashes.sort((a, b) => view(a).localeCompare(view(b)) || a.collection.id.localeCompare(b.collection.id));
-			collided = false;
-			for (let i = 1; i < hashes.length; i++) {
-				const prev = hashes[i - 1];
-				const curr = hashes[i];
-				if (view(prev) === view(curr) && curr.view + collectionPrefixLen < curr.hash.length) {
-					curr.view++;
-					collided = true;
-				}
-			}
-		} while (collided);
-
-		return Object.fromEntries(hashes.map(h => [h.collection.id, view(h) + '.']));
-	});
-
 	private readonly _workspaceStorage = new Lazy(() => this._register(this._instantiationService.createInstance(McpRegistryInputStorage, StorageScope.WORKSPACE, StorageTarget.USER)));
 	private readonly _profileStorage = new Lazy(() => this._register(this._instantiationService.createInstance(McpRegistryInputStorage, StorageScope.PROFILE, StorageTarget.USER)));
 
@@ -150,7 +111,8 @@ export class McpRegistry extends Disposable implements IMcpRegistry {
 		if (toReplace) {
 			this._collections.set(currentCollections.map(c => c === toReplace ? collection : c), undefined);
 		} else {
-			this._collections.set([...currentCollections, collection], undefined);
+			this._collections.set([...currentCollections, collection]
+				.sort((a, b) => (a.presentation?.order || 0) - (b.presentation?.order || 0)), undefined);
 		}
 
 		return {
@@ -169,10 +131,6 @@ export class McpRegistry extends Disposable implements IMcpRegistry {
 		});
 	}
 
-	public collectionToolPrefix(collection: McpCollectionReference): IObservable<string> {
-		return this._collectionToPrefixes.map(p => p[collection.id] ?? '');
-	}
-
 	public async discoverCollections(): Promise<McpCollectionDefinition[]> {
 		const toDiscover = this._collections.get().filter(c => c.lazy && !c.lazy.isCached);
 

src/vs/workbench/contrib/mcp/common/mcpRegistryTypes.ts

@@ -55,9 +55,6 @@ export interface IMcpRegistry {
 	/** Whether there are new collections that can be resolved with a discover() call */
 	readonly lazyCollectionState: IObservable<LazyCollectionState>;
 
-	/** Gets the prefix that should be applied to a collection's tools in order to avoid ID conflicts */
-	collectionToolPrefix(collection: McpCollectionReference): IObservable<string>;
-
 	/** Helper function to observe a definition by its reference. */
 	getServerDefinition(collectionRef: McpDefinitionReference, definitionRef: McpDefinitionReference): IObservable<{ server: McpServerDefinition | undefined; collection: McpCollectionDefinition | undefined }>;
 

src/vs/workbench/contrib/mcp/common/mcpServer.ts

@@ -29,7 +29,7 @@ import { mcpActivationEvent } from './mcpConfiguration.js';
 import { McpDevModeServerAttache } from './mcpDevMode.js';
 import { IMcpRegistry } from './mcpRegistryTypes.js';
 import { McpServerRequestHandler } from './mcpServerRequestHandler.js';
-import { extensionMcpCollectionPrefix, IMcpPrompt, IMcpPromptMessage, IMcpResource, IMcpServer, IMcpServerConnection, IMcpServerStartOpts, IMcpTool, McpCollectionDefinition, McpCollectionReference, McpConnectionFailedError, McpConnectionState, McpDefinitionReference, McpResourceURI, McpServerDefinition, McpServerCacheState, McpServerTransportType, McpCapability, IMcpResourceTemplate } from './mcpTypes.js';
+import { extensionMcpCollectionPrefix, IMcpPrompt, IMcpPromptMessage, IMcpResource, IMcpServer, IMcpServerConnection, IMcpServerStartOpts, IMcpTool, McpCollectionDefinition, McpCollectionReference, McpConnectionFailedError, McpConnectionState, McpDefinitionReference, McpResourceURI, McpServerDefinition, McpServerCacheState, McpServerTransportType, McpCapability, IMcpResourceTemplate, McpToolName } from './mcpTypes.js';
 import { MCP } from './modelContextProtocol.js';
 import { UriTemplate } from './uriTemplate.js';
 
@@ -266,7 +266,7 @@ export class McpServer extends Disposable implements IMcpServer {
 
 		const fromServer = this._tools.fromServerPromise.read(reader);
 		const connectionState = this.connectionState.read(reader);
-		const isIdle = McpConnectionState.canBeStarted(connectionState.state) && !fromServer;
+		const isIdle = McpConnectionState.canBeStarted(connectionState.state) || !fromServer;
 		if (isIdle) {
 			return stateWhenServingFromCache();
 		}
@@ -297,6 +297,7 @@ export class McpServer extends Disposable implements IMcpServer {
 		explicitRoots: URI[] | undefined,
 		private readonly _requiresExtensionActivation: boolean | undefined,
 		private readonly _primitiveCache: McpServerMetadataCache,
+		toolPrefix: string,
 		@IMcpRegistry private readonly _mcpRegistry: IMcpRegistry,
 		@IWorkspaceContextService workspacesService: IWorkspaceContextService,
 		@IExtensionService private readonly _extensionService: IExtensionService,
@@ -354,15 +355,11 @@ export class McpServer extends Disposable implements IMcpServer {
 		}));
 
 		// 3. Publish tools
-		const toolPrefix = this._mcpRegistry.collectionToolPrefix(this.collection);
 		this._tools = new CachedPrimitive<IMcpTool, IValidatedMcpTool>(
 			this.definition.id,
 			this._primitiveCache,
 			(entry) => entry.tools,
-			(entry, reader) => {
-				const prefix = toolPrefix.read(reader);
-				return entry.map(def => new McpTool(this, prefix, def)).sort((a, b) => a.compare(b));
-			},
+			(entry) => entry.map(def => new McpTool(this, toolPrefix, def)).sort((a, b) => a.compare(b)),
 		);
 
 		// 4. Publish promtps
@@ -708,7 +705,7 @@ export class McpTool implements IMcpTool {
 		idPrefix: string,
 		private readonly _definition: IValidatedMcpTool,
 	) {
-		this.id = (idPrefix + _definition.name).replaceAll('.', '_');
+		this.id = (idPrefix + _definition.name).replaceAll('.', '_').slice(0, McpToolName.MaxLength);
 	}
 
 	call(params: Record<string, unknown>, token?: CancellationToken): Promise<MCP.CallToolResult> {

src/vs/workbench/contrib/mcp/common/mcpService.ts

@@ -20,14 +20,14 @@ import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocatio
 import { McpCommandIds } from './mcpCommandIds.js';
 import { IMcpRegistry } from './mcpRegistryTypes.js';
 import { McpServer, McpServerMetadataCache } from './mcpServer.js';
-import { IMcpServer, IMcpService, IMcpTool, McpCollectionDefinition, McpServerCacheState, McpServerDefinition } from './mcpTypes.js';
+import { IMcpServer, IMcpService, IMcpTool, McpCollectionDefinition, McpServerCacheState, McpServerDefinition, McpToolName } from './mcpTypes.js';
 
 interface ISyncedToolData {
 	toolData: IToolData;
 	store: DisposableStore;
 }
 
-type IMcpServerRec = IReference<IMcpServer>;
+type IMcpServerRec = IReference<IMcpServer> & { toolPrefix: string };
 
 export class McpService extends Disposable implements IMcpService {
 
@@ -94,6 +94,15 @@ export class McpService extends Disposable implements IMcpService {
 
 		store.add(autorun(reader => {
 			const toDelete = new Set(tools.keys());
+
+			// toRegister is deferred until deleting tools that moving a tool between
+			// servers (or deleting one instance of a multi-instance server) doesn't cause an error.
+			const toRegister: (() => void)[] = [];
+			const registerTool = (tool: IMcpTool, toolData: IToolData, store: DisposableStore) => {
+				store.add(this._toolsService.registerToolData(toolData));
+				store.add(this._toolsService.registerToolImplementation(tool.id, this._instantiationService.createInstance(McpToolImplementation, tool, server)));
+			};
+
 			for (const tool of server.tools.read(reader)) {
 				const existing = tools.get(tool.id);
 				const collection = this._mcpRegistry.collections.get().find(c => c.id === server.collection.id);
@@ -102,7 +111,7 @@ export class McpService extends Disposable implements IMcpService {
 					source: { type: 'mcp', label: server.definition.label, collectionId: server.collection.id, definitionId: server.definition.id },
 					icon: Codicon.tools,
 					displayName: tool.definition.annotations?.title || tool.definition.name,
-					toolReferenceName: tool.definition.name,
+					toolReferenceName: tool.id,
 					modelDescription: tool.definition.description ?? '',
 					userDescription: tool.definition.description ?? '',
 					inputSchema: tool.definition.inputSchema,
@@ -112,26 +121,20 @@ export class McpService extends Disposable implements IMcpService {
 					tags: ['mcp'],
 				};
 
-				const registerTool = (store: DisposableStore) => {
-					store.add(this._toolsService.registerToolData(toolData));
-					store.add(this._toolsService.registerToolImplementation(tool.id, this._instantiationService.createInstance(McpToolImplementation, tool, server)));
-				};
-
 				if (existing) {
 					if (!equals(existing.toolData, toolData)) {
 						existing.toolData = toolData;
 						existing.store.clear();
 						// We need to re-register both the data and implementation, as the
 						// implementation is discarded when the data is removed (#245921)
-						registerTool(store);
+						registerTool(tool, toolData, store);
 					}
 					toDelete.delete(tool.id);
 				} else {
 					const store = new DisposableStore();
-					registerTool(store);
+					toRegister.push(() => registerTool(tool, toolData, store));
 					tools.set(tool.id, { toolData, store });
 				}
-
 			}
 
 			for (const id of toDelete) {
@@ -141,6 +144,10 @@ export class McpService extends Disposable implements IMcpService {
 					tools.delete(id);
 				}
 			}
+
+			for (const fn of toRegister) {
+				fn();
+			}
 		}));
 
 		store.add(toDisposable(() => {
@@ -151,11 +158,12 @@ export class McpService extends Disposable implements IMcpService {
 	}
 
 	public updateCollectedServers() {
+		const prefixGenerator = new McpPrefixGenerator();
 		const definitions = this._mcpRegistry.collections.get().flatMap(collectionDefinition =>
-			collectionDefinition.serverDefinitions.get().map(serverDefinition => ({
-				serverDefinition,
-				collectionDefinition,
-			}))
+			collectionDefinition.serverDefinitions.get().map(serverDefinition => {
+				const toolPrefix = prefixGenerator.generate(serverDefinition.label);
+				return { serverDefinition, collectionDefinition, toolPrefix };
+			})
 		);
 
 		const nextDefinitions = new Set(definitions);
@@ -174,7 +182,7 @@ export class McpService extends Disposable implements IMcpService {
 
 		// Transfer over any servers that are still valid.
 		for (const server of currentServers) {
-			const match = definitions.find(d => defsEqual(server.object, d));
+			const match = definitions.find(d => defsEqual(server.object, d) && server.toolPrefix === d.toolPrefix);
 			if (match) {
 				pushMatch(match, server);
 			} else {
@@ -192,12 +200,13 @@ export class McpService extends Disposable implements IMcpService {
 				def.serverDefinition.roots,
 				!!def.collectionDefinition.lazy,
 				def.collectionDefinition.scope === StorageScope.WORKSPACE ? this.workspaceCache : this.userCache,
+				def.toolPrefix,
 			);
 
 			store.add(object);
 			this._syncTools(object, store);
 
-			nextServers.push({ object, dispose: () => store.dispose() });
+			nextServers.push({ object, dispose: () => store.dispose(), toolPrefix: def.toolPrefix });
 		}
 
 		transaction(tx => {
@@ -292,3 +301,18 @@ class McpToolImplementation implements IToolImpl {
 		return result;
 	}
 }
+
+// Helper class for generating unique MCP tool prefixes
+class McpPrefixGenerator {
+	private readonly seenPrefixes = new Set<string>();
+
+	generate(label: string): string {
+		const baseToolPrefix = McpToolName.Prefix + label.toLowerCase().replace(/[^a-z0-9_.-]+/g, '_').slice(0, McpToolName.MaxPrefixLen - McpToolName.Prefix.length - 1);
+		let toolPrefix = baseToolPrefix + '_';
+		for (let i = 2; this.seenPrefixes.has(toolPrefix); i++) {
+			toolPrefix = baseToolPrefix + i + '_';
+		}
+		this.seenPrefixes.add(toolPrefix);
+		return toolPrefix;
+	}
+}

src/vs/workbench/contrib/mcp/common/mcpTypes.ts

@@ -639,3 +639,9 @@ export const enum McpCapability {
 	Tools = 1 << 7,
 	ToolsListChanged = 1 << 8,
 }
+
+export const enum McpToolName {
+	Prefix = 'mcp_',
+	MaxPrefixLen = 18,
+	MaxLength = 64,
+}