Add a confirmation dialog for tools not running in a chat session

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

@@ -34,7 +34,7 @@ import { IChatVariablesService } from '../common/chatVariables.js';
 import { ChatWidgetHistoryService, IChatWidgetHistoryService } from '../common/chatWidgetHistoryService.js';
 import { ILanguageModelsService, LanguageModelsService } from '../common/languageModels.js';
 import { ILanguageModelStatsService, LanguageModelStatsService } from '../common/languageModelStats.js';
-import { ILanguageModelToolsService, LanguageModelToolsService } from '../common/languageModelToolsService.js';
+import { ILanguageModelToolsService } from '../common/languageModelToolsService.js';
 import { LanguageModelToolsExtensionPointHandler } from '../common/tools/languageModelToolsContribution.js';
 import { IVoiceChatService, VoiceChatService } from '../common/voiceChatService.js';
 import { PanelChatAccessibilityHelp, QuickChatAccessibilityHelp } from './actions/chatAccessibilityHelp.js';
@@ -71,6 +71,7 @@ import './contrib/chatInputEditorContrib.js';
 import './contrib/chatInputEditorHover.js';
 import { EditorContributionInstantiation, registerEditorContribution } from '../../../../editor/browser/editorExtensions.js';
 import { ChatEditorController } from './chatEditorController.js';
+import { LanguageModelToolsService } from './languageModelToolsService.js';
 
 // Register configuration
 const configurationRegistry = Registry.as<IConfigurationRegistry>(ConfigurationExtensions.Configuration);

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

