diff --git a/src/vs/workbench/contrib/chat/browser/chatContentParts/chatConfirmationWidget.ts b/src/vs/workbench/contrib/chat/browser/chatContentParts/chatConfirmationWidget.ts index 761e83ff51c..7bdb9b23499 100644 --- a/src/vs/workbench/contrib/chat/browser/chatContentParts/chatConfirmationWidget.ts +++ b/src/vs/workbench/contrib/chat/browser/chatContentParts/chatConfirmationWidget.ts @@ -54,7 +54,7 @@ abstract class BaseChatConfirmationWidget extends Disposable { this._domNode = elements.root; this.markdownRenderer = this.instantiationService.createInstance(MarkdownRenderer, {}); - const renderedTitle = this._register(this.markdownRenderer.render(new MarkdownString(title), { + const renderedTitle = this._register(this.markdownRenderer.render(new MarkdownString(title, { supportThemeIcons: true }), { asyncRenderCallback: () => this._onDidChangeHeight.fire(), })); elements.title.append(renderedTitle.element); diff --git a/src/vs/workbench/contrib/mcp/browser/mcpCommands.ts b/src/vs/workbench/contrib/mcp/browser/mcpCommands.ts index 364c5c2b4b4..0514c67931b 100644 --- a/src/vs/workbench/contrib/mcp/browser/mcpCommands.ts +++ b/src/vs/workbench/contrib/mcp/browser/mcpCommands.ts @@ -6,11 +6,11 @@ import { reset } from '../../../../base/browser/dom.js'; import { renderLabelWithIcons } from '../../../../base/browser/ui/iconLabel/iconLabels.js'; import { Codicon } from '../../../../base/common/codicons.js'; -import { groupBy } from '../../../../base/common/collections.js'; +import { diffSets, groupBy } from '../../../../base/common/collections.js'; import { Event } from '../../../../base/common/event.js'; import { KeyMod, KeyCode } from '../../../../base/common/keyCodes.js'; import { Disposable, DisposableStore } from '../../../../base/common/lifecycle.js'; -import { autorun, derived } from '../../../../base/common/observable.js'; +import { autorun, derived, transaction } from '../../../../base/common/observable.js'; import { assertType } from '../../../../base/common/types.js'; import { ILocalizedString, localize, localize2 } from '../../../../nls.js'; import { IActionViewItemService } from '../../../../platform/actions/browser/actionViewItemService.js'; @@ -26,7 +26,7 @@ import { CHAT_CATEGORY } from '../../chat/browser/actions/chatActions.js'; import { ChatAgentLocation } from '../../chat/common/chatAgents.js'; import { ChatContextKeys } from '../../chat/common/chatContextKeys.js'; import { McpContextKeys } from '../common/mcpContextKeys.js'; -import { IMcpService, IMcpTool, McpConnectionState } from '../common/mcpTypes.js'; +import { IMcpServer, IMcpService, IMcpTool, McpConnectionState } from '../common/mcpTypes.js'; // acroynms do not get localized const category: ILocalizedString = { @@ -180,12 +180,12 @@ export class AttachMCPToolsAction extends Action2 { f1: false, category: CHAT_CATEGORY, precondition: ContextKeyExpr.and( - McpContextKeys.serverCount.greater(0), + McpContextKeys.toolsCount.greater(0), ChatContextKeys.location.isEqualTo(ChatAgentLocation.EditingSession) ), menu: { when: ContextKeyExpr.and( - McpContextKeys.serverCount.greater(0), + McpContextKeys.toolsCount.greater(0), ChatContextKeys.location.isEqualTo(ChatAgentLocation.EditingSession) ), id: MenuId.ChatInputAttachmentToolbar, @@ -204,48 +204,131 @@ export class AttachMCPToolsAction extends Action2 { const quickPickService = accessor.get(IQuickInputService); const mcpService = accessor.get(IMcpService); - type IToolPick = IQuickPickItem & { tool: IMcpTool }; - const picks: (IToolPick | IQuickPickSeparator)[] = []; + type ToolPick = IQuickPickItem & { picked: boolean; tool: IMcpTool; parent: ServerPick }; + type ServerPick = IQuickPickItem & { picked: boolean; server: IMcpServer; toolPicks: ToolPick[] }; + type McpPick = ToolPick | ServerPick; + + function isServerPick(obj: any): obj is ServerPick { + return Boolean((obj as ServerPick).server); + } + function isToolPick(obj: any): obj is ToolPick { + return Boolean((obj as ToolPick).tool); + } + + const store = new DisposableStore(); + const picker = store.add(quickPickService.createQuickPick({ useSeparators: true })); + + const picks: (McpPick | IQuickPickSeparator)[] = []; for (const server of mcpService.servers.get()) { + const tools = server.tools.get(); + + if (tools.length === 0) { + continue; + } picks.push({ type: 'separator', - label: server.definition.label + label: server.collection.label }); - for (const tool of server.tools.get()) { - picks.push({ - type: 'item', - label: tool.definition.name, - detail: tool.definition.description, - tooltip: tool.definition.description, - picked: tool.enabled.get(), + const item: ServerPick = { + server, + type: 'item', + label: `$(server) ${server.definition.label}`, + description: McpConnectionState.toString(server.state.get()), + picked: tools.some(tool => tool.enabled.get()), + toolPicks: [] + }; + + picks.push(item); + + for (const tool of tools) { + const toolItem: ToolPick = { tool, - }); + parent: item, + type: 'item', + label: `$(tools) ${tool.definition.name}`, + description: tool.definition.description, + picked: tool.enabled.get() + }; + item.toolPicks.push(toolItem); + picks.push(toolItem); } } - const result = await quickPickService.pick(picks, { - placeHolder: localize('placeholder', "Select tools that are available to chat"), - canPickMany: true - }); + picker.placeholder = localize('placeholder', "Select tools that are available to chat"); + picker.canSelectMany = true; - if (!result) { - return; - } + let lastSelectedItems = new Set(); + let ignoreEvent = false; - const seen = new Set(); - for (const item of result) { - item.tool.updateEnablement(true); - seen.add(item.tool); - } - - for (const pick of picks) { - if (pick.type === 'item' && !seen.has(pick.tool)) { - pick.tool.updateEnablement(false); + const _update = () => { + ignoreEvent = true; + try { + const items = picks.filter((p): p is McpPick => p.type === 'item' && Boolean(p.picked)); + lastSelectedItems = new Set(items); + picker.items = picks; + picker.selectedItems = items; + } finally { + ignoreEvent = false; } - } + }; + + _update(); + picker.show(); + + store.add(picker.onDidChangeSelection(selectedPicks => { + if (ignoreEvent) { + return; + } + + const { added, removed } = diffSets(lastSelectedItems, new Set(selectedPicks)); + + for (const item of added) { + item.picked = true; + + if (isServerPick(item)) { + // add server -> add back tools + for (const toolPick of item.toolPicks) { + toolPick.picked = true; + } + } else if (isToolPick(item)) { + // add server when tool is picked + item.parent.picked = true; + } + } + + for (const item of removed) { + item.picked = false; + + if (isServerPick(item)) { + // removed server -> remove tools + for (const toolPick of item.toolPicks) { + toolPick.picked = false; + } + } else if (isToolPick(item) && item.parent.toolPicks.every(child => !child.picked)) { + // remove LAST tool -> remove server + item.parent.picked = false; + } + } + + transaction(tx => { + for (const item of picks) { + if (isToolPick(item)) { + item.tool.updateEnablement(item.picked, tx); + } + } + }); + + _update(); + })); + + + await Promise.race([Event.toPromise(Event.any(picker.onDidAccept, picker.onDidHide))]); + + store.dispose(); + } } diff --git a/src/vs/workbench/contrib/mcp/common/mcpContextKeys.ts b/src/vs/workbench/contrib/mcp/common/mcpContextKeys.ts index c49fd058f55..545c678a5c1 100644 --- a/src/vs/workbench/contrib/mcp/common/mcpContextKeys.ts +++ b/src/vs/workbench/contrib/mcp/common/mcpContextKeys.ts @@ -15,11 +15,13 @@ import { IMcpService } from './mcpTypes.js'; export namespace McpContextKeys { export const serverCount = new RawContextKey('mcp.serverCount', undefined, { type: 'number', description: localize('mcp.serverCount.description', "Context key that has the number of registered MCP servers") }); + export const toolsCount = new RawContextKey('mcp.toolsCount', undefined, { type: 'number', description: localize('mcp.toolsCount.description', "Context key that has the number of registered MCP tools") }); } export class McpContextKeysController extends Disposable implements IWorkbenchContribution { - public static readonly ID = 'workbench.contrib.mcp.contextKey'; + + static readonly ID = 'workbench.contrib.mcp.contextKey'; constructor( @IMcpService mcpService: IMcpService, @@ -28,9 +30,11 @@ export class McpContextKeysController extends Disposable implements IWorkbenchCo super(); const ctxServerCount = McpContextKeys.serverCount.bindTo(contextKeyService); + const ctxToolsCount = McpContextKeys.toolsCount.bindTo(contextKeyService); this._store.add(autorun(r => { ctxServerCount.set(mcpService.servers.read(r).length); + ctxToolsCount.set(mcpService.servers.read(r).reduce((count, server) => count + server.tools.read(r).length, 0)); })); } } diff --git a/src/vs/workbench/contrib/mcp/common/mcpServer.ts b/src/vs/workbench/contrib/mcp/common/mcpServer.ts index f884386642d..082d62363f1 100644 --- a/src/vs/workbench/contrib/mcp/common/mcpServer.ts +++ b/src/vs/workbench/contrib/mcp/common/mcpServer.ts @@ -230,8 +230,8 @@ export class McpTool implements IMcpTool { this.id = `${_server.definition.id}_${definition.name}`.replaceAll('.', '_'); } - updateEnablement(value: boolean): void { - this._enabled.set(value, undefined); + updateEnablement(value: boolean, tx?: ITransaction): void { + this._enabled.set(value, tx); } call(params: Record, token?: CancellationToken): Promise { diff --git a/src/vs/workbench/contrib/mcp/common/mcpService.ts b/src/vs/workbench/contrib/mcp/common/mcpService.ts index f46e7fd57b5..f08e7af0f48 100644 --- a/src/vs/workbench/contrib/mcp/common/mcpService.ts +++ b/src/vs/workbench/contrib/mcp/common/mcpService.ts @@ -4,11 +4,13 @@ *--------------------------------------------------------------------------------------------*/ import { RunOnceScheduler } from '../../../../base/common/async.js'; +import { MarkdownString } from '../../../../base/common/htmlContent.js'; import { Disposable, DisposableStore, MutableDisposable, toDisposable } from '../../../../base/common/lifecycle.js'; import { autorun, autorunWithStore, derived, IObservable, observableValue } from '../../../../base/common/observable.js'; import { localize } from '../../../../nls.js'; import { IContextKeyService, RawContextKey } from '../../../../platform/contextkey/common/contextkey.js'; import { IInstantiationService } from '../../../../platform/instantiation/common/instantiation.js'; +import { IProductService } from '../../../../platform/product/common/productService.js'; import { StorageScope } from '../../../../platform/storage/common/storage.js'; import { ILanguageModelToolsService, IToolResult } from '../../chat/common/languageModelToolsService.js'; import { IMcpRegistry } from './mcpRegistryTypes.js'; @@ -29,7 +31,8 @@ export class McpService extends Disposable implements IMcpService { @IInstantiationService instantiationService: IInstantiationService, @IMcpRegistry private readonly _mcpRegistry: IMcpRegistry, @ILanguageModelToolsService toolsService: ILanguageModelToolsService, - @IContextKeyService contextKeyService: IContextKeyService + @IContextKeyService contextKeyService: IContextKeyService, + @IProductService productService: IProductService ) { super(); @@ -96,6 +99,7 @@ export class McpService extends Disposable implements IMcpService { ctxInst.set(tool.enabled.read(reader)); })); + newStore.add(toolsService.registerToolData({ id: tool.id, displayName: tool.definition.name, @@ -107,11 +111,20 @@ export class McpService extends Disposable implements IMcpService { newStore.add(toolsService.registerToolImplementation(tool.id, { async prepareToolInvocation(parameters, token) { + + const mcpToolWarning = localize( + 'mcp.tool.warning', + "MCP servers or malicious conversation content may attempt to misuse '{0}' through the installed tools. Please carefully review any requested actions.", + productService.nameShort + ); + return { confirmationMessages: { - title: localize('aaa', "Run tool from {0}", server.definition.label), - message: localize('ddd', "Do you allow to run `{0}` from `{1}`?", tool.definition.name, server.definition.label) - } + title: localize('msg.title', "Run `{0}` from $(server) `{1}` (MCP server)", tool.definition.name, server.definition.label), + message: new MarkdownString(localize('msg.msg', "{0}\n\nInput:\n\n```json\n{1}\n```\n\n$(warning) {2}", tool.definition.description, JSON.stringify(parameters, undefined, 2), mcpToolWarning), { supportThemeIcons: true }) + }, + invocationMessage: new MarkdownString(localize('msg.run', "Running `{0}`", tool.definition.name, server.definition.label)), + pastTenseMessage: new MarkdownString(localize('msg.ran', "Ran `{0}` ", tool.definition.name, server.definition.label)) }; }, @@ -133,6 +146,8 @@ export class McpService extends Disposable implements IMcpService { } } + // result.toolResultMessage = new MarkdownString(localize('reuslt.pattern', "```json\n{0}\n```", JSON.stringify(callResult, undefined, 2))); + return result; }, })); diff --git a/src/vs/workbench/contrib/mcp/common/mcpTypes.ts b/src/vs/workbench/contrib/mcp/common/mcpTypes.ts index 5baf2ca137d..ec1b2e00c4b 100644 --- a/src/vs/workbench/contrib/mcp/common/mcpTypes.ts +++ b/src/vs/workbench/contrib/mcp/common/mcpTypes.ts @@ -7,7 +7,7 @@ import { assertNever } from '../../../../base/common/assert.js'; import { CancellationToken } from '../../../../base/common/cancellation.js'; import { IDisposable } from '../../../../base/common/lifecycle.js'; import { equals as objectsEqual } from '../../../../base/common/objects.js'; -import { IObservable } from '../../../../base/common/observable.js'; +import { IObservable, ITransaction } from '../../../../base/common/observable.js'; import { URI } from '../../../../base/common/uri.js'; import { localize } from '../../../../nls.js'; import { ConfigurationTarget } from '../../../../platform/configuration/common/configuration.js'; @@ -113,7 +113,7 @@ export interface IMcpTool { readonly enabled: IObservable; - updateEnablement(value: boolean): void; + updateEnablement(value: boolean, tx?: ITransaction): void; /** * Calls a tool