From 2dee4671afdabdf336b6123492e6d197d17d3b5c Mon Sep 17 00:00:00 2001
From: Johannes Rieken <johannes.rieken@gmail.com>
Date: Tue, 15 Apr 2025 11:38:26 +0200
Subject: [PATCH] fix wrong bucket assignment when extension implements tool
 and MCP server config (#246599)

---
 .../chat/browser/actions/chatToolActions.ts   | 19 +++++++------------
 1 file changed, 7 insertions(+), 12 deletions(-)

diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatToolActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatToolActions.ts
index ce4d1e6a6b6..b70f7e04a70 100644
--- a/src/vs/workbench/contrib/chat/browser/actions/chatToolActions.ts
+++ b/src/vs/workbench/contrib/chat/browser/actions/chatToolActions.ts
@@ -20,7 +20,6 @@ import { ExtensionIdentifier } from '../../../../../platform/extensions/common/e
 import { KeybindingWeight } from '../../../../../platform/keybinding/common/keybindingsRegistry.js';
 import { IQuickInputService, IQuickPickItem, IQuickPickSeparator } from '../../../../../platform/quickinput/common/quickInput.js';
 import { ITelemetryService } from '../../../../../platform/telemetry/common/telemetry.js';
-import { IExtensionService } from '../../../../services/extensions/common/extensions.js';
 import { IExtensionsWorkbenchService } from '../../../extensions/common/extensions.js';
 import { AddConfigurationAction } from '../../../mcp/browser/mcpCommands.js';
 import { IMcpService, IMcpServer, McpConnectionState } from '../../../mcp/common/mcpTypes.js';
@@ -111,7 +110,6 @@ export class AttachToolsAction extends Action2 {
 		const quickPickService = accessor.get(IQuickInputService);
 		const mcpService = accessor.get(IMcpService);
 		const toolsService = accessor.get(ILanguageModelToolsService);
-		const extensionService = accessor.get(IExtensionService);
 		const chatWidgetService = accessor.get(IChatWidgetService);
 		const telemetryService = accessor.get(ITelemetryService);
 		const commandService = accessor.get(ICommandService);
@@ -185,7 +183,8 @@ export class AttachToolsAction extends Action2 {
 				if (!mcpServer) {
 					continue;
 				}
-				bucket = toolBuckets.get(mcpServer.definition.id) ?? {
+				const key = tool.source.type + mcpServer.definition.id;
+				bucket = toolBuckets.get(key) ?? {
 					type: 'item',
 					label: localize('mcplabel', "MCP Server: {0}", mcpServer?.definition.label),
 					status: localize('mcpstatus', "From {0} ({1})", mcpServer.collection.label, McpConnectionState.toString(mcpServer.connectionState.get())),
@@ -194,23 +193,19 @@ export class AttachToolsAction extends Action2 {
 					picked: false,
 					children: []
 				};
-				toolBuckets.set(mcpServer.definition.id, bucket);
+				toolBuckets.set(key, bucket);
 			} else if (tool.source.type === 'extension') {
-				const extensionId = tool.source.extensionId;
-				const ext = extensionService.extensions.find(value => ExtensionIdentifier.equals(value.identifier, extensionId));
-				if (!ext) {
-					continue;
-				}
+				const key = tool.source.type + ExtensionIdentifier.toKey(tool.source.extensionId);
 
-				bucket = toolBuckets.get(ExtensionIdentifier.toKey(extensionId)) ?? {
+				bucket = toolBuckets.get(key) ?? {
 					type: 'item',
-					label: ext.displayName ?? ext.name,
+					label: tool.source.label,
 					ordinal: BucketOrdinal.Extension,
 					picked: false,
 					source: tool.source,
 					children: []
 				};
-				toolBuckets.set(ExtensionIdentifier.toKey(ext.identifier), bucket);
+				toolBuckets.set(key, bucket);
 			} else if (tool.source.type === 'internal') {
 				bucket = defaultBucket;
 			} else {