@@ -0,0 +1,194 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { renderStringAsPlaintext } from '../../../../base/browser/markdownRenderer.js';
+import { RunOnceScheduler } from '../../../../base/common/async.js';
+import { CancellationToken } from '../../../../base/common/cancellation.js';
+import { CancellationError } from '../../../../base/common/errors.js';
+import { Emitter } from '../../../../base/common/event.js';
+import { Iterable } from '../../../../base/common/iterator.js';
+import { Disposable, IDisposable, toDisposable } from '../../../../base/common/lifecycle.js';
+import { localize } from '../../../../nls.js';
+import { IContextKeyService } from '../../../../platform/contextkey/common/contextkey.js';
+import { IDialogService } from '../../../../platform/dialogs/common/dialogs.js';
+import { IExtensionService } from '../../../services/extensions/common/extensions.js';
+import { ChatModel } from '../common/chatModel.js';
+import { ChatToolInvocation } from '../common/chatProgressTypes/chatToolInvocation.js';
+import { IChatService } from '../common/chatService.js';
+import { CountTokensCallback, ILanguageModelToolsService, IToolData, IToolImpl, IToolInvocation, IToolResult } from '../common/languageModelToolsService.js';
+
+interface IToolEntry {
+	data: IToolData;
+	impl?: IToolImpl;
+}
+
+export class LanguageModelToolsService extends Disposable implements ILanguageModelToolsService {
+	_serviceBrand: undefined;
+
+	private _onDidChangeTools = new Emitter<void>();
+	readonly onDidChangeTools = this._onDidChangeTools.event;
+
+	/** Throttle tools updates because it sends all tools and runs on context key updates */
+	private _onDidChangeToolsScheduler = new RunOnceScheduler(() => this._onDidChangeTools.fire(), 750);
+
+	private _tools = new Map<string, IToolEntry>();
+	private _toolContextKeys = new Set<string>();
+
+	constructor(
+		@IExtensionService private readonly _extensionService: IExtensionService,
+		@IContextKeyService private readonly _contextKeyService: IContextKeyService,
+		@IChatService private readonly _chatService: IChatService,
+		@IDialogService private readonly _dialogService: IDialogService,
+	) {
+		super();
+
+		this._register(this._contextKeyService.onDidChangeContext(e => {
+			if (e.affectsSome(this._toolContextKeys)) {
+				// Not worth it to compute a delta here unless we have many tools changing often
+				this._onDidChangeToolsScheduler.schedule();
+			}
+		}));
+	}
+
+	registerToolData(toolData: IToolData): IDisposable {
+		if (this._tools.has(toolData.id)) {
+			throw new Error(`Tool "${toolData.id}" is already registered.`);
+		}
+
+		// Ensure that text/plain is supported
+		if (!toolData.supportedContentTypes.includes('text/plain')) {
+			toolData.supportedContentTypes.push('text/plain');
+		}
+
+		this._tools.set(toolData.id, { data: toolData });
+		this._onDidChangeToolsScheduler.schedule();
+
+		toolData.when?.keys().forEach(key => this._toolContextKeys.add(key));
+
+		return toDisposable(() => {
+			this._tools.delete(toolData.id);
+			this._refreshAllToolContextKeys();
+			this._onDidChangeToolsScheduler.schedule();
+		});
+	}
+
+	private _refreshAllToolContextKeys() {
+		this._toolContextKeys.clear();
+		for (const tool of this._tools.values()) {
+			tool.data.when?.keys().forEach(key => this._toolContextKeys.add(key));
+		}
+	}
+
+	registerToolImplementation(id: string, tool: IToolImpl): IDisposable {
+		const entry = this._tools.get(id);
+		if (!entry) {
+			throw new Error(`Tool "${id}" was not contributed.`);
+		}
+
+		if (entry.impl) {
+			throw new Error(`Tool "${id}" already has an implementation.`);
+		}
+
+		entry.impl = tool;
+		return toDisposable(() => {
+			entry.impl = undefined;
+		});
+	}
+
+	getTools(): Iterable<Readonly<IToolData>> {
+		const toolDatas = Iterable.map(this._tools.values(), i => i.data);
+		return Iterable.filter(toolDatas, toolData => !toolData.when || this._contextKeyService.contextMatchesRules(toolData.when));
+	}
+
+	getTool(id: string): IToolData | undefined {
+		return this._getToolEntry(id)?.data;
+	}
+
+	private _getToolEntry(id: string): IToolEntry | undefined {
+		const entry = this._tools.get(id);
+		if (entry && (!entry.data.when || this._contextKeyService.contextMatchesRules(entry.data.when))) {
+			return entry;
+		} else {
+			return undefined;
+		}
+	}
+
+	getToolByName(name: string): IToolData | undefined {
+		for (const toolData of this.getTools()) {
+			if (toolData.name === name) {
+				return toolData;
+			}
+		}
+		return undefined;
+	}
+
+	async invokeTool(dto: IToolInvocation, countTokens: CountTokensCallback, token: CancellationToken): Promise<IToolResult> {
+		// When invoking a tool, don't validate the "when" clause. An extension may have invoked a tool just as it was becoming disabled, and just let it go through rather than throw and break the chat.
+		let tool = this._tools.get(dto.toolId);
+		if (!tool) {
+			throw new Error(`Tool ${dto.toolId} was not contributed`);
+		}
+
+		if (!tool.impl) {
+			await this._extensionService.activateByEvent(`onLanguageModelTool:${dto.toolId}`);
+
+			// Extension should activate and register the tool implementation
+			tool = this._tools.get(dto.toolId);
+			if (!tool?.impl) {
+				throw new Error(`Tool ${dto.toolId} does not have an implementation registered.`);
+			}
+		}
+
+		// Shortcut to write to the model directly here, but could call all the way back to use the real stream.
+		let toolInvocation: ChatToolInvocation | undefined;
+		if (dto.context) {
+			const model = this._chatService.getSession(dto.context?.sessionId) as ChatModel;
+			const request = model.getRequests().at(-1)!;
+
+			const participantName = request.response?.agent?.fullName ?? ''; // This should always be set in this scenario with a new live request
+
+			const prepared = tool.impl.prepareToolInvocation ?
+				await tool.impl.prepareToolInvocation(participantName, dto.parameters, token)
+				: undefined;
+			const confirmationMessages = tool.data.requiresConfirmation ?
+				prepared?.confirmationMessages ?? {
+					title: localize('toolConfirmTitle', "Use {0}?", `"${tool.data.displayName ?? tool.data.id}"`),
+					message: localize('toolConfirmMessage', "{0} will use {1}.", participantName, `"${tool.data.displayName ?? tool.data.id}"`),
+				}
+				: undefined;
+
+			const defaultMessage = localize('toolInvocationMessage', "Using {0}", `"${tool.data.displayName ?? tool.data.id}"`);
+			const invocationMessage = prepared?.invocationMessage ?? defaultMessage;
+			toolInvocation = new ChatToolInvocation(invocationMessage, confirmationMessages);
+			token.onCancellationRequested(() => {
+				toolInvocation!.confirmed.complete(false);
+			});
+			model.acceptResponseProgress(request, toolInvocation);
+			if (tool.data.requiresConfirmation) {
+				const userConfirmed = await toolInvocation.confirmed.p;
+				if (!userConfirmed) {
+					throw new CancellationError();
+				}
+			}
+		} else if (tool.data.requiresConfirmation) {
+			const prepared = tool.impl.prepareToolInvocation ?
+				await tool.impl.prepareToolInvocation('Some Extension', dto.parameters, token)
+				: undefined;
+
+			const confirmationMessages = prepared?.confirmationMessages ?? {
+				title: localize('toolConfirmTitle', "Use {0}?", `"${tool.data.displayName ?? tool.data.id}"`),
+				message: localize('toolConfirmMessage', "{0} will use {1}.", 'Some Extension', `"${tool.data.displayName ?? tool.data.id}"`),
+			};
+
+			await this._dialogService.confirm({ message: confirmationMessages.title, detail: renderStringAsPlaintext(confirmationMessages.message) });
+		}
+
+		try {
+			return await tool.impl.invoke(dto, countTokens, token);
+		} finally {
+			toolInvocation?.isCompleteDeferred.complete();
+		}
+	}
+}