From 8d3270de420bd4320fdfbf2c18b681b6abd680a2 Mon Sep 17 00:00:00 2001 From: Connor Peet Date: Tue, 13 Jan 2026 16:02:29 -0800 Subject: [PATCH] chat: wip on model-specific tools This is a rethinking of the initial API proposed in `languageModelToolSupportsModel.d.ts`. 1. This switches to `models?: LanguageModelChatSelector[];` to control enablement. definitely open to switching this out, but I think a synchronously-analyzable expression is important to retain the data flows in core without too many races. 2. The extension is able to define a tool at runtime via registerToolDefinition. This should let us have entirely service-driven tools from model providers without requiring a static definition for each one. We can also have model-specific variants of tools without a ton of package.json work for each variant of the tool (as initially proposed using `when` clauses) This then propagates that down into the tools service. Currently I have this as just compiling to a `when` expression once it reaches the main thread. Then, for the tools service, it takes an IContextKeyService in cases where tools should be enumerated, and the chat input sets the model keys in its scoped context key service. This allows the tools to be filtered correctly in the tool picker. I initially thought about allowing multiple definitions be registered for the same tool name/id for model-specific variants of tools but I realized that gets really gnarly and we already have a `toolReferenceName` that multiple tools can register into. Todo for tomorrow morning: - Tools don't make it to the ChatRequest yet, or something, still need to investigate - Need to make sure tools in prompts/models all work. For a first pass I think we can let prompts/modes reference all tools by toolReferenceName. - Validate that multiple tools actually can safely share a reference name (and do some priority ordering?) - General further validation - Some unit tests --- .../browser/mainThreadLanguageModelTools.ts | 118 ++++++++++++++++-- .../workbench/api/common/extHost.api.impl.ts | 3 + .../workbench/api/common/extHost.protocol.ts | 16 ++- .../api/common/extHostChatAgents2.ts | 17 +-- .../api/common/extHostLanguageModelTools.ts | 56 +++++---- .../chat/browser/actions/chatActions.ts | 2 +- .../chat/browser/actions/chatToolActions.ts | 2 +- .../chat/browser/actions/chatToolPicker.ts | 41 +----- .../attachments/chatAttachmentWidgets.ts | 5 +- .../contrib/chat/browser/chat.contribution.ts | 2 +- .../browser/chatSetup/chatSetupProviders.ts | 4 +- .../promptToolsCodeLensProvider.ts | 8 +- .../tools/languageModelToolsService.ts | 104 +++++++++------ .../browser/tools/toolSetsContribution.ts | 2 +- .../contrib/chat/browser/widget/chatWidget.ts | 2 +- .../browser/widget/input/chatInputPart.ts | 14 +++ .../browser/widget/input/chatSelectedTools.ts | 16 ++- .../chat/common/actions/chatContextKeys.ts | 7 ++ .../languageProviders/promptValidator.ts | 2 +- .../tools/builtinTools/runSubagentTool.ts | 4 +- .../tools/languageModelToolsContribution.ts | 2 +- .../common/tools/languageModelToolsService.ts | 73 +++++++++-- .../tools/languageModelToolsService.test.ts | 74 +++++------ .../widget/input/chatSelectedTools.test.ts | 7 +- .../tools/mockLanguageModelToolsService.ts | 17 +-- .../browser/alternativeRecommendation.ts | 5 +- .../browser/tools/runInTerminalTool.ts | 4 +- ...ode.proposed.chatParticipantAdditions.d.ts | 12 -- ...oposed.languageModelToolSupportsModel.d.ts | 61 +++++++-- 29 files changed, 455 insertions(+), 225 deletions(-) diff --git a/src/vs/workbench/api/browser/mainThreadLanguageModelTools.ts b/src/vs/workbench/api/browser/mainThreadLanguageModelTools.ts index 11a10d82cb7..934fcab4942 100644 --- a/src/vs/workbench/api/browser/mainThreadLanguageModelTools.ts +++ b/src/vs/workbench/api/browser/mainThreadLanguageModelTools.ts @@ -6,10 +6,57 @@ import { CancellationToken } from '../../../base/common/cancellation.js'; import { Disposable, DisposableMap } from '../../../base/common/lifecycle.js'; import { revive } from '../../../base/common/marshalling.js'; -import { CountTokensCallback, ILanguageModelToolsService, IToolInvocation, IToolProgressStep, IToolResult, ToolProgress, toolResultHasBuffers } from '../../contrib/chat/common/tools/languageModelToolsService.js'; +import { ThemeIcon } from '../../../base/common/themables.js'; +import { isDefined } from '../../../base/common/types.js'; +import { isUriComponents, URI, UriComponents } from '../../../base/common/uri.js'; +import { ContextKeyExpr, ContextKeyExpression } from '../../../platform/contextkey/common/contextkey.js'; +import { ChatContextKeys } from '../../contrib/chat/common/actions/chatContextKeys.js'; +import { CountTokensCallback, ILanguageModelToolsService, IToolData, IToolInvocation, IToolProgressStep, IToolResult, ToolDataSource, ToolProgress, toolResultHasBuffers } from '../../contrib/chat/common/tools/languageModelToolsService.js'; import { IExtHostContext, extHostNamedCustomer } from '../../services/extensions/common/extHostCustomers.js'; import { Dto, SerializableObjectWithBuffers } from '../../services/extensions/common/proxyIdentifier.js'; -import { ExtHostContext, ExtHostLanguageModelToolsShape, IToolDataDto, MainContext, MainThreadLanguageModelToolsShape } from '../common/extHost.protocol.js'; +import { ExtHostContext, ExtHostLanguageModelToolsShape, ILanguageModelChatSelectorDto, IToolDataDto, IToolDefinitionDto, MainContext, MainThreadLanguageModelToolsShape } from '../common/extHost.protocol.js'; + +/** + * Compile a single model selector to a ContextKeyExpression. + * All specified fields must match (AND). + */ +function selectorToContextKeyExpr(selector: ILanguageModelChatSelectorDto): ContextKeyExpression | undefined { + const conditions: ContextKeyExpression[] = []; + if (selector.id) { + conditions.push(ContextKeyExpr.equals(ChatContextKeys.Model.id.key, selector.id)); + } + if (selector.vendor) { + conditions.push(ContextKeyExpr.equals(ChatContextKeys.Model.vendor.key, selector.vendor)); + } + if (selector.family) { + conditions.push(ContextKeyExpr.equals(ChatContextKeys.Model.family.key, selector.family)); + } + if (selector.version) { + conditions.push(ContextKeyExpr.equals(ChatContextKeys.Model.version.key, selector.version)); + } + if (conditions.length === 0) { + return undefined; + } + return ContextKeyExpr.and(...conditions); +} + +/** + * Compile multiple model selectors to a ContextKeyExpression. + * Any selector may match (OR). + */ +function selectorsToContextKeyExpr(selectors: ILanguageModelChatSelectorDto[]): ContextKeyExpression | undefined { + if (selectors.length === 0) { + return undefined; + } + const expressions = selectors.map(selectorToContextKeyExpr).filter(isDefined); + if (expressions.length === 0) { + return undefined; + } + if (expressions.length === 1) { + return expressions[0]; + } + return ContextKeyExpr.or(...expressions); +} @extHostNamedCustomer(MainContext.MainThreadLanguageModelTools) export class MainThreadLanguageModelTools extends Disposable implements MainThreadLanguageModelToolsShape { @@ -32,7 +79,7 @@ export class MainThreadLanguageModelTools extends Disposable implements MainThre } private getToolDtos(): IToolDataDto[] { - return Array.from(this._languageModelToolsService.getTools()) + return Array.from(this._languageModelToolsService.getAllToolsIncludingDisabled()) .map(tool => ({ id: tool.id, displayName: tool.displayName, @@ -93,16 +140,71 @@ export class MainThreadLanguageModelTools extends Disposable implements MainThre } }, prepareToolInvocation: (context, token) => this._proxy.$prepareToolInvocation(id, context, token), - supportsModel: (modelId, token) => this._proxy.$supportsModel(id, modelId, token), }); this._tools.set(id, disposable); } + $registerToolWithDefinition(definition: IToolDefinitionDto): void { + let icon: IToolData['icon'] | undefined; + if (definition.icon) { + if (ThemeIcon.isThemeIcon(definition.icon)) { + icon = definition.icon; + } else if (typeof definition.icon === 'object' && definition.icon !== null && isUriComponents(definition.icon)) { + icon = { dark: URI.revive(definition.icon as UriComponents) }; + } else { + const iconObj = definition.icon as { light?: UriComponents; dark: UriComponents }; + icon = { dark: URI.revive(iconObj.dark), light: iconObj.light ? URI.revive(iconObj.light) : undefined }; + } + } + + // Compile model selectors to when clause + let when: ContextKeyExpression | undefined; + if (definition.models?.length) { + when = selectorsToContextKeyExpr(definition.models); + } + + // Convert source from DTO + const source = revive(definition.source); + + // Create the tool data + const toolData: IToolData = { + id: definition.id, + displayName: definition.displayName, + toolReferenceName: definition.toolReferenceName, + legacyToolReferenceFullNames: definition.legacyToolReferenceFullNames, + tags: definition.tags, + userDescription: definition.userDescription, + modelDescription: definition.modelDescription, + inputSchema: definition.inputSchema, + source, + icon, + when, + models: definition.models, + canBeReferencedInPrompt: true, + }; + + // Register both tool data and implementation + const id = definition.id; + const disposable = this._languageModelToolsService.registerTool( + toolData, + { + invoke: async (dto, countTokens, progress, token) => { + try { + this._runningToolCalls.set(dto.callId, { countTokens, progress }); + const resultSerialized = await this._proxy.$invokeTool(dto, token); + const resultDto: Dto = resultSerialized instanceof SerializableObjectWithBuffers ? resultSerialized.value : resultSerialized; + return revive(resultDto); + } finally { + this._runningToolCalls.delete(dto.callId); + } + }, + prepareToolInvocation: (context, token) => this._proxy.$prepareToolInvocation(id, context, token), + } + ); + this._tools.set(id, disposable); + } + $unregisterTool(name: string): void { this._tools.deleteAndDispose(name); } - - $supportsModel(toolId: string, modelId: string, token: CancellationToken): Promise { - return this._languageModelToolsService.supportsModel(toolId, modelId, token); - } } diff --git a/src/vs/workbench/api/common/extHost.api.impl.ts b/src/vs/workbench/api/common/extHost.api.impl.ts index 830708d6cfe..0a7ba21b116 100644 --- a/src/vs/workbench/api/common/extHost.api.impl.ts +++ b/src/vs/workbench/api/common/extHost.api.impl.ts @@ -1607,6 +1607,9 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I registerTool(name: string, tool: vscode.LanguageModelTool) { return extHostLanguageModelTools.registerTool(extension, name, tool); }, + registerToolDefinition(definition: vscode.LanguageModelToolDefinition, tool: vscode.LanguageModelTool) { + return extHostLanguageModelTools.registerToolDefinition(extension, definition, tool); + }, invokeTool(nameOrInfo: string | vscode.LanguageModelToolInformation, parameters: vscode.LanguageModelToolInvocationOptions, token?: vscode.CancellationToken) { if (typeof nameOrInfo !== 'string') { checkProposedApiEnabled(extension, 'chatParticipantAdditions'); diff --git a/src/vs/workbench/api/common/extHost.protocol.ts b/src/vs/workbench/api/common/extHost.protocol.ts index c6382b263b6..3e45e3e4694 100644 --- a/src/vs/workbench/api/common/extHost.protocol.ts +++ b/src/vs/workbench/api/common/extHost.protocol.ts @@ -1489,14 +1489,26 @@ export interface IToolDataDto { inputSchema?: IJSONSchema; } +export interface ILanguageModelChatSelectorDto { + vendor?: string; + family?: string; + version?: string; + id?: string; +} + +export interface IToolDefinitionDto extends IToolDataDto { + icon?: IconPathDto; + models?: ILanguageModelChatSelectorDto[]; +} + export interface MainThreadLanguageModelToolsShape extends IDisposable { $getTools(): Promise[]>; $acceptToolProgress(callId: string, progress: IToolProgressStep): void; $invokeTool(dto: Dto, token?: CancellationToken): Promise | SerializableObjectWithBuffers>>; $countTokensForInvocation(callId: string, input: string, token: CancellationToken): Promise; $registerTool(id: string): void; + $registerToolWithDefinition(definition: IToolDefinitionDto): void; $unregisterTool(name: string): void; - $supportsModel(toolId: string, modelId: string, token: CancellationToken): Promise; } export type IChatRequestVariableValueDto = Dto; @@ -1505,9 +1517,7 @@ export interface ExtHostLanguageModelToolsShape { $onDidChangeTools(tools: IToolDataDto[]): void; $invokeTool(dto: Dto, token: CancellationToken): Promise | SerializableObjectWithBuffers>>; $countTokensForInvocation(callId: string, input: string, token: CancellationToken): Promise; - $prepareToolInvocation(toolId: string, context: IToolInvocationPreparationContext, token: CancellationToken): Promise; - $supportsModel(toolId: string, modelId: string, token: CancellationToken): Promise; } export interface MainThreadUrlsShape extends IDisposable { diff --git a/src/vs/workbench/api/common/extHostChatAgents2.ts b/src/vs/workbench/api/common/extHostChatAgents2.ts index 674633d2416..d9b6230da4e 100644 --- a/src/vs/workbench/api/common/extHostChatAgents2.ts +++ b/src/vs/workbench/api/common/extHostChatAgents2.ts @@ -738,22 +738,7 @@ export class ExtHostChatAgents2 extends Disposable implements ExtHostChatAgentsS return new Map(); } const result = new Map(); - const allTools = this._tools.getTools(extension); - - // Check model support for all tools in parallel - const toolChecks = await Promise.all( - Array.from(allTools).map(async (tool) => { - const supports = await this._tools.supportsModel(tool.name, modelId, token); - // undefined means no supportsModel impl, treat as supported - // false means explicitly not supported - return { tool, supported: supports === true }; - }) - ); - - for (const { tool, supported } of toolChecks) { - if (!supported) { - continue; - } + for (const tool of this._tools.getTools(extension)) { if (typeof tools[tool.name] === 'boolean') { result.set(tool, tools[tool.name]); } diff --git a/src/vs/workbench/api/common/extHostLanguageModelTools.ts b/src/vs/workbench/api/common/extHostLanguageModelTools.ts index fa599c1bcd5..65073d77ea9 100644 --- a/src/vs/workbench/api/common/extHostLanguageModelTools.ts +++ b/src/vs/workbench/api/common/extHostLanguageModelTools.ts @@ -18,7 +18,7 @@ import { InternalFetchWebPageToolId } from '../../contrib/chat/common/tools/buil import { SearchExtensionsToolId } from '../../contrib/extensions/common/searchExtensionsTool.js'; import { checkProposedApiEnabled, isProposedApiEnabled } from '../../services/extensions/common/extensions.js'; import { Dto, SerializableObjectWithBuffers } from '../../services/extensions/common/proxyIdentifier.js'; -import { ExtHostLanguageModelToolsShape, IMainContext, IToolDataDto, MainContext, MainThreadLanguageModelToolsShape } from './extHost.protocol.js'; +import { ExtHostLanguageModelToolsShape, IMainContext, IToolDataDto, IToolDefinitionDto, MainContext, MainThreadLanguageModelToolsShape } from './extHost.protocol.js'; import { ExtHostLanguageModels } from './extHostLanguageModels.js'; import * as typeConvert from './extHostTypeConverters.js'; @@ -172,14 +172,6 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape }); } - /** - * Check if a tool supports a specific model. - * @returns `true` if supported, `false` if not, `undefined` if no supportsModel implementation (treat as supported) - */ - async supportsModel(toolId: string, modelId: string, token: CancellationToken): Promise { - return this._proxy.$supportsModel(toolId, modelId, token); - } - async $invokeTool(dto: Dto, token: CancellationToken): Promise | SerializableObjectWithBuffers>> { const item = this._registeredTools.get(dto.toolId); if (!item) { @@ -287,21 +279,6 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape return undefined; } - async $supportsModel(toolId: string, modelId: string, token: CancellationToken): Promise { - const item = this._registeredTools.get(toolId); - if (!item) { - throw new Error(`Unknown tool ${toolId}`); - } - - // supportsModel is a proposed API - const supportsModelFn = item.tool.supportsModel; - if (supportsModelFn) { - return supportsModelFn(modelId); - } - - return undefined; - } - registerTool(extension: IExtensionDescription, id: string, tool: vscode.LanguageModelTool): IDisposable { this._registeredTools.set(id, { extension, tool }); this._proxy.$registerTool(id); @@ -311,4 +288,35 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape this._proxy.$unregisterTool(id); }); } + + registerToolDefinition(extension: IExtensionDescription, definition: vscode.LanguageModelToolDefinition, tool: vscode.LanguageModelTool): IDisposable { + checkProposedApiEnabled(extension, 'languageModelToolSupportsModel'); + + const id = definition.name; + + // Convert the definition to a DTO + const dto: IToolDefinitionDto = { + id, + displayName: definition.displayName, + toolReferenceName: definition.toolReferenceName, + userDescription: definition.userDescription, + modelDescription: definition.description, + inputSchema: definition.inputSchema as object, + source: { + type: 'extension', + label: extension.displayName ?? extension.name, + extensionId: extension.identifier, + }, + icon: typeConvert.IconPath.from(definition.icon), + models: definition.models, + }; + + this._registeredTools.set(id, { extension, tool }); + this._proxy.$registerToolWithDefinition(dto); + + return toDisposable(() => { + this._registeredTools.delete(id); + this._proxy.$unregisterTool(id); + }); + } } diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatActions.ts index 22ca202b3fe..2eeef50c655 100644 --- a/src/vs/workbench/contrib/chat/browser/actions/chatActions.ts +++ b/src/vs/workbench/contrib/chat/browser/actions/chatActions.ts @@ -1167,7 +1167,7 @@ registerAction2(class EditToolApproval extends Action2 { async run(accessor: ServicesAccessor, scope?: 'workspace' | 'profile' | 'session'): Promise { const confirmationService = accessor.get(ILanguageModelToolsConfirmationService); const toolsService = accessor.get(ILanguageModelToolsService); - confirmationService.manageConfirmationPreferences([...toolsService.getTools()], scope ? { defaultScope: scope } : undefined); + confirmationService.manageConfirmationPreferences([...toolsService.getAllToolsIncludingDisabled()], scope ? { defaultScope: scope } : undefined); } }); diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatToolActions.ts b/src/vs/workbench/contrib/chat/browser/actions/chatToolActions.ts index 4038fa42061..a8ad2887cac 100644 --- a/src/vs/workbench/contrib/chat/browser/actions/chatToolActions.ts +++ b/src/vs/workbench/contrib/chat/browser/actions/chatToolActions.ts @@ -188,7 +188,7 @@ class ConfigureToolsAction extends Action2 { }); try { - const result = await instaService.invokeFunction(showToolsPicker, placeholder, description, () => entriesMap.get(), widget.input.selectedLanguageModel?.metadata.id, cts.token); + const result = await instaService.invokeFunction(showToolsPicker, placeholder, description, () => entriesMap.get(), cts.token); if (result) { widget.input.selectedToolsModel.set(result, false); } diff --git a/src/vs/workbench/contrib/chat/browser/actions/chatToolPicker.ts b/src/vs/workbench/contrib/chat/browser/actions/chatToolPicker.ts index 7dc01e206e0..4c9b65ddfc6 100644 --- a/src/vs/workbench/contrib/chat/browser/actions/chatToolPicker.ts +++ b/src/vs/workbench/contrib/chat/browser/actions/chatToolPicker.ts @@ -194,7 +194,6 @@ export async function showToolsPicker( placeHolder: string, description?: string, getToolsEntries?: () => ReadonlyMap, - modelId?: string, token?: CancellationToken ): Promise | undefined> { @@ -215,40 +214,15 @@ export async function showToolsPicker( } } - // Pre-compute which tools support the model (if modelId is provided) - const supportedTools = new Set(); - if (modelId) { - const allTools = Array.from(toolsService.getTools()); - const checks = await Promise.all( - allTools.map(async (tool) => { - const supports = await toolsService.supportsModel(tool.id, modelId, CancellationToken.None); - // undefined means no supportsModel impl, treat as supported - // false means explicitly not supported - return { toolId: tool.id, supported: supports !== false }; - }) - ); - for (const { toolId, supported } of checks) { - if (supported) { - supportedTools.add(toolId); - } - } - } - - const isToolSupportedForModel = (toolId: string): boolean => { - // If no modelId specified, all tools are available - if (!modelId) { - return true; - } - return supportedTools.has(toolId); - }; + const contextKeyService = accessor.get(IContextKeyService); function computeItems(previousToolsEntries?: ReadonlyMap) { // Create default entries if none provided let toolsEntries = getToolsEntries ? new Map(getToolsEntries()) : undefined; if (!toolsEntries) { const defaultEntries = new Map(); - for (const tool of toolsService.getTools()) { - if (tool.canBeReferencedInPrompt && isToolSupportedForModel(tool.id)) { + for (const tool of toolsService.getTools(contextKeyService)) { + if (tool.canBeReferencedInPrompt) { defaultEntries.set(tool, false); } } @@ -431,9 +405,6 @@ export async function showToolsPicker( bucket.children.push(treeItem); const children = []; for (const tool of toolSet.getTools()) { - if (!isToolSupportedForModel(tool.id)) { - continue; - } const toolChecked = toolSetChecked || toolsEntries.get(tool) === true; const toolTreeItem = createToolTreeItemFromData(tool, toolChecked); children.push(toolTreeItem); @@ -443,13 +414,11 @@ export async function showToolsPicker( } } } - for (const tool of toolsService.getTools()) { + // getting potentially disabled tools is fine here because we filter `toolsEntries.has` + for (const tool of toolsService.getAllToolsIncludingDisabled()) { if (!tool.canBeReferencedInPrompt || !toolsEntries.has(tool)) { continue; } - if (!isToolSupportedForModel(tool.id)) { - continue; - } const bucket = getBucket(tool.source); if (!bucket) { continue; diff --git a/src/vs/workbench/contrib/chat/browser/attachments/chatAttachmentWidgets.ts b/src/vs/workbench/contrib/chat/browser/attachments/chatAttachmentWidgets.ts index df8d5a18201..d72d283da78 100644 --- a/src/vs/workbench/contrib/chat/browser/attachments/chatAttachmentWidgets.ts +++ b/src/vs/workbench/contrib/chat/browser/attachments/chatAttachmentWidgets.ts @@ -721,12 +721,13 @@ export class ToolSetOrToolItemAttachmentWidget extends AbstractChatAttachmentWid @ILanguageModelToolsService toolsService: ILanguageModelToolsService, @ICommandService commandService: ICommandService, @IOpenerService openerService: IOpenerService, - @IHoverService hoverService: IHoverService + @IHoverService hoverService: IHoverService, + @IContextKeyService contextKeyService: IContextKeyService, ) { super(attachment, options, container, contextResourceLabels, currentLanguageModel, commandService, openerService); - const toolOrToolSet = Iterable.find(toolsService.getTools(), tool => tool.id === attachment.id) ?? Iterable.find(toolsService.toolSets.get(), toolSet => toolSet.id === attachment.id); + const toolOrToolSet = Iterable.find(toolsService.getTools(contextKeyService), tool => tool.id === attachment.id) ?? Iterable.find(toolsService.toolSets.get(), toolSet => toolSet.id === attachment.id); let name = attachment.name; const icon = attachment.icon ?? Codicon.tools; diff --git a/src/vs/workbench/contrib/chat/browser/chat.contribution.ts b/src/vs/workbench/contrib/chat/browser/chat.contribution.ts index 4269ca6a4a8..db6e9c2282b 100644 --- a/src/vs/workbench/contrib/chat/browser/chat.contribution.ts +++ b/src/vs/workbench/contrib/chat/browser/chat.contribution.ts @@ -1058,7 +1058,7 @@ class ToolReferenceNamesContribution extends Disposable implements IWorkbenchCon private _updateToolReferenceNames(): void { const tools = - Array.from(this._languageModelToolsService.getTools()) + Array.from(this._languageModelToolsService.getAllToolsIncludingDisabled()) .filter((tool): tool is typeof tool & { toolReferenceName: string } => typeof tool.toolReferenceName === 'string') .sort((a, b) => a.toolReferenceName.localeCompare(b.toolReferenceName)); toolReferenceNameEnumValues.length = 0; diff --git a/src/vs/workbench/contrib/chat/browser/chatSetup/chatSetupProviders.ts b/src/vs/workbench/contrib/chat/browser/chatSetup/chatSetupProviders.ts index 37dbdf9e408..2bdd04b83dc 100644 --- a/src/vs/workbench/contrib/chat/browser/chatSetup/chatSetupProviders.ts +++ b/src/vs/workbench/contrib/chat/browser/chatSetup/chatSetupProviders.ts @@ -422,14 +422,14 @@ export class SetupAgent extends Disposable implements IChatAgentImplementation { } // check that tools other than setup. and internal tools are registered. - for (const tool of languageModelToolsService.getTools()) { + for (const tool of languageModelToolsService.getAllToolsIncludingDisabled()) { if (tool.id.startsWith('copilot_')) { return; // we have tools! } } return Event.toPromise(Event.filter(languageModelToolsService.onDidChangeTools, () => { - for (const tool of languageModelToolsService.getTools()) { + for (const tool of languageModelToolsService.getAllToolsIncludingDisabled()) { if (tool.id.startsWith('copilot_')) { return true; // we have tools! } diff --git a/src/vs/workbench/contrib/chat/browser/promptSyntax/promptToolsCodeLensProvider.ts b/src/vs/workbench/contrib/chat/browser/promptSyntax/promptToolsCodeLensProvider.ts index 4c9a0c63203..1e5ed8bf103 100644 --- a/src/vs/workbench/contrib/chat/browser/promptSyntax/promptToolsCodeLensProvider.ts +++ b/src/vs/workbench/contrib/chat/browser/promptSyntax/promptToolsCodeLensProvider.ts @@ -22,6 +22,7 @@ import { Range } from '../../../../../editor/common/core/range.js'; import { IEditorModel } from '../../../../../editor/common/editorCommon.js'; import { PromptHeaderAttributes } from '../../common/promptSyntax/promptFileParser.js'; import { isGithubTarget } from '../../common/promptSyntax/languageProviders/promptValidator.js'; +import { IContextKeyService } from '../../../../../platform/contextkey/common/contextkey.js'; class PromptToolsCodeLensProvider extends Disposable implements CodeLensProvider { @@ -32,7 +33,8 @@ class PromptToolsCodeLensProvider extends Disposable implements CodeLensProvider @IPromptsService private readonly promptsService: IPromptsService, @ILanguageFeaturesService private readonly languageService: ILanguageFeaturesService, @ILanguageModelToolsService private readonly languageModelToolsService: ILanguageModelToolsService, - @IInstantiationService private readonly instantiationService: IInstantiationService + @IInstantiationService private readonly instantiationService: IInstantiationService, + @IContextKeyService private readonly contextKeyService: IContextKeyService, ) { super(); @@ -84,12 +86,12 @@ class PromptToolsCodeLensProvider extends Disposable implements CodeLensProvider } private async updateTools(model: ITextModel, range: Range, selectedTools: readonly string[], target: string | undefined): Promise { - const selectedToolsNow = () => this.languageModelToolsService.toToolAndToolSetEnablementMap(selectedTools, target); + const selectedToolsNow = () => this.languageModelToolsService.toToolAndToolSetEnablementMap(selectedTools, target, this.contextKeyService); const newSelectedAfter = await this.instantiationService.invokeFunction(showToolsPicker, localize('placeholder', "Select tools"), undefined, selectedToolsNow); if (!newSelectedAfter) { return; } - await this.instantiationService.createInstance(PromptFileRewriter).rewriteTools(model, newSelectedAfter, range); + this.instantiationService.createInstance(PromptFileRewriter).rewriteTools(model, newSelectedAfter, range); } } diff --git a/src/vs/workbench/contrib/chat/browser/tools/languageModelToolsService.ts b/src/vs/workbench/contrib/chat/browser/tools/languageModelToolsService.ts index 5fe0cb5981d..e02e787523a 100644 --- a/src/vs/workbench/contrib/chat/browser/tools/languageModelToolsService.ts +++ b/src/vs/workbench/contrib/chat/browser/tools/languageModelToolsService.ts @@ -16,7 +16,7 @@ import { Emitter, Event } from '../../../../../base/common/event.js'; import { createMarkdownCommandLink, MarkdownString } from '../../../../../base/common/htmlContent.js'; import { Iterable } from '../../../../../base/common/iterator.js'; import { combinedDisposable, Disposable, DisposableStore, IDisposable, toDisposable } from '../../../../../base/common/lifecycle.js'; -import { derived, IObservable, IReader, observableFromEventOpts, ObservableSet } from '../../../../../base/common/observable.js'; +import { derived, derivedOpts, IObservable, IReader, observableFromEventOpts, ObservableSet, observableSignal, transaction } from '../../../../../base/common/observable.js'; import Severity from '../../../../../base/common/severity.js'; import { StopWatch } from '../../../../../base/common/stopwatch.js'; import { ThemeIcon } from '../../../../../base/common/themables.js'; @@ -35,11 +35,11 @@ import { IStorageService, StorageScope, StorageTarget } from '../../../../../pla import { ITelemetryService } from '../../../../../platform/telemetry/common/telemetry.js'; import { IExtensionService } from '../../../../services/extensions/common/extensions.js'; import { ChatContextKeys } from '../../common/actions/chatContextKeys.js'; -import { IVariableReference } from '../../common/chatModes.js'; -import { ChatToolInvocation } from '../../common/model/chatProgressTypes/chatToolInvocation.js'; -import { ConfirmedReason, IChatService, IChatToolInvocation, ToolConfirmKind } from '../../common/chatService/chatService.js'; import { ChatRequestToolReferenceEntry, toToolSetVariableEntry, toToolVariableEntry } from '../../common/attachments/chatVariableEntries.js'; +import { IVariableReference } from '../../common/chatModes.js'; +import { ConfirmedReason, IChatService, IChatToolInvocation, ToolConfirmKind } from '../../common/chatService/chatService.js'; import { ChatConfiguration } from '../../common/constants.js'; +import { ChatToolInvocation } from '../../common/model/chatProgressTypes/chatToolInvocation.js'; import { ILanguageModelToolsConfirmationService } from '../../common/tools/languageModelToolsConfirmationService.js'; import { CountTokensCallback, createToolSchemaUri, ILanguageModelToolsService, IPreparedToolInvocation, IToolAndToolSetEnablementMap, IToolData, IToolImpl, IToolInvocation, IToolResult, IToolResultInputOutputDetails, SpecedToolAliases, stringifyPromptTsxPart, ToolDataSource, ToolSet, VSCodeToolReference } from '../../common/tools/languageModelToolsService.js'; import { getToolConfirmationAlert } from '../accessibility/chatAccessibilityProvider.js'; @@ -258,15 +258,6 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo }); } - async supportsModel(toolId: string, modelId: string, token: CancellationToken): Promise { - const entry = this._tools.get(toolId); - if (!entry?.impl?.supportsModel) { - return undefined; - } - - return entry.impl.supportsModel(modelId, token); - } - registerTool(toolData: IToolData, tool: IToolImpl): IDisposable { return combinedDisposable( this.registerToolData(toolData), @@ -274,36 +265,52 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo ); } - getTools(includeDisabled?: boolean): Iterable { + getTools(contextKeyService: IContextKeyService): Iterable { const toolDatas = Iterable.map(this._tools.values(), i => i.data); const extensionToolsEnabled = this._configurationService.getValue(ChatConfiguration.ExtensionToolsEnabled); return Iterable.filter( toolDatas, toolData => { - const satisfiesWhenClause = includeDisabled || !toolData.when || this._contextKeyService.contextMatchesRules(toolData.when); + const satisfiesWhenClause = !toolData.when || contextKeyService.contextMatchesRules(toolData.when); const satisfiesExternalToolCheck = toolData.source.type !== 'extension' || !!extensionToolsEnabled; - const satisfiesPermittedCheck = includeDisabled || this.isPermitted(toolData); + const satisfiesPermittedCheck = this.isPermitted(toolData); return satisfiesWhenClause && satisfiesExternalToolCheck && satisfiesPermittedCheck; }); } - readonly toolsObservable = observableFromEventOpts({ equalsFn: arrayEqualsC() }, this.onDidChangeTools, () => Array.from(this.getTools())); + observeTools(contextKeyService: IContextKeyService): IObservable { + const meta = derived(reader => { + const signal = observableSignal('observeToolsContext'); + const trigger = () => transaction(tx => signal.trigger(tx)); + + const scheduler = reader.store.add(new RunOnceScheduler(trigger, 750)); + + this._register(this.onDidChangeTools(trigger)); + this._register(contextKeyService.onDidChangeContext(e => { + if (e.affectsSome(this._toolContextKeys) && !scheduler.isScheduled()) { + this._onDidChangeToolsScheduler.schedule(); + } + })); + + return signal; + }); + + return derivedOpts({ equalsFn: arrayEqualsC() }, reader => { + meta.read(reader).read(reader); + return Array.from(this.getTools(contextKeyService)); + }); + } + + getAllToolsIncludingDisabled(): Iterable { + return Iterable.map(this._tools.values(), i => i.data); + } getTool(id: string): IToolData | undefined { - return this._getToolEntry(id)?.data; + return this._tools.get(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, includeDisabled?: boolean): IToolData | undefined { - for (const tool of this.getTools(!!includeDisabled)) { + getToolByName(name: string): IToolData | undefined { + for (const tool of this.getAllToolsIncludingDisabled()) { if (tool.toolReferenceName === name) { return tool; } @@ -822,7 +829,7 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo * @param fullReferenceNames A list of tool or toolset by their full reference names that are enabled. * @returns A map of tool or toolset instances to their enablement state. */ - toToolAndToolSetEnablementMap(fullReferenceNames: readonly string[], _target: string | undefined): IToolAndToolSetEnablementMap { + toToolAndToolSetEnablementMap(fullReferenceNames: readonly string[], _target: string | undefined, contextKeyService: IContextKeyService | undefined): IToolAndToolSetEnablementMap { const toolOrToolSetNames = new Set(fullReferenceNames); const result = new Map(); for (const [tool, fullReferenceName] of this.toolsWithFullReferenceName.get()) { @@ -835,16 +842,21 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo } } } else { - if (!result.has(tool)) { // already set via an enabled toolset - const enabled = toolOrToolSetNames.has(fullReferenceName) - || Iterable.some(this.getToolAliases(tool, fullReferenceName), name => toolOrToolSetNames.has(name)) - || !!tool.legacyToolReferenceFullNames?.some(toolFullName => { - // enable tool if just the legacy tool set name is present - const index = toolFullName.lastIndexOf('/'); - return index !== -1 && toolOrToolSetNames.has(toolFullName.substring(0, index)); - }); - result.set(tool, enabled); + if (result.has(tool)) { // already set via an enabled toolset + continue; } + if (contextKeyService && tool.when && !contextKeyService.contextMatchesRules(tool.when)) { + continue; + } + + const enabled = toolOrToolSetNames.has(fullReferenceName) + || Iterable.some(this.getToolAliases(tool, fullReferenceName), name => toolOrToolSetNames.has(name)) + || !!tool.legacyToolReferenceFullNames?.some(toolFullName => { + // enable tool if just the legacy tool set name is present + const index = toolFullName.lastIndexOf('/'); + return index !== -1 && toolOrToolSetNames.has(toolFullName.substring(0, index)); + }); + result.set(tool, enabled); } } @@ -954,6 +966,12 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo return result; } + private readonly allToolsIncludingDisableObs = observableFromEventOpts( + { equalsFn: arrayEqualsC() }, + this.onDidChangeTools, + () => Array.from(this.getAllToolsIncludingDisabled()), + ); + readonly toolsWithFullReferenceName = derived<[IToolData | ToolSet, string][]>(reader => { const result: [IToolData | ToolSet, string][] = []; const coveredByToolSets = new Set(); @@ -966,7 +984,13 @@ export class LanguageModelToolsService extends Disposable implements ILanguageMo } } } - for (const tool of this.toolsObservable.read(reader)) { + for (const tool of this.allToolsIncludingDisableObs.read(reader)) { + // todo@connor4312/aeschil: this effectively hides model-specific tools + // for prompt referencing. Should we eventually enable this? (If so how?) + if (tool.when && !this._contextKeyService.contextMatchesRules(tool.when)) { + continue; + } + if (tool.canBeReferencedInPrompt && !coveredByToolSets.has(tool) && this.isPermitted(tool, reader)) { result.push([tool, getToolFullReferenceName(tool)]); } diff --git a/src/vs/workbench/contrib/chat/browser/tools/toolSetsContribution.ts b/src/vs/workbench/contrib/chat/browser/tools/toolSetsContribution.ts index 6b959e113ab..18b9eea7023 100644 --- a/src/vs/workbench/contrib/chat/browser/tools/toolSetsContribution.ts +++ b/src/vs/workbench/contrib/chat/browser/tools/toolSetsContribution.ts @@ -146,7 +146,7 @@ export class UserToolSetsContributions extends Disposable implements IWorkbenchC lifecycleService.when(LifecyclePhase.Restored) ]).then(() => this._initToolSets()); - const toolsObs = observableFromEvent(this, _languageModelToolsService.onDidChangeTools, () => Array.from(_languageModelToolsService.getTools())); + const toolsObs = observableFromEvent(this, _languageModelToolsService.onDidChangeTools, () => Array.from(_languageModelToolsService.getAllToolsIncludingDisabled())); const store = this._store.add(new DisposableStore()); this._store.add(autorun(r => { diff --git a/src/vs/workbench/contrib/chat/browser/widget/chatWidget.ts b/src/vs/workbench/contrib/chat/browser/widget/chatWidget.ts index 42f602ace9f..f7e686633f9 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/chatWidget.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/chatWidget.ts @@ -2606,7 +2606,7 @@ export class ChatWidget extends Disposable implements IChatWidget { // if not tools to enable are present, we are done if (tools !== undefined && this.input.currentModeKind === ChatModeKind.Agent) { - const enablementMap = this.toolsService.toToolAndToolSetEnablementMap(tools, Target.VSCode); + const enablementMap = this.toolsService.toToolAndToolSetEnablementMap(tools, Target.VSCode, this.contextKeyService); this.input.selectedToolsModel.set(enablementMap, true); } diff --git a/src/vs/workbench/contrib/chat/browser/widget/input/chatInputPart.ts b/src/vs/workbench/contrib/chat/browser/widget/input/chatInputPart.ts index 9669d88db00..40a6ec71bb8 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/input/chatInputPart.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/input/chatInputPart.ts @@ -330,6 +330,10 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge private withinEditSessionKey: IContextKey; private filePartOfEditSessionKey: IContextKey; private chatSessionHasOptions: IContextKey; + private modelIdKey: IContextKey; + private modelVendorKey: IContextKey; + private modelFamilyKey: IContextKey; + private modelVersionKey: IContextKey; private modelWidget: ModelPickerActionItem | undefined; private modeWidget: ModePickerActionItem | undefined; private chatSessionPickerWidgets: Map = new Map(); @@ -518,6 +522,10 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge this.withinEditSessionKey = ChatContextKeys.withinEditSessionDiff.bindTo(contextKeyService); this.filePartOfEditSessionKey = ChatContextKeys.filePartOfEditSession.bindTo(contextKeyService); this.chatSessionHasOptions = ChatContextKeys.chatSessionHasModels.bindTo(contextKeyService); + this.modelIdKey = ChatContextKeys.Model.id.bindTo(contextKeyService); + this.modelVendorKey = ChatContextKeys.Model.vendor.bindTo(contextKeyService); + this.modelFamilyKey = ChatContextKeys.Model.family.bindTo(contextKeyService); + this.modelVersionKey = ChatContextKeys.Model.version.bindTo(contextKeyService); const chatToolCount = ChatContextKeys.chatToolCount.bindTo(contextKeyService); @@ -935,6 +943,12 @@ export class ChatInputPart extends Disposable implements IHistoryNavigationWidge public setCurrentLanguageModel(model: ILanguageModelChatMetadataAndIdentifier) { this._currentLanguageModel = model; + // Update model context keys for tool filtering + this.modelIdKey.set(model.metadata.id); + this.modelVendorKey.set(model.metadata.vendor); + this.modelFamilyKey.set(model.metadata.family); + this.modelVersionKey.set(model.metadata.version ?? ''); + if (this.cachedDimensions) { // For quick chat and editor chat, relayout because the input may need to shrink to accomodate the model name this.layout(this.cachedDimensions.height, this.cachedDimensions.width); diff --git a/src/vs/workbench/contrib/chat/browser/widget/input/chatSelectedTools.ts b/src/vs/workbench/contrib/chat/browser/widget/input/chatSelectedTools.ts index 21a4c37ece4..58a969c15f6 100644 --- a/src/vs/workbench/contrib/chat/browser/widget/input/chatSelectedTools.ts +++ b/src/vs/workbench/contrib/chat/browser/widget/input/chatSelectedTools.ts @@ -8,17 +8,20 @@ import { Disposable } from '../../../../../../base/common/lifecycle.js'; import { derived, IObservable, ObservableMap } from '../../../../../../base/common/observable.js'; import { isObject } from '../../../../../../base/common/types.js'; import { URI } from '../../../../../../base/common/uri.js'; +import { IContextKeyService } from '../../../../../../platform/contextkey/common/contextkey.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 { UserSelectedTools } from '../../../common/participants/chatAgents.js'; import { IChatMode } from '../../../common/chatModes.js'; import { ChatModeKind } from '../../../common/constants.js'; -import { ILanguageModelToolsService, IToolAndToolSetEnablementMap, IToolData, ToolSet } from '../../../common/tools/languageModelToolsService.js'; +import { UserSelectedTools } from '../../../common/participants/chatAgents.js'; import { PromptsStorage } from '../../../common/promptSyntax/service/promptsService.js'; +import { ILanguageModelToolsService, IToolAndToolSetEnablementMap, IToolData, ToolSet } from '../../../common/tools/languageModelToolsService.js'; import { PromptFileRewriter } from '../../promptSyntax/promptFileRewriter.js'; +// todo@connor4312/bhavyaus: make tools key off displayName so model-specific tool +// enablement can stick between models with different underlying tool definitions type ToolEnablementStates = { readonly toolSets: ReadonlyMap; readonly tools: ReadonlyMap; @@ -98,9 +101,11 @@ export class ChatSelectedTools extends Disposable { private readonly _globalState: ObservableMemento; private readonly _sessionStates = new ObservableMap(); + private readonly _currentTools: IObservable; constructor( private readonly _mode: IObservable, + @IContextKeyService private readonly _contextKeyService: IContextKeyService, @ILanguageModelToolsService private readonly _toolsService: ILanguageModelToolsService, @IStorageService _storageService: IStorageService, @IInstantiationService private readonly _instantiationService: IInstantiationService, @@ -115,10 +120,12 @@ export class ChatSelectedTools extends Disposable { }); this._globalState = this._store.add(globalStateMemento(StorageScope.PROFILE, StorageTarget.MACHINE, _storageService)); + this._currentTools = _toolsService.observeTools(this._contextKeyService); } /** * All tools and tool sets with their enabled state. + * Tools are filtered based on the current model context. */ public readonly entriesMap: IObservable = derived(r => { const map = new Map(); @@ -130,13 +137,14 @@ export class ChatSelectedTools extends Disposable { const modeTools = currentMode.customTools?.read(r); if (modeTools) { const target = currentMode.target?.read(r); - currentMap = ToolEnablementStates.fromMap(this._toolsService.toToolAndToolSetEnablementMap(modeTools, target)); + currentMap = ToolEnablementStates.fromMap(this._toolsService.toToolAndToolSetEnablementMap(modeTools, target, this._contextKeyService)); } } if (!currentMap) { currentMap = this._globalState.read(r); } - for (const tool of this._toolsService.toolsObservable.read(r)) { + // Use getTools with contextKeyService to filter tools by current model + for (const tool of this._currentTools.read(r)) { if (tool.canBeReferencedInPrompt) { map.set(tool, currentMap.tools.get(tool.id) !== false); // if unknown, it's enabled } diff --git a/src/vs/workbench/contrib/chat/common/actions/chatContextKeys.ts b/src/vs/workbench/contrib/chat/common/actions/chatContextKeys.ts index 62efc919657..a6501cc4108 100644 --- a/src/vs/workbench/contrib/chat/common/actions/chatContextKeys.ts +++ b/src/vs/workbench/contrib/chat/common/actions/chatContextKeys.ts @@ -85,6 +85,13 @@ export namespace ChatContextKeys { toolsCount: new RawContextKey('toolsCount', 0, { type: 'number', description: localize('toolsCount', "The count of tools available in the chat.") }) }; + export const Model = { + id: new RawContextKey('chatModelId', '', { type: 'string', description: localize('chatModelId', "The identifier of the currently selected language model.") }), + vendor: new RawContextKey('chatModelVendor', '', { type: 'string', description: localize('chatModelVendor', "The vendor of the currently selected language model.") }), + family: new RawContextKey('chatModelFamily', '', { type: 'string', description: localize('chatModelFamily', "The family of the currently selected language model.") }), + version: new RawContextKey('chatModelVersion', '', { type: 'string', description: localize('chatModelVersion', "The version of the currently selected language model.") }), + }; + export const Modes = { hasCustomChatModes: new RawContextKey('chatHasCustomAgents', false, { type: 'boolean', description: localize('chatHasAgents', "True when the chat has custom agents available.") }), agentModeDisabledByPolicy: new RawContextKey('chatAgentModeDisabledByPolicy', false, { type: 'boolean', description: localize('chatAgentModeDisabledByPolicy', "True when agent mode is disabled by organization policy.") }), diff --git a/src/vs/workbench/contrib/chat/common/promptSyntax/languageProviders/promptValidator.ts b/src/vs/workbench/contrib/chat/common/promptSyntax/languageProviders/promptValidator.ts index b7e91c439c1..3e31dee613b 100644 --- a/src/vs/workbench/contrib/chat/common/promptSyntax/languageProviders/promptValidator.ts +++ b/src/vs/workbench/contrib/chat/common/promptSyntax/languageProviders/promptValidator.ts @@ -92,7 +92,7 @@ export class PromptValidator { if (body.variableReferences.length && !isGitHubTarget) { const headerTools = promptAST.header?.tools; const headerTarget = promptAST.header?.target; - const headerToolsMap = headerTools ? this.languageModelToolsService.toToolAndToolSetEnablementMap(headerTools, headerTarget) : undefined; + const headerToolsMap = headerTools ? this.languageModelToolsService.toToolAndToolSetEnablementMap(headerTools, headerTarget, undefined) : undefined; const available = new Set(this.languageModelToolsService.getFullReferenceNames()); const deprecatedNames = this.languageModelToolsService.getDeprecatedFullReferenceNames(); diff --git a/src/vs/workbench/contrib/chat/common/tools/builtinTools/runSubagentTool.ts b/src/vs/workbench/contrib/chat/common/tools/builtinTools/runSubagentTool.ts index e9001142f86..b5de8005f78 100644 --- a/src/vs/workbench/contrib/chat/common/tools/builtinTools/runSubagentTool.ts +++ b/src/vs/workbench/contrib/chat/common/tools/builtinTools/runSubagentTool.ts @@ -162,7 +162,7 @@ export class RunSubagentTool extends Disposable implements IToolImpl { const modeCustomTools = mode.customTools?.get(); if (modeCustomTools) { // Convert the mode's custom tools (array of qualified names) to UserSelectedTools format - const enablementMap = this.languageModelToolsService.toToolAndToolSetEnablementMap(modeCustomTools, mode.target?.get()); + const enablementMap = this.languageModelToolsService.toToolAndToolSetEnablementMap(modeCustomTools, mode.target?.get(), undefined); // Convert enablement map to UserSelectedTools (Record) modeTools = {}; for (const [tool, enabled] of enablementMap) { @@ -277,7 +277,7 @@ export class RunSubagentTool extends Disposable implements IToolImpl { const fullReferenceNames = tools.map(tool => this.languageModelToolsService.getFullReferenceName(tool)); if (fullReferenceNames.length > 0) { - enabledTools = this.languageModelToolsService.toToolAndToolSetEnablementMap(fullReferenceNames, undefined); + enabledTools = this.languageModelToolsService.toToolAndToolSetEnablementMap(fullReferenceNames, undefined, undefined); } } diff --git a/src/vs/workbench/contrib/chat/common/tools/languageModelToolsContribution.ts b/src/vs/workbench/contrib/chat/common/tools/languageModelToolsContribution.ts index ac3977f1424..9850dee5ee2 100644 --- a/src/vs/workbench/contrib/chat/common/tools/languageModelToolsContribution.ts +++ b/src/vs/workbench/contrib/chat/common/tools/languageModelToolsContribution.ts @@ -329,7 +329,7 @@ export class LanguageModelToolsExtensionPointHandler implements IWorkbenchContri const toolSets: ToolSet[] = []; for (const toolName of toolSet.tools) { - const toolObj = languageModelToolsService.getToolByName(toolName, true); + const toolObj = languageModelToolsService.getToolByName(toolName); if (toolObj) { tools.push(toolObj); continue; diff --git a/src/vs/workbench/contrib/chat/common/tools/languageModelToolsService.ts b/src/vs/workbench/contrib/chat/common/tools/languageModelToolsService.ts index 2bf02155cd7..12c3413ee33 100644 --- a/src/vs/workbench/contrib/chat/common/tools/languageModelToolsService.ts +++ b/src/vs/workbench/contrib/chat/common/tools/languageModelToolsService.ts @@ -17,7 +17,7 @@ import { ThemeIcon } from '../../../../../base/common/themables.js'; import { URI } from '../../../../../base/common/uri.js'; import { Location } from '../../../../../editor/common/languages.js'; import { localize } from '../../../../../nls.js'; -import { ContextKeyExpression } from '../../../../../platform/contextkey/common/contextkey.js'; +import { ContextKeyExpression, IContextKeyService } from '../../../../../platform/contextkey/common/contextkey.js'; import { ExtensionIdentifier } from '../../../../../platform/extensions/common/extensions.js'; import { ByteSize } from '../../../../../platform/files/common/files.js'; import { createDecorator } from '../../../../../platform/instantiation/common/instantiation.js'; @@ -29,6 +29,17 @@ import { ChatRequestToolReferenceEntry } from '../attachments/chatVariableEntrie import { LanguageModelPartAudience } from '../languageModels.js'; import { PromptElementJSON, stringifyPromptElementJSON } from './promptTsxTypes.js'; +/** + * Selector for matching language models by vendor, family, version, or id. + * Used to filter tools to specific models or model families. + */ +export interface ILanguageModelChatSelector { + readonly vendor?: string; + readonly family?: string; + readonly version?: string; + readonly id?: string; +} + export interface IToolData { readonly id: string; readonly source: ToolDataSource; @@ -52,6 +63,11 @@ export interface IToolData { readonly canRequestPreApproval?: boolean; /** True if this tool might ask for post-approval */ readonly canRequestPostApproval?: boolean; + /** + * Model selectors that this tool is available for. + * If defined, the tool is only available when the selected model matches one of the selectors. + */ + readonly models?: readonly ILanguageModelChatSelector[]; } export interface IToolProgressStep { @@ -298,7 +314,6 @@ export interface IPreparedToolInvocation { export interface IToolImpl { invoke(invocation: IToolInvocation, countTokens: CountTokensCallback, progress: ToolProgress, token: CancellationToken): Promise; prepareToolInvocation?(context: IToolInvocationPreparationContext, token: CancellationToken): Promise; - supportsModel?(modelId: string, token: CancellationToken): Promise; } export type IToolAndToolSetEnablementMap = ReadonlyMap; @@ -369,11 +384,43 @@ export interface ILanguageModelToolsService { registerToolData(toolData: IToolData): IDisposable; registerToolImplementation(id: string, tool: IToolImpl): IDisposable; registerTool(toolData: IToolData, tool: IToolImpl): IDisposable; - supportsModel(toolId: string, modelId: string, token: CancellationToken): Promise; - getTools(): Iterable; - readonly toolsObservable: IObservable; + + /** + * Get all tools currently enabled (matching the context key service's context). + * @param contextKeyService The context key service to evaluate `when` clauses against + */ + getTools(contextKeyService: IContextKeyService): Iterable; + + /** + * Creats an observable of enabled tools in the context. Note the observable + * should be created and reused, not created per reader, for example: + * + * ``` + * const toolsObs = toolsService.observeTools(contextKeyService); + * autorun(reader => { + * const tools = toolsObs.read(reader); + * ... + * }); + * ``` + */ + observeTools(contextKeyService: IContextKeyService): IObservable; + + /** + * Get all registered tools regardless of enablement state. + * Use this for configuration UIs, completions, etc. where all tools should be visible. + */ + getAllToolsIncludingDisabled(): Iterable; + + /** + * Get a tool by its ID. Does not check when clauses. + */ getTool(id: string): IToolData | undefined; - getToolByName(name: string, includeDisabled?: boolean): IToolData | undefined; + + /** + * Get a tool by its reference name. Does not check when clauses. + */ + getToolByName(name: string): IToolData | undefined; + invokeTool(invocation: IToolInvocation, countTokens: CountTokensCallback, token: CancellationToken): Promise; cancelToolCallsForRequest(requestId: string): void; /** Flush any pending tool updates to the extension hosts. */ @@ -390,7 +437,19 @@ export interface ILanguageModelToolsService { getToolByFullReferenceName(fullReferenceName: string): IToolData | ToolSet | undefined; getDeprecatedFullReferenceNames(): Map>; - toToolAndToolSetEnablementMap(fullReferenceNames: readonly string[], target: string | undefined): IToolAndToolSetEnablementMap; + /** + * Gets the enablement maps based on the given set of references. + * @param fullReferenceNames The full reference names of the tools and tool sets to enable. + * @param target Optional target to filter tools by. + * @param contextKeyService Context key service to evaluate tool enablement. + * If undefined is passed, all tools will be returned, even if normally disabled. + */ + toToolAndToolSetEnablementMap( + fullReferenceNames: readonly string[], + target: string | undefined, + contextKeyService: IContextKeyService | undefined, + ): IToolAndToolSetEnablementMap; + toFullReferenceNames(map: IToolAndToolSetEnablementMap): string[]; toToolReferences(variableReferences: readonly IVariableReference[]): ChatRequestToolReferenceEntry[]; } diff --git a/src/vs/workbench/contrib/chat/test/browser/tools/languageModelToolsService.test.ts b/src/vs/workbench/contrib/chat/test/browser/tools/languageModelToolsService.test.ts index 8e2751e2fe0..0b05b0cae12 100644 --- a/src/vs/workbench/contrib/chat/test/browser/tools/languageModelToolsService.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/tools/languageModelToolsService.test.ts @@ -276,7 +276,7 @@ suite('LanguageModelToolsService', () => { store.add(service.registerToolData(toolData2)); store.add(service.registerToolData(toolData3)); - const tools = Array.from(service.getTools()); + const tools = Array.from(service.getTools(contextKeyService)); assert.strictEqual(tools.length, 2); assert.strictEqual(tools[0].id, 'testTool2'); assert.strictEqual(tools[1].id, 'testTool3'); @@ -314,8 +314,8 @@ suite('LanguageModelToolsService', () => { store.add(service.registerToolData(toolData2)); store.add(service.registerToolData(toolData3)); - assert.strictEqual(service.getToolByName('testTool1'), undefined); - assert.strictEqual(service.getToolByName('testTool1', true)?.id, 'testTool1'); + // getToolByName searches all tools regardless of when clause + assert.strictEqual(service.getToolByName('testTool1')?.id, 'testTool1'); assert.strictEqual(service.getToolByName('testTool2')?.id, 'testTool2'); assert.strictEqual(service.getToolByName('testTool3')?.id, 'testTool3'); }); @@ -572,7 +572,7 @@ suite('LanguageModelToolsService', () => { // Test with enabled tool { const fullReferenceNames = ['tool1RefName']; - const result1 = service.toToolAndToolSetEnablementMap(fullReferenceNames, undefined); + const result1 = service.toToolAndToolSetEnablementMap(fullReferenceNames, undefined, undefined); assert.strictEqual(result1.size, numOfTools, `Expected ${numOfTools} tools and tool sets`); assert.strictEqual([...result1.entries()].filter(([_, enabled]) => enabled).length, 1, 'Expected 1 tool to be enabled'); assert.strictEqual(result1.get(tool1), true, 'tool1 should be enabled'); @@ -584,7 +584,7 @@ suite('LanguageModelToolsService', () => { // Test with multiple enabled tools { const fullReferenceNames = ['my.extension/extTool1RefName', 'mcpToolSetRefName/*', 'internalToolSetRefName/internalToolSetTool1RefName']; - const result1 = service.toToolAndToolSetEnablementMap(fullReferenceNames, undefined); + const result1 = service.toToolAndToolSetEnablementMap(fullReferenceNames, undefined, undefined); assert.strictEqual(result1.size, numOfTools, `Expected ${numOfTools} tools and tool sets`); assert.strictEqual([...result1.entries()].filter(([_, enabled]) => enabled).length, 4, 'Expected 4 tools to be enabled'); assert.strictEqual(result1.get(extTool1), true, 'extTool1 should be enabled'); @@ -597,7 +597,7 @@ suite('LanguageModelToolsService', () => { } // Test with all enabled tools, redundant names { - const result1 = service.toToolAndToolSetEnablementMap(allFullReferenceNames, undefined); + const result1 = service.toToolAndToolSetEnablementMap(allFullReferenceNames, undefined, undefined); assert.strictEqual(result1.size, numOfTools, `Expected ${numOfTools} tools and tool sets`); assert.strictEqual([...result1.entries()].filter(([_, enabled]) => enabled).length, 11, 'Expected 11 tools to be enabled'); // +3 including the vscode, execute, read toolsets @@ -608,7 +608,7 @@ suite('LanguageModelToolsService', () => { // Test with no enabled tools { const fullReferenceNames: string[] = []; - const result1 = service.toToolAndToolSetEnablementMap(fullReferenceNames, undefined); + const result1 = service.toToolAndToolSetEnablementMap(fullReferenceNames, undefined, undefined); assert.strictEqual(result1.size, numOfTools, `Expected ${numOfTools} tools and tool sets`); assert.strictEqual([...result1.entries()].filter(([_, enabled]) => enabled).length, 0, 'Expected 0 tools to be enabled'); @@ -618,7 +618,7 @@ suite('LanguageModelToolsService', () => { // Test with unknown tool { const fullReferenceNames: string[] = ['unknownToolRefName']; - const result1 = service.toToolAndToolSetEnablementMap(fullReferenceNames, undefined); + const result1 = service.toToolAndToolSetEnablementMap(fullReferenceNames, undefined, undefined); assert.strictEqual(result1.size, numOfTools, `Expected ${numOfTools} tools and tool sets`); assert.strictEqual([...result1.entries()].filter(([_, enabled]) => enabled).length, 0, 'Expected 0 tools to be enabled'); @@ -628,7 +628,7 @@ suite('LanguageModelToolsService', () => { // Test with legacy tool names { const fullReferenceNames: string[] = ['extTool1RefName', 'mcpToolSetRefName', 'internalToolSetTool1RefName']; - const result1 = service.toToolAndToolSetEnablementMap(fullReferenceNames, undefined); + const result1 = service.toToolAndToolSetEnablementMap(fullReferenceNames, undefined, undefined); assert.strictEqual(result1.size, numOfTools, `Expected ${numOfTools} tools and tool sets`); assert.strictEqual([...result1.entries()].filter(([_, enabled]) => enabled).length, 4, 'Expected 4 tools to be enabled'); assert.strictEqual(result1.get(extTool1), true, 'extTool1 should be enabled'); @@ -643,7 +643,7 @@ suite('LanguageModelToolsService', () => { // Test with tool in user tool set { const fullReferenceNames = ['Tool2 Display Name']; - const result1 = service.toToolAndToolSetEnablementMap(fullReferenceNames, undefined); + const result1 = service.toToolAndToolSetEnablementMap(fullReferenceNames, undefined, undefined); assert.strictEqual(result1.size, numOfTools, `Expected ${numOfTools} tools and tool sets`); assert.strictEqual([...result1.entries()].filter(([_, enabled]) => enabled).length, 2, 'Expected 1 tool and user tool set to be enabled'); assert.strictEqual(result1.get(tool2), true, 'tool2 should be enabled'); @@ -670,7 +670,7 @@ suite('LanguageModelToolsService', () => { // Test enabling the tool set const enabledNames = [toolData1].map(t => service.getFullReferenceName(t)); - const result = service.toToolAndToolSetEnablementMap(enabledNames, undefined); + const result = service.toToolAndToolSetEnablementMap(enabledNames, undefined, undefined); assert.strictEqual(result.get(toolData1), true, 'individual tool should be enabled'); @@ -730,7 +730,7 @@ suite('LanguageModelToolsService', () => { // Test enabling the tool set const enabledNames = [toolSet, toolData1].map(t => service.getFullReferenceName(t)); - const result = service.toToolAndToolSetEnablementMap(enabledNames, undefined); + const result = service.toToolAndToolSetEnablementMap(enabledNames, undefined, undefined); assert.strictEqual(result.get(toolData1), true, 'individual tool should be enabled'); assert.strictEqual(result.get(toolData2), false); @@ -765,7 +765,7 @@ suite('LanguageModelToolsService', () => { // Test with non-existent tool names const enabledNames = [toolData, unregisteredToolData].map(t => service.getFullReferenceName(t)); - const result = service.toToolAndToolSetEnablementMap(enabledNames, undefined); + const result = service.toToolAndToolSetEnablementMap(enabledNames, undefined, undefined); assert.strictEqual(result.get(toolData), true, 'existing tool should be enabled'); // Non-existent tools should not appear in the result map @@ -814,7 +814,7 @@ suite('LanguageModelToolsService', () => { // Test 1: Using legacy tool reference name should enable the tool { - const result = service.toToolAndToolSetEnablementMap(['oldToolName'], undefined); + const result = service.toToolAndToolSetEnablementMap(['oldToolName'], undefined, undefined); assert.strictEqual(result.get(toolWithLegacy), true, 'tool should be enabled via legacy name'); const fullReferenceNames = service.toFullReferenceNames(result); @@ -823,7 +823,7 @@ suite('LanguageModelToolsService', () => { // Test 2: Using another legacy tool reference name should also work { - const result = service.toToolAndToolSetEnablementMap(['deprecatedToolName'], undefined); + const result = service.toToolAndToolSetEnablementMap(['deprecatedToolName'], undefined, undefined); assert.strictEqual(result.get(toolWithLegacy), true, 'tool should be enabled via another legacy name'); const fullReferenceNames = service.toFullReferenceNames(result); @@ -832,7 +832,7 @@ suite('LanguageModelToolsService', () => { // Test 3: Using legacy toolset name should enable the entire toolset { - const result = service.toToolAndToolSetEnablementMap(['oldToolSet'], undefined); + const result = service.toToolAndToolSetEnablementMap(['oldToolSet'], undefined, undefined); assert.strictEqual(result.get(toolSetWithLegacy), true, 'toolset should be enabled via legacy name'); assert.strictEqual(result.get(toolInSet), true, 'tool in set should be enabled when set is enabled via legacy name'); @@ -842,7 +842,7 @@ suite('LanguageModelToolsService', () => { // Test 4: Using deprecated toolset name should also work { - const result = service.toToolAndToolSetEnablementMap(['deprecatedToolSet'], undefined); + const result = service.toToolAndToolSetEnablementMap(['deprecatedToolSet'], undefined, undefined); assert.strictEqual(result.get(toolSetWithLegacy), true, 'toolset should be enabled via another legacy name'); assert.strictEqual(result.get(toolInSet), true, 'tool in set should be enabled when set is enabled via legacy name'); @@ -852,7 +852,7 @@ suite('LanguageModelToolsService', () => { // Test 5: Mix of current and legacy names { - const result = service.toToolAndToolSetEnablementMap(['newToolRef', 'oldToolSet'], undefined); + const result = service.toToolAndToolSetEnablementMap(['newToolRef', 'oldToolSet'], undefined, undefined); assert.strictEqual(result.get(toolWithLegacy), true, 'tool should be enabled via current name'); assert.strictEqual(result.get(toolSetWithLegacy), true, 'toolset should be enabled via legacy name'); assert.strictEqual(result.get(toolInSet), true, 'tool in set should be enabled'); @@ -863,7 +863,7 @@ suite('LanguageModelToolsService', () => { // Test 6: Using legacy names and current names together (redundant but should work) { - const result = service.toToolAndToolSetEnablementMap(['newToolRef', 'oldToolName', 'deprecatedToolName'], undefined); + const result = service.toToolAndToolSetEnablementMap(['newToolRef', 'oldToolName', 'deprecatedToolName'], undefined, undefined); assert.strictEqual(result.get(toolWithLegacy), true, 'tool should be enabled (redundant legacy names should not cause issues)'); const fullReferenceNames = service.toFullReferenceNames(result); @@ -889,7 +889,7 @@ suite('LanguageModelToolsService', () => { // Test 1: Using the full legacy name should enable the tool { - const result = service.toToolAndToolSetEnablementMap(['oldToolSet/oldToolName'], undefined); + const result = service.toToolAndToolSetEnablementMap(['oldToolSet/oldToolName'], undefined, undefined); assert.strictEqual(result.get(toolWithOrphanedToolSet), true, 'tool should be enabled via full legacy name'); const fullReferenceNames = service.toFullReferenceNames(result); @@ -898,7 +898,7 @@ suite('LanguageModelToolsService', () => { // Test 2: Using just the orphaned toolset name should also enable the tool { - const result = service.toToolAndToolSetEnablementMap(['oldToolSet'], undefined); + const result = service.toToolAndToolSetEnablementMap(['oldToolSet'], undefined, undefined); assert.strictEqual(result.get(toolWithOrphanedToolSet), true, 'tool should be enabled via orphaned toolset name'); const fullReferenceNames = service.toFullReferenceNames(result); @@ -918,7 +918,7 @@ suite('LanguageModelToolsService', () => { store.add(service.registerToolData(anotherToolFromOrphanedSet)); { - const result = service.toToolAndToolSetEnablementMap(['oldToolSet'], undefined); + const result = service.toToolAndToolSetEnablementMap(['oldToolSet'], undefined, undefined); assert.strictEqual(result.get(toolWithOrphanedToolSet), true, 'first tool should be enabled via orphaned toolset name'); assert.strictEqual(result.get(anotherToolFromOrphanedSet), true, 'second tool should also be enabled via orphaned toolset name'); @@ -939,7 +939,7 @@ suite('LanguageModelToolsService', () => { store.add(service.registerToolData(unrelatedTool)); { - const result = service.toToolAndToolSetEnablementMap(['oldToolSet'], undefined); + const result = service.toToolAndToolSetEnablementMap(['oldToolSet'], undefined, undefined); assert.strictEqual(result.get(toolWithOrphanedToolSet), true, 'tool from oldToolSet should be enabled'); assert.strictEqual(result.get(anotherToolFromOrphanedSet), true, 'another tool from oldToolSet should be enabled'); assert.strictEqual(result.get(unrelatedTool), false, 'tool from different toolset should NOT be enabled'); @@ -967,7 +967,7 @@ suite('LanguageModelToolsService', () => { store.add(newToolSetWithSameName.addTool(toolInRecreatedSet)); { - const result = service.toToolAndToolSetEnablementMap(['oldToolSet'], undefined); + const result = service.toToolAndToolSetEnablementMap(['oldToolSet'], undefined, undefined); // Now 'oldToolSet' should enable BOTH the recreated toolset AND the tools with legacy names pointing to oldToolSet assert.strictEqual(result.get(newToolSetWithSameName), true, 'recreated toolset should be enabled'); assert.strictEqual(result.get(toolInRecreatedSet), true, 'tool in recreated set should be enabled'); @@ -1064,7 +1064,7 @@ suite('LanguageModelToolsService', () => { { const toolNames = ['custom-agent', 'shell']; - const result = service.toToolAndToolSetEnablementMap(toolNames, undefined); + const result = service.toToolAndToolSetEnablementMap(toolNames, undefined, undefined); assert.strictEqual(result.get(service.executeToolSet), true, 'execute should be enabled'); assert.strictEqual(result.get(agentSet), true, 'agent should be enabled'); @@ -1079,7 +1079,7 @@ suite('LanguageModelToolsService', () => { } { const toolNames = ['github/*', 'playwright/*']; - const result = service.toToolAndToolSetEnablementMap(toolNames, undefined); + const result = service.toToolAndToolSetEnablementMap(toolNames, undefined, undefined); assert.strictEqual(result.get(githubMcpToolSet), true, 'githubMcpToolSet should be enabled'); assert.strictEqual(result.get(playwrightMcpToolSet), true, 'playwrightMcpToolSet should be enabled'); @@ -1095,7 +1095,7 @@ suite('LanguageModelToolsService', () => { { // the speced names should work and not be altered const toolNames = ['github/create_branch', 'playwright/browser_click']; - const result = service.toToolAndToolSetEnablementMap(toolNames, undefined); + const result = service.toToolAndToolSetEnablementMap(toolNames, undefined, undefined); assert.strictEqual(result.get(githubMcpTool1), true, 'githubMcpTool1 should be enabled'); assert.strictEqual(result.get(playwrightMcpTool1), true, 'playwrightMcpTool1 should be enabled'); @@ -1111,7 +1111,7 @@ suite('LanguageModelToolsService', () => { { // using the old MCP full names should also work const toolNames = ['github/github-mcp-server/*', 'microsoft/playwright-mcp/*']; - const result = service.toToolAndToolSetEnablementMap(toolNames, undefined); + const result = service.toToolAndToolSetEnablementMap(toolNames, undefined, undefined); assert.strictEqual(result.get(githubMcpToolSet), true, 'githubMcpToolSet should be enabled'); assert.strictEqual(result.get(playwrightMcpToolSet), true, 'playwrightMcpToolSet should be enabled'); @@ -1126,7 +1126,7 @@ suite('LanguageModelToolsService', () => { { // using the old MCP full names should also work const toolNames = ['github/github-mcp-server/create_branch', 'microsoft/playwright-mcp/browser_click']; - const result = service.toToolAndToolSetEnablementMap(toolNames, undefined); + const result = service.toToolAndToolSetEnablementMap(toolNames, undefined, undefined); assert.strictEqual(result.get(githubMcpTool1), true, 'githubMcpTool1 should be enabled'); assert.strictEqual(result.get(playwrightMcpTool1), true, 'playwrightMcpTool1 should be enabled'); @@ -1142,7 +1142,7 @@ suite('LanguageModelToolsService', () => { { // using the latest MCP full names should also work const toolNames = ['io.github.github/github-mcp-server/*', 'com.microsoft/playwright-mcp/*']; - const result = service.toToolAndToolSetEnablementMap(toolNames, undefined); + const result = service.toToolAndToolSetEnablementMap(toolNames, undefined, undefined); assert.strictEqual(result.get(githubMcpToolSet), true, 'githubMcpToolSet should be enabled'); assert.strictEqual(result.get(playwrightMcpToolSet), true, 'playwrightMcpToolSet should be enabled'); @@ -1158,7 +1158,7 @@ suite('LanguageModelToolsService', () => { { // using the latest MCP full names should also work const toolNames = ['io.github.github/github-mcp-server/create_branch', 'com.microsoft/playwright-mcp/browser_click']; - const result = service.toToolAndToolSetEnablementMap(toolNames, undefined); + const result = service.toToolAndToolSetEnablementMap(toolNames, undefined, undefined); assert.strictEqual(result.get(githubMcpTool1), true, 'githubMcpTool1 should be enabled'); assert.strictEqual(result.get(playwrightMcpTool1), true, 'playwrightMcpTool1 should be enabled'); @@ -1174,7 +1174,7 @@ suite('LanguageModelToolsService', () => { { // using the old MCP full names should also work const toolNames = ['github-mcp-server/create_branch']; - const result = service.toToolAndToolSetEnablementMap(toolNames, undefined); + const result = service.toToolAndToolSetEnablementMap(toolNames, undefined, undefined); assert.strictEqual(result.get(githubMcpTool1), true, 'githubMcpTool1 should be enabled'); const fullReferenceNames = service.toFullReferenceNames(result).sort(); @@ -1780,12 +1780,12 @@ suite('LanguageModelToolsService', () => { store.add(service.registerToolData(disabledTool)); store.add(service.registerToolData(enabledTool)); - const enabledTools = Array.from(service.getTools()); + const enabledTools = Array.from(service.getTools(contextKeyService)); assert.strictEqual(enabledTools.length, 1, 'Should only return enabled tools'); assert.strictEqual(enabledTools[0].id, 'enabledTool'); - const allTools = Array.from(service.getTools(true)); - assert.strictEqual(allTools.length, 2, 'includeDisabled should return all tools'); + const allTools = Array.from(service.getAllToolsIncludingDisabled()); + assert.strictEqual(allTools.length, 2, 'getAllToolsIncludingDisabled should return all tools'); }); test('tool registration duplicate error', () => { @@ -2025,7 +2025,7 @@ suite('LanguageModelToolsService', () => { // Enable the MCP toolset { const enabledNames = [mcpToolSet].map(t => service.getFullReferenceName(t)); - const result = service.toToolAndToolSetEnablementMap(enabledNames, undefined); + const result = service.toToolAndToolSetEnablementMap(enabledNames, undefined, undefined); assert.strictEqual(result.get(mcpToolSet), true, 'MCP toolset should be enabled'); // Ensure the toolset is in the map assert.strictEqual(result.get(mcpTool), true, 'MCP tool should be enabled when its toolset is enabled'); // Ensure the tool is in the map @@ -2036,7 +2036,7 @@ suite('LanguageModelToolsService', () => { // Enable a tool from the MCP toolset { const enabledNames = [mcpTool].map(t => service.getFullReferenceName(t, mcpToolSet)); - const result = service.toToolAndToolSetEnablementMap(enabledNames, undefined); + const result = service.toToolAndToolSetEnablementMap(enabledNames, undefined, undefined); assert.strictEqual(result.get(mcpToolSet), false, 'MCP toolset should be disabled'); // Ensure the toolset is in the map assert.strictEqual(result.get(mcpTool), true, 'MCP tool should be enabled'); // Ensure the tool is in the map diff --git a/src/vs/workbench/contrib/chat/test/browser/widget/input/chatSelectedTools.test.ts b/src/vs/workbench/contrib/chat/test/browser/widget/input/chatSelectedTools.test.ts index 3462416b317..2598c65a1e1 100644 --- a/src/vs/workbench/contrib/chat/test/browser/widget/input/chatSelectedTools.test.ts +++ b/src/vs/workbench/contrib/chat/test/browser/widget/input/chatSelectedTools.test.ts @@ -4,6 +4,7 @@ *--------------------------------------------------------------------------------------------*/ import assert from 'assert'; +import { IContextKeyService } from '../../../../../../../platform/contextkey/common/contextkey.js'; import { TestConfigurationService } from '../../../../../../../platform/configuration/test/common/testConfigurationService.js'; import { ContextKeyService } from '../../../../../../../platform/contextkey/browser/contextKeyService.js'; import { workbenchInstantiationService } from '../../../../../../test/browser/workbenchTestServices.js'; @@ -27,6 +28,7 @@ suite('ChatSelectedTools', () => { let toolsService: ILanguageModelToolsService; let selectedTools: ChatSelectedTools; + let contextKeyService: IContextKeyService; setup(() => { @@ -40,6 +42,7 @@ suite('ChatSelectedTools', () => { store.add(instaService); toolsService = instaService.get(ILanguageModelToolsService); + contextKeyService = instaService.get(IContextKeyService); selectedTools = store.add(instaService.createInstance(ChatSelectedTools, constObservable(ChatMode.Agent))); }); @@ -95,7 +98,7 @@ suite('ChatSelectedTools', () => { store.add(toolset.addTool(toolData2)); store.add(toolset.addTool(toolData3)); - assert.strictEqual(Iterable.length(toolsService.getTools()), 3); + assert.strictEqual(Iterable.length(toolsService.getTools(contextKeyService)), 3); const size = Iterable.length(toolset.getTools()); assert.strictEqual(size, 3); @@ -159,7 +162,7 @@ suite('ChatSelectedTools', () => { store.add(toolset.addTool(toolData2)); store.add(toolset.addTool(toolData3)); - assert.strictEqual(Iterable.length(toolsService.getTools()), 3); + assert.strictEqual(Iterable.length(toolsService.getTools(contextKeyService)), 3); const size = Iterable.length(toolset.getTools()); assert.strictEqual(size, 3); diff --git a/src/vs/workbench/contrib/chat/test/common/tools/mockLanguageModelToolsService.ts b/src/vs/workbench/contrib/chat/test/common/tools/mockLanguageModelToolsService.ts index be847426f80..9076a8defee 100644 --- a/src/vs/workbench/contrib/chat/test/common/tools/mockLanguageModelToolsService.ts +++ b/src/vs/workbench/contrib/chat/test/common/tools/mockLanguageModelToolsService.ts @@ -9,6 +9,7 @@ import { Event } from '../../../../../../base/common/event.js'; import { Disposable, IDisposable } from '../../../../../../base/common/lifecycle.js'; import { constObservable, IObservable } from '../../../../../../base/common/observable.js'; import { ThemeIcon } from '../../../../../../base/common/themables.js'; +import { IContextKeyService } from '../../../../../../platform/contextkey/common/contextkey.js'; import { IProgressStep } from '../../../../../../platform/progress/common/progress.js'; import { IVariableReference } from '../../../common/chatModes.js'; import { ChatRequestToolReferenceEntry } from '../../../common/attachments/chatVariableEntries.js'; @@ -62,25 +63,27 @@ export class MockLanguageModelToolsService implements ILanguageModelToolsService return Disposable.None; } - async supportsModel(toolId: string, modelId: string, token: CancellationToken): Promise { - return undefined; - } - registerTool(toolData: IToolData, tool: IToolImpl): IDisposable { return Disposable.None; } - getTools(): Iterable { + getTools(contextKeyService: IContextKeyService): Iterable { return []; } - toolsObservable: IObservable = constObservable([]); + getAllToolsIncludingDisabled(): Iterable { + return []; + } getTool(id: string): IToolData | undefined { return undefined; } - getToolByName(name: string, includeDisabled?: boolean): IToolData | undefined { + observeTools(contextKeyService: IContextKeyService): IObservable { + return constObservable([]); + } + + getToolByName(name: string): IToolData | undefined { return undefined; } diff --git a/src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/alternativeRecommendation.ts b/src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/alternativeRecommendation.ts index da3f3980009..5b9b450b093 100644 --- a/src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/alternativeRecommendation.ts +++ b/src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/alternativeRecommendation.ts @@ -3,6 +3,7 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ +import { IContextKeyService } from '../../../../../platform/contextkey/common/contextkey.js'; import type { ILanguageModelToolsService } from '../../../chat/common/tools/languageModelToolsService.js'; let previouslyRecommededInSession = false; @@ -26,8 +27,8 @@ const terminalCommands: { commands: RegExp[]; tags: string[] }[] = [ } ]; -export function getRecommendedToolsOverRunInTerminal(commandLine: string, languageModelToolsService: ILanguageModelToolsService): string | undefined { - const tools = languageModelToolsService.getTools(); +export function getRecommendedToolsOverRunInTerminal(commandLine: string, contextKeyService: IContextKeyService, languageModelToolsService: ILanguageModelToolsService): string | undefined { + const tools = languageModelToolsService.getTools(contextKeyService); if (!tools || previouslyRecommededInSession) { return; } diff --git a/src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/tools/runInTerminalTool.ts b/src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/tools/runInTerminalTool.ts index fdfbcb4308f..d11ccc599bf 100644 --- a/src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/tools/runInTerminalTool.ts +++ b/src/vs/workbench/contrib/terminalContrib/chatAgentTools/browser/tools/runInTerminalTool.ts @@ -58,6 +58,7 @@ import { isNumber, isString } from '../../../../../../base/common/types.js'; import { ChatConfiguration } from '../../../../chat/common/constants.js'; import { IChatWidgetService } from '../../../../chat/browser/chat.js'; import { TerminalChatCommandId } from '../../../chat/browser/terminalChat.js'; +import { IContextKeyService } from '../../../../../../platform/contextkey/common/contextkey.js'; // #region Tool data @@ -308,6 +309,7 @@ export class RunInTerminalTool extends Disposable implements IToolImpl { @ITerminalService private readonly _terminalService: ITerminalService, @IWorkspaceContextService private readonly _workspaceContextService: IWorkspaceContextService, @IChatWidgetService private readonly _chatWidgetService: IChatWidgetService, + @IContextKeyService private readonly _contextKeyService: IContextKeyService, ) { super(); @@ -407,7 +409,7 @@ export class RunInTerminalTool extends Disposable implements IToolImpl { // HACK: Exit early if there's an alternative recommendation, this is a little hacky but // it's the current mechanism for re-routing terminal tool calls to something else. - const alternativeRecommendation = getRecommendedToolsOverRunInTerminal(args.command, this._languageModelToolsService); + const alternativeRecommendation = getRecommendedToolsOverRunInTerminal(args.command, this._contextKeyService, this._languageModelToolsService); if (alternativeRecommendation) { toolSpecificData.alternativeRecommendation = alternativeRecommendation; return { diff --git a/src/vscode-dts/vscode.proposed.chatParticipantAdditions.d.ts b/src/vscode-dts/vscode.proposed.chatParticipantAdditions.d.ts index 89f4504e68e..a6e645fbd1b 100644 --- a/src/vscode-dts/vscode.proposed.chatParticipantAdditions.d.ts +++ b/src/vscode-dts/vscode.proposed.chatParticipantAdditions.d.ts @@ -412,18 +412,6 @@ declare module 'vscode' { * Fired when the set of tools on a chat request changes. */ export const onDidChangeChatRequestTools: Event; - - /** - * Invoke a tool by its full information object rather than just name. - * This allows disambiguation when multiple tools have the same name - * (e.g., from different MCP servers or model-specific implementations). - * - * @param tool The tool information object, typically obtained from {@link lm.tools}. - * @param options The options to use when invoking the tool. - * @param token A cancellation token. - * @returns The result of the tool invocation. - */ - export function invokeTool(tool: LanguageModelToolInformation, options: LanguageModelToolInvocationOptions, token?: CancellationToken): Thenable; } export class LanguageModelToolExtensionSource { diff --git a/src/vscode-dts/vscode.proposed.languageModelToolSupportsModel.d.ts b/src/vscode-dts/vscode.proposed.languageModelToolSupportsModel.d.ts index b4f2af57e70..015ae3eeed3 100644 --- a/src/vscode-dts/vscode.proposed.languageModelToolSupportsModel.d.ts +++ b/src/vscode-dts/vscode.proposed.languageModelToolSupportsModel.d.ts @@ -5,17 +5,58 @@ declare module 'vscode' { - export interface LanguageModelTool { + export interface LanguageModelToolDefinition extends LanguageModelToolInformation { /** - * Called to check if this tool supports a specific language model. If this method is not implemented, - * the tool is assumed to support all models. - * - * This method allows extensions to dynamically determine which models a tool can work with, - * enabling fine-grained control over tool availability based on model capabilities. - * - * @param modelId The identifier of the language model (e.g., 'gpt-4o', 'claude-3-5-sonnet') - * @returns `true` if the tool supports the given model, `false` otherwise + * Display name for the tool. */ - supportsModel?(modelId: string): Thenable; + displayName: string; + /** + * Name of the tools that can users can reference in the prompt. If not + * provided, the tool will not be able to be referenced. Must not contain whitespace. + */ + toolReferenceName?: string; + /** + * Description for the tool shown to the user. + */ + userDescription?: string; + /** + * Icon for the tool shown to the user. + */ + icon?: IconPath; + /** + * If defined, the tool will only be available for language models that match + * the selector. + */ + models?: LanguageModelChatSelector[]; + } + + export namespace lm { + /** + * Registers a language model tool along with its definition. Unlike {@link lm.registerTool}, + * this does not require the tool to be present first in the extension's `package.json` contributions. + * + * Multiple tools may be registered with the the same name using the API. In any given context, + * the most specific tool (based on the {@link LanguageModelToolDefinition.models}) will be used. + * + * @param definition The definition of the tool to register. + * @param tool The implementation of the tool. + * @returns A disposable that unregisters the tool when disposed. + */ + export function registerToolDefinition( + definition: LanguageModelToolDefinition, + tool: LanguageModelTool, + ): Disposable; + + /** + * Invoke a tool by its full information object rather than just name. + * This allows disambiguation when multiple tools have the same name + * (e.g., from different MCP servers or model-specific implementations). + * + * @param tool The tool information object, typically obtained from {@link lm.tools}. + * @param options The options to use when invoking the tool. + * @param token A cancellation token. + * @returns The result of the tool invocation. + */ + export function invokeTool(tool: LanguageModelToolInformation, options: LanguageModelToolInvocationOptions, token?: CancellationToken): Thenable; } }