chat: improve tool selection rememberance

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

@@ -3,6 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
+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';
@@ -23,7 +24,7 @@ import { ChatContextKeys } from '../../common/chatContextKeys.js';
 import { IChatToolInvocation } from '../../common/chatService.js';
 import { isResponseVM } from '../../common/chatViewModel.js';
 import { ChatMode } from '../../common/constants.js';
-import { ILanguageModelToolsService, IToolData } from '../../common/languageModelToolsService.js';
+import { ILanguageModelToolsService, IToolData, ToolDataSource } from '../../common/languageModelToolsService.js';
 import { IChatWidget, IChatWidgetService } from '../chat.js';
 import { CHAT_CATEGORY } from './chatActions.js';
 
@@ -140,14 +141,15 @@ export class AttachToolsAction extends Action2 {
 		}
 
 		const enum BucketOrdinal { Extension, Mcp, Other }
-		type BucketPick = IQuickPickItem & { picked: boolean; ordinal: BucketOrdinal; status?: string; children: ToolPick[] };
+		type BucketPick = IQuickPickItem & { picked: boolean; ordinal: BucketOrdinal; status?: string; children: ToolPick[]; source: ToolDataSource };
 		type ToolPick = IQuickPickItem & { picked: boolean; tool: IToolData; parent: BucketPick };
 		type MyPick = ToolPick | BucketPick;
 
 		const defaultBucket: BucketPick = {
 			type: 'item',
 			children: [],
 			label: localize('defaultBucketLabel', "Other Tools"),
+			source: { type: 'internal' },
 			ordinal: BucketOrdinal.Other,
 			picked: true,
 		};
