Fix enablement of MCP tool sets and tools (#255747)

fix enablement of mcp tool sets and tools

Author: aeschli

src/vs/workbench/contrib/chat/browser/actions/chatToolPicker.ts

@@ -6,7 +6,6 @@ import { assertNever } from '../../../../../base/common/assert.js';
 import { Codicon } from '../../../../../base/common/codicons.js';
 import { diffSets } from '../../../../../base/common/collections.js';
 import { Event } from '../../../../../base/common/event.js';
-import { Iterable } from '../../../../../base/common/iterator.js';
 import { DisposableStore } from '../../../../../base/common/lifecycle.js';
 import { ThemeIcon } from '../../../../../base/common/themables.js';
 import { assertType } from '../../../../../base/common/types.js';
@@ -108,7 +107,9 @@ export async function showToolsPicker(
 	if (!toolsEntries) {
 		const defaultEntries = new Map();
 		for (const tool of toolsService.getTools()) {
-			defaultEntries.set(tool, false);
+			if (tool.canBeReferencedInPrompt) {
+				defaultEntries.set(tool, false);
+			}
 		}
 		for (const toolSet of toolsService.toolSets.get()) {
 			defaultEntries.set(toolSet, false);
@@ -203,8 +204,8 @@ export async function showToolsPicker(
 			} else {
 				// stash the MCP toolset into the bucket item
 				bucket.toolset = toolSetOrTool;
+				bucket.picked = picked;
 			}
-
 		} else if (toolSetOrTool.canBeReferencedInPrompt) {
 			bucket.children.push({
 				parent: bucket,
@@ -216,10 +217,6 @@ export async function showToolsPicker(
 				indented: true,
 			});
 		}
-
-		if (picked) {
-			bucket.picked = true;
-		}
 	}
 
 	for (const bucket of [builtinBucket, userBucket]) {
@@ -228,6 +225,21 @@ export async function showToolsPicker(
 		}
 	}
 
+	// set the checkmarks in the UI:
+	// bucket is checked if at least one of the children is checked
+	// tool is checked if the bucket is checked or the tool itself is checked
+	for (const bucket of toolBuckets.values()) {
+		if (bucket.picked) {
+			// check all children if the bucket is checked
+			for (const child of bucket.children) {
+				child.picked = true;
+			}
+		} else {
+			// check the bucket if one of the children is checked
+			bucket.picked = bucket.children.some(child => child.picked);
+		}
+	}
+
 	const store = new DisposableStore();
 
 	const picks: (MyPick | IQuickPickSeparator)[] = [];
@@ -385,22 +397,12 @@ export async function showToolsPicker(
 
 	store.dispose();
 
-	const mcpToolSets = new Set<ToolSet>();
-
+	// in the result, a MCP toolset is only enabled if all tools in the toolset are enabled
 	for (const item of toolsService.toolSets.get()) {
 		if (item.source.type === 'mcp') {
-			mcpToolSets.add(item);
-
-			if (Iterable.every(item.getTools(), tool => result.get(tool))) {
-				// ALL tools from the MCP tool set are here, replace them with just the toolset
-				// but only when computing the final result
-				for (const tool of item.getTools()) {
-					result.delete(tool);
-				}
-				result.set(item, true);
-			}
+			const toolsInSet = Array.from(item.getTools());
+			result.set(item, toolsInSet.every(tool => result.get(tool)));
 		}
 	}
-
 	return didAccept ? result : undefined;
 }

src/vs/workbench/contrib/chat/browser/languageModelToolsService.ts

@@ -492,11 +492,23 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 	toToolAndToolSetEnablementMap(enabledToolOrToolSetNames: readonly string[] | undefined): Map<ToolSet | IToolData, boolean> {
 		const toolOrToolSetNames = enabledToolOrToolSetNames ? new Set(enabledToolOrToolSetNames) : undefined;
 		const result = new Map<ToolSet | IToolData, boolean>();
-		for (const tool of this._tools.values()) {
-			result.set(tool.data, tool.data.toolReferenceName !== undefined && (toolOrToolSetNames === undefined || toolOrToolSetNames.has(tool.data.toolReferenceName)));
+		for (const tool of this.getTools()) {
+			if (tool.canBeReferencedInPrompt) {
+				result.set(tool, toolOrToolSetNames === undefined || toolOrToolSetNames.has(tool.toolReferenceName ?? tool.displayName));
+			}
 		}
 		for (const toolSet of this._toolSets) {
-			result.set(toolSet, (toolOrToolSetNames === undefined || toolOrToolSetNames.has(toolSet.referenceName)));
+			const enabled = toolOrToolSetNames === undefined || toolOrToolSetNames.has(toolSet.referenceName);
+			result.set(toolSet, enabled);
+
+			// if a mcp toolset is enabled, all tools in it are enabled
+			if (enabled && toolSet.source.type === 'mcp') {
+				for (const tool of toolSet.getTools()) {
+					if (tool.canBeReferencedInPrompt) {
+						result.set(tool, enabled);
+					}
+				}
+			}
 		}
 		return result;
 	}

src/vs/workbench/contrib/chat/browser/promptSyntax/promptFileRewriter.ts

@@ -9,7 +9,7 @@ import { ICodeEditorService } from '../../../../../editor/browser/services/codeE
 import { EditOperation } from '../../../../../editor/common/core/editOperation.js';
 import { Range } from '../../../../../editor/common/core/range.js';
 import { ITextModel } from '../../../../../editor/common/model.js';
-import { IToolAndToolSetEnablementMap, ToolSet } from '../../common/languageModelToolsService.js';
+import { IToolAndToolSetEnablementMap, IToolData, ToolSet } from '../../common/languageModelToolsService.js';
 import { IPromptsService } from '../../common/promptSyntax/service/promptsService.js';
 
 export class PromptFileRewriter {
@@ -59,11 +59,19 @@ export class PromptFileRewriter {
 			model.pushStackElement();
 			return;
 		}
+		const toolsCoveredBySets = new Set<IToolData>();
+		for (const [item, picked] of newTools) {
+			if (picked && item instanceof ToolSet) {
+				for (const tool of item.getTools()) {
+					toolsCoveredBySets.add(tool);
+				}
+			}
+		}
 		for (const [item, picked] of newTools) {
 			if (picked) {
 				if (item instanceof ToolSet) {
 					newToolNames.push(item.referenceName);
-				} else {
+				} else if (!toolsCoveredBySets.has(item)) {
 					newToolNames.push(item.toolReferenceName ?? item.displayName);
 				}
 			}