Enhance session management for selected tools in chat (microsoft#249878)

* have `ChatSelectedTools#disable|enable` and also consider tool sets for prompts

* * chat selected tools should have session enablement state
* use session enablement state for prompts

Author: jrieken

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

@@ -122,7 +122,7 @@ class ConfigureToolsAction extends Action2 {
 					}
 				}
 			}
-			widget.input.selectedToolsModel.update(disableToolSets, disableTools);
+			widget.input.selectedToolsModel.disable(disableToolSets, disableTools, false);
 		});
 
 		telemetryService.publicLog2<SelectedToolData, SelectedToolClassification>('chat/selectedTools', {
@@ -132,8 +132,6 @@ class ConfigureToolsAction extends Action2 {
 	}
 }
 
-
-
 export function registerChatToolActions() {
 	registerAction2(AcceptToolConfirmation);
 	registerAction2(ConfigureToolsAction);

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

@@ -290,7 +290,19 @@ export async function showToolsPicker(
 				}
 			}
 
-			onUpdate?.(result);
+			if (onUpdate) {
+				let didChange = toolsEntries.size !== result.size;
+				for (const [key, value] of toolsEntries) {
+					if (didChange) {
+						break;
+					}
+					didChange = result.get(key) !== value;
+				}
+
+				if (didChange) {
+					onUpdate(result);
+				}
+			}
 
 		} finally {
 			ignoreEvent = false;

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

@@ -659,6 +659,8 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge
 		const attachments = state.inputState?.chatContextAttachments ?? [];
 		this._attachmentModel.clearAndSetContext(...attachments);
 
+		this.selectedToolsModel.resetSessionEnablementState();
+
 		if (state.inputValue) {
 			this.setValue(state.inputValue, false);
 		}

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

@@ -4,12 +4,13 @@
  *--------------------------------------------------------------------------------------------*/
 
 import { Disposable } from '../../../../base/common/lifecycle.js';
-import { autorun, IObservable, observableFromEvent, ObservableMap, transaction } from '../../../../base/common/observable.js';
+import { autorun, derived, IObservable, observableFromEvent, ObservableMap, observableValue, transaction } from '../../../../base/common/observable.js';
 import { ObservableMemento, observableMemento } from '../../../../platform/observable/common/observableMemento.js';
 import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js';
 import { ChatMode } from '../common/constants.js';
 import { ILanguageModelToolsService, IToolData, ToolSet, ToolDataSource } from '../common/languageModelToolsService.js';
 
+
 /**
  * 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 tool sets and
@@ -18,8 +19,8 @@ import { ILanguageModelToolsService, IToolData, ToolSet, ToolDataSource } from '
  * also enabled.
  */
 type StoredData = {
-	disabledToolSets?: readonly string[];
-	disabledTools?: readonly string[];
+	readonly disabledToolSets?: readonly string[];
+	readonly disabledTools?: readonly string[];
 };
 
 const storedTools = observableMemento<StoredData>({
@@ -31,6 +32,8 @@ export class ChatSelectedTools extends Disposable {
 
 	private readonly _selectedTools: ObservableMemento<StoredData>;
 
+	private readonly _sessionSelectedTools = observableValue<StoredData>(this, {});
+
 	private readonly _allTools: IObservable<Readonly<IToolData>[]>;
 
 	/**
@@ -53,22 +56,39 @@ export class ChatSelectedTools extends Disposable {
 
 	constructor(
 		mode: IObservable<ChatMode>,
-		@ILanguageModelToolsService toolsService: ILanguageModelToolsService,
+		@ILanguageModelToolsService private readonly _toolsService: ILanguageModelToolsService,
 		@IStorageService storageService: IStorageService,
 	) {
 		super();
 
 		this._selectedTools = this._store.add(storedTools(StorageScope.WORKSPACE, StorageTarget.MACHINE, storageService));
 
-		this._allTools = observableFromEvent(toolsService.onDidChangeTools, () => Array.from(toolsService.getTools()));
+		this._allTools = observableFromEvent(_toolsService.onDidChangeTools, () => Array.from(_toolsService.getTools()));
+
+		const disabledDataObs = derived(r => {
+			const globalData = this._selectedTools.read(r);
+			const sessionData = this._sessionSelectedTools.read(r);
+
+			const toolSetIds = new Set<string>();
+			const toolIds = new Set<string>();
 
-		const disabledDataObs = this._selectedTools.map(data => {
-			return (data.disabledToolSets?.length || data.disabledTools?.length)
-				? {
-					toolSetIds: new Set(data.disabledToolSets),
-					toolIds: new Set(data.disabledTools),
+			for (const data of [globalData, sessionData]) {
+				if (data.disabledToolSets) {
+					for (const id of data.disabledToolSets) {
+						toolSetIds.add(id);
+					}
 				}
-				: undefined;
+				if (data.disabledTools) {
+					for (const id of data.disabledTools) {
+						toolIds.add(id);
+					}
+				}
+			}
+
+			if (toolSetIds.size === 0 && toolIds.size === 0) {
+				return undefined;
+			}
+			return { toolSetIds, toolIds };
 		});
 
 		this._store.add(autorun(r => {
@@ -82,12 +102,12 @@ export class ChatSelectedTools extends Disposable {
 				sourceByTool.set(tool, tool.source);
 			}
 
-			const toolSets = toolsService.toolSets.read(r);
+			const toolSets = _toolsService.toolSets.read(r);
 
 			for (const toolSet of toolSets) {
 
 				if (!toolSet.isHomogenous.read(r)) {
-					// only homogenous tool sets can shallow tools
+					// only homogenous tool sets can swallow tools
 					continue;
 				}
 
@@ -126,16 +146,27 @@ export class ChatSelectedTools extends Disposable {
 		}));
 	}
 
-	selectOnly(toolIds: readonly string[]): void {
-		const uniqueTools = new Set(toolIds);
+	resetSessionEnablementState() {
+		this._sessionSelectedTools.set({}, undefined);
+	}
 
-		const disabledTools = this._allTools.get().filter(tool => !uniqueTools.has(tool.id));
+	enable(toolSets: readonly ToolSet[], tools: readonly IToolData[], sessionOnly: boolean): void {
+		const toolIds = new Set(tools.map(t => t.id));
+		const toolsetIds = new Set(toolSets.map(t => t.id));
 
-		this.update([], disabledTools);
+		const disabledTools = this._allTools.get().filter(tool => !toolIds.has(tool.id));
+		const disabledToolSets = Array.from(this._toolsService.toolSets.get()).filter(toolset => !toolsetIds.has(toolset.id));
+
+		this.disable(disabledToolSets, disabledTools, sessionOnly);
 	}
 
-	update(disabledToolSets: readonly ToolSet[], disableTools: readonly IToolData[]): void {
-		this._selectedTools.set({
+	disable(disabledToolSets: readonly ToolSet[], disableTools: readonly IToolData[], sessionOnly: boolean): void {
+
+		const target = sessionOnly
+			? this._sessionSelectedTools
+			: this._selectedTools;
+
+		target.set({
 			disabledToolSets: disabledToolSets.map(t => t.id),
 			disabledTools: disableTools.map(t => t.id)
 		}, undefined);

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

@@ -51,7 +51,7 @@ import { ChatViewModel, IChatResponseViewModel, isRequestVM, isResponseVM } from
 import { IChatInputState } from '../common/chatWidgetHistoryService.js';
 import { CodeBlockModelCollection } from '../common/codeBlockModelCollection.js';
 import { ChatAgentLocation, ChatMode } from '../common/constants.js';
-import { ILanguageModelToolsService, ToolSet } from '../common/languageModelToolsService.js';
+import { ILanguageModelToolsService, IToolData, ToolSet } from '../common/languageModelToolsService.js';
 import { IPromptMetadata } from '../common/promptSyntax/parsers/types.js';
 import { IMetadata, IPromptsService } from '../common/promptSyntax/service/types.js';
 import { handleModeSwitch } from './actions/chatActions.js';
@@ -1568,27 +1568,22 @@ export class ChatWidget extends Disposable implements IChatWidget {
 		// if we have some tools present, the mode must have been equal to `agent`
 		assert(this.inputPart.currentMode === ChatMode.Agent, `Chat mode must be 'agent' when there are 'tools' defined, got ${this.inputPart.currentMode}.`);
 
-
-		// convert tools names to tool IDs
-		const toolIds = tools
-			.map((toolName) => {
-				const tool = this.toolsService.getToolByName(toolName);
-
-				if (tool === undefined) {
-					this.logService.warn(
-						`[setup tools]: cannot to find tool '${toolName}'`,
-					);
-				}
-
-				return tool;
-			})
-			.filter(isDefined)
-			.map(tool => tool.id);
-
-		// if there are some tools defined in the prompt files, select only the specified tools
-		this.inputPart
-			.selectedToolsModel
-			.selectOnly(toolIds);
+		// convert names to tools or tool sets
+		const enabledTools: IToolData[] = [];
+		const enabledToolSets: ToolSet[] = [];
+		for (const name of tools) {
+			const tool = this.toolsService.getToolByName(name);
+			if (tool) {
+				enabledTools.push(tool);
+				continue;
+			}
+			const toolset = this.toolsService.getToolSetByName(name);
+			if (toolset) {
+				enabledToolSets.push(toolset);
+				continue;
+			}
+		}
+		this.inputPart.selectedToolsModel.enable(enabledToolSets, enabledTools, true);
 	}
 
 	/**