@@ -156,36 +158,47 @@ export class AttachToolsAction extends Action2 {
 		const toolBuckets = new Map<string, BucketPick>();
 
 		for (const tool of toolsService.getTools()) {
-
 			if (!tool.canBeReferencedInPrompt) {
 				continue;
 			}
 
 			let bucket: BucketPick;
 
-			const mcpServer = mcpServerByTool.get(tool.id);
-			const ext = extensionService.extensions.find(value => ExtensionIdentifier.equals(value.identifier, tool.extensionId));
-			if (mcpServer) {
-				bucket = toolBuckets.get(mcpServer.definition.id) ?? {
+			if (tool.source.type === 'mcp') {
+				const mcpServer = mcpServerByTool.get(tool.id);
+				if (!mcpServer) {
+					continue;
+				}
+				bucket = toolBuckets.get(tool.source.collectionId) ?? {
 					type: 'item',
-					label: localize('mcplabel', "MCP Server: {0}", mcpServer.definition.label),
+					label: localize('mcplabel', "MCP Server: {0}", mcpServer?.definition.label),
 					status: localize('mcpstatus', "From {0} ({1})", mcpServer.collection.label, McpConnectionState.toString(mcpServer.connectionState.get())),
 					ordinal: BucketOrdinal.Mcp,
+					source: tool.source,
 					picked: false,
 					children: []
 				};
-				toolBuckets.set(mcpServer.definition.id, bucket);
-			} else if (ext) {
-				bucket = toolBuckets.get(ExtensionIdentifier.toKey(ext.identifier)) ?? {
+				toolBuckets.set(tool.source.collectionId, 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;
+				}
+
+				bucket = toolBuckets.get(ExtensionIdentifier.toKey(extensionId)) ?? {
 					type: 'item',
 					label: ext.displayName ?? ext.name,
 					ordinal: BucketOrdinal.Extension,
 					picked: false,
+					source: tool.source,
 					children: []
 				};
 				toolBuckets.set(ExtensionIdentifier.toKey(ext.identifier), bucket);
-			} else {
+			} else if (tool.source.type === 'internal') {
 				bucket = defaultBucket;
+			} else {
+				assertNever(tool.source);
 			}
 
 			const picked = nowSelectedTools.has(tool);
@@ -243,13 +256,19 @@ export class AttachToolsAction extends Action2 {
 				lastSelectedItems = new Set(items);
 				picker.selectedItems = items;
 
-				const toolPicks = items.filter(isToolPick);
-				const allTools = picks.filter(isToolPick);
-				if (toolPicks.length === allTools.length) {
-					widget.input.selectedToolsModel.reset();
-				} else {
-					widget.input.selectedToolsModel.update(items.filter(isToolPick).map(tool => tool.tool));
+				const disableBuckets: ToolDataSource[] = [];
+				const disableTools: IToolData[] = [];
+				for (const item of picks) {
+					if (item.type === 'item' && !item.picked) {
+						if (isBucketPick(item)) {
+							disableBuckets.push(item.source);
+						} else if (isToolPick(item) && item.parent.picked) {
+							disableTools.push(item.tool);
+						}
+					}
 				}
+
+				widget.input.selectedToolsModel.update(disableBuckets, disableTools);
 			} finally {
 				ignoreEvent = false;
 			}

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

@@ -240,13 +240,18 @@ configurationRegistry.registerConfiguration({
 			tags: ['experimental'],
 		},
 		[mcpDiscoverySection]: {
-			type: 'object',
-			default: Object.fromEntries(allDiscoverySources.map(k => [k, true])),
-			properties: Object.fromEntries(allDiscoverySources.map(k => [
-				k,
-				{ type: 'boolean', description: nls.localize('mcp.discovery.source', "Enables discovery of {0} servers", discoverySourceLabel[k]) }
-			])),
-			markdownDescription: nls.localize('mpc.discovery.enabled', "Configures discovery of Model Context Protocol servers on the machine. It may be set to `true` or `false` to disable or enable all sources, and an array of sources you wish to enable."),
+			oneOf: [
+				{ type: 'boolean' },
+				{
+					type: 'object',
+					default: Object.fromEntries(allDiscoverySources.map(k => [k, true])),
+					properties: Object.fromEntries(allDiscoverySources.map(k => [
+						k,
+						{ type: 'boolean', description: nls.localize('mcp.discovery.source', "Enables discovery of {0} servers", discoverySourceLabel[k]) }
+					])),
+				}
+			],
+			markdownDescription: nls.localize('mpc.discovery.enabled', "Configures discovery of Model Context Protocol servers on the machine. It may be set to `true` or `false` to disable or enable all sources, and an mapping sources you wish to enable."),
 		},
 		[PromptsConfig.KEY]: {
 			type: 'boolean',

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

@@ -16,12 +16,19 @@ import { MenuItemAction } from '../../../../platform/actions/common/actions.js';
 import { IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js';
 import { ObservableMemento, observableMemento } from '../../../../platform/observable/common/observableMemento.js';
 import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js';
-import { ILanguageModelToolsService, IToolData } from '../common/languageModelToolsService.js';
+import { ILanguageModelToolsService, IToolData, ToolDataSource } from '../common/languageModelToolsService.js';
 
-type StoredData = { all: boolean; ids?: string[] };
+/**
+ * New tools and new tool sources that come in should generally be enabled until
+ * the user disables them. To store things, we store only the buckets and
+ * individual tools that were disabled, so the new data sources that come in
+ * are enabled, and new tools that come in for data sources not disabled are
+ * also enabled.
+ */
+type StoredData = { disabledBuckets?: /* ToolDataSource.toKey */ readonly string[]; disabledTools?: readonly string[] };
 
 const storedTools = observableMemento<StoredData>({
-	defaultValue: { all: true },
+	defaultValue: {},
 	key: 'chat/selectedTools',
 });
 
@@ -47,15 +54,23 @@ export class ChatSelectedTools extends Disposable {
 			() => Array.from(toolsService.getTools()).filter(t => t.canBeReferencedInPrompt)
 		);
 
+		const disabledData = this._selectedTools.map(data => {
+			return (data.disabledBuckets?.length || data.disabledTools?.length) && {
+				buckets: new Set(data.disabledBuckets),
+				toolIds: new Set(data.disabledTools),
+			};
+		});
+
 		this.tools = derived(r => {
-			const stored = this._selectedTools.read(r);
+			const disabled = disabledData.read(r);
 			const tools = allTools.read(r);
-			if (stored.all) {
+			if (!disabled) {
 				return tools;
 			}
 
-			const ids = new Set(stored.ids);
-			return tools.filter(t => ids.has(t.id));
+			return tools.filter(t =>
+				!(disabled.toolIds.has(t.id) || disabled.buckets.has(ToolDataSource.toKey(t.source)))
+			);
 		});
 
 		const toolsCount = derived(r => {
@@ -101,11 +116,10 @@ export class ChatSelectedTools extends Disposable {
 		};
 	}
 
-	update(tools: readonly IToolData[]): void {
-		this._selectedTools.set({ all: false, ids: tools.map(t => t.id) }, undefined);
-	}
-
-	reset(): void {
-		this._selectedTools.set({ all: true }, undefined);
+	update(disableBuckets: readonly ToolDataSource[], disableTools: readonly IToolData[]): void {
+		this._selectedTools.set({
+			disabledBuckets: disableBuckets.map(ToolDataSource.toKey),
+			disabledTools: disableTools.map(t => t.id)
+		}, undefined);
 	}
 }

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

@@ -281,7 +281,8 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 					result: 'success',
 					chatSessionId: dto.context?.sessionId,
 					toolId: tool.data.id,
-					toolExtensionId: tool.data.extensionId?.value,
+					toolExtensionId: tool.data.source.type === 'extension' ? tool.data.source.extensionId.value : undefined,
+					toolSourceKind: tool.data.source.type,
 				});
 			return toolResult;
 		} catch (err) {
@@ -292,7 +293,8 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo
 					result,
 					chatSessionId: dto.context?.sessionId,
 					toolId: tool.data.id,
-					toolExtensionId: tool.data.extensionId?.value,
+					toolExtensionId: tool.data.source.type === 'extension' ? tool.data.source.extensionId.value : undefined,
+					toolSourceKind: tool.data.source.type,
 				});
 			throw err;
 		} finally {
@@ -359,13 +361,15 @@ type LanguageModelToolInvokedEvent = {
 	chatSessionId: string | undefined;
 	toolId: string;
 	toolExtensionId: string | undefined;
+	toolSourceKind: string;
 };
 
 type LanguageModelToolInvokedClassification = {
 	result: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'Whether invoking the LanguageModelTool resulted in an error.' };
 	chatSessionId: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The ID of the chat session that the tool was used within, if applicable.' };
 	toolId: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The ID of the tool used.' };
 	toolExtensionId: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The extension that contributed the tool.' };
+	toolSourceKind: { classification: 'SystemMetaData'; purpose: 'FeatureInsight'; comment: 'The source (mcp/extension/internal) of the tool.' };
 	owner: 'roblourens';
 	comment: 'Provides insight into the usage of language model tools.';
 };

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

@@ -19,7 +19,7 @@ import { Schemas } from '../../../../base/common/network.js';
 
 export interface IToolData {
 	id: string;
-	extensionId?: ExtensionIdentifier;
+	source: ToolDataSource;
 	toolReferenceName?: string;
 	icon?: { dark: URI; light?: URI } | ThemeIcon;
 	when?: ContextKeyExpression;
@@ -36,6 +36,21 @@ export interface IToolData {
 	runsInWorkspace?: boolean;
 }
 
+export type ToolDataSource =
+	| { type: 'extension'; extensionId: ExtensionIdentifier }
+	| { type: 'mcp'; collectionId: string }
+	| { type: 'internal' };
+
+export namespace ToolDataSource {
+	export function toKey(source: ToolDataSource): string {
+		switch (source.type) {
+			case 'extension': return `extension:${source.extensionId.value}`;
+			case 'mcp': return `mcp:${source.collectionId}`;
+			case 'internal': return 'internal';
+		}
+	}
+}
+
 export interface IToolInvocation {
 	callId: string;
 	toolId: string;

src/vs/workbench/contrib/chat/common/tools/editFileTool.ts

@@ -50,6 +50,7 @@ export const EditToolData: IToolData = {
 	id: InternalEditToolId,
 	displayName: localize('chat.tools.editFile', "Edit File"),
 	modelDescription: `Edit a file in the workspace. Use this tool once per file that needs to be modified, even if there are multiple changes for a file. Generate the "explanation" property first. ${codeInstructions}`,
+	source: { type: 'internal' },
 	inputSchema: {
 		type: 'object',
 		properties: {

src/vs/workbench/contrib/chat/common/tools/insertNotebookCellsTool.ts

@@ -33,6 +33,7 @@ export const EditToolData: IToolData = {
 	id: InternalEditToolId,
 	displayName: localize('chat.tools.editFile', "Edit File"),
 	modelDescription: `Insert cells into a new notebook n the workspace. Use this tool once per file that needs to be modified, even if there are multiple changes for a file. Generate the "explanation" property first. ${codeInstructions}`,
+	source: { type: 'internal' },
 	inputSchema: {
 		type: 'object',
 		properties: {

src/vs/workbench/contrib/chat/common/tools/languageModelToolsContribution.ts

@@ -197,7 +197,7 @@ export class LanguageModelToolsExtensionPointHandler implements IWorkbenchContri
 
 					const tool: IToolData = {
 						...rawTool,
-						extensionId: extension.description.identifier,
+						source: { type: 'extension', extensionId: extension.description.identifier },
 						inputSchema: rawTool.inputSchema,
 						id: rawTool.name,
 						icon,

src/vs/workbench/contrib/chat/electron-sandbox/tools/fetchPageTool.ts

@@ -17,6 +17,7 @@ export const FetchWebPageToolData: IToolData = {
 	displayName: 'Fetch Web Page',
 	canBeReferencedInPrompt: false,
 	modelDescription: localize('fetchWebPage.modelDescription', 'Fetches the main content from a web page. This tool is useful for summarizing or analyzing the content of a webpage.'),
+	source: { type: 'internal' },
 	inputSchema: {
 		type: 'object',
 		properties: {

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

@@ -39,7 +39,8 @@ suite('LanguageModelToolsService', () => {
 		const toolData: IToolData = {
 			id: 'testTool',
 			modelDescription: 'Test Tool',
-			displayName: 'Test Tool'
+			displayName: 'Test Tool',
+			source: { type: 'internal' },
 		};
 
 		const disposable = service.registerToolData(toolData);
@@ -52,7 +53,8 @@ suite('LanguageModelToolsService', () => {
 		const toolData: IToolData = {
 			id: 'testTool',
 			modelDescription: 'Test Tool',
-			displayName: 'Test Tool'
+			displayName: 'Test Tool',
+			source: { type: 'internal' },
 		};
 
 		store.add(service.registerToolData(toolData));
@@ -71,20 +73,23 @@ suite('LanguageModelToolsService', () => {
 			id: 'testTool1',
 			modelDescription: 'Test Tool 1',
 			when: ContextKeyEqualsExpr.create('testKey', false),
-			displayName: 'Test Tool'
+			displayName: 'Test Tool',
+			source: { type: 'internal' },
 		};
 
 		const toolData2: IToolData = {
 			id: 'testTool2',
 			modelDescription: 'Test Tool 2',
 			when: ContextKeyEqualsExpr.create('testKey', true),
-			displayName: 'Test Tool'
+			displayName: 'Test Tool',
+			source: { type: 'internal' },
 		};
 
 		const toolData3: IToolData = {
 			id: 'testTool3',
 			modelDescription: 'Test Tool 3',
-			displayName: 'Test Tool'
+			displayName: 'Test Tool',
+			source: { type: 'internal' },
 		};
 
 		store.add(service.registerToolData(toolData1));
@@ -101,7 +106,8 @@ suite('LanguageModelToolsService', () => {
 		const toolData: IToolData = {
 			id: 'testTool',
 			modelDescription: 'Test Tool',
-			displayName: 'Test Tool'
+			displayName: 'Test Tool',
+			source: { type: 'internal' },
 		};
 
 		store.add(service.registerToolData(toolData));
@@ -135,7 +141,8 @@ suite('LanguageModelToolsService', () => {
 		const toolData: IToolData = {
 			id: 'testTool',
 			modelDescription: 'Test Tool',
-			displayName: 'Test Tool'
+			displayName: 'Test Tool',
+			source: { type: 'internal' },
 		};
 
 		store.add(service.registerToolData(toolData));