Skip to content

Commit 73944df

Browse files
connor4312connor4312
authored
mcp: initial sampling implementation (#249946)
* mcp: initial sampling implementation - MCP servers get a prompt when they ask for sampling. Usage in tool calls (i.e. a sampling request within a concurrent tool request) is a separate prompt from 'ambient' usage. Sampling is 'just' a request to the client and can technically happen at any time. - By default MCP servers will get the default model in chat. But that is configurable by users. Multiple models can be configured and if they are then the model will be picked based on the rough 'model hints' we get from the server. - Not sure where the config should live. I have it in user settings in this demo. Having it in mcp.json would be a natural place to put it, but that doesn't work for extension-provided servers or servers from other config files (e.g. claude desktop configs) so that could not be the only place. - There are some limitations from MCP such as modelcontextprotocol/modelcontextprotocol#91 * rationalize configuration
2 parents 95fa805 + 2dc20eb commit 73944df

20 files changed

+569
-79
lines changed

‎src/vs/workbench/api/common/extHostMcp.ts

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,8 @@ import { AUTH_SERVER_METADATA_DISCOVERY_PATH, getDefaultMetadataForUrl, getMetad
2020
import { URI } from '../../../base/common/uri.js';
2121
import { MCP } from '../../contrib/mcp/common/modelContextProtocol.js';
2222
import { CancellationError } from '../../../base/common/errors.js';
23+
import { ConfigurationTarget } from '../../../platform/configuration/common/configuration.js';
24+
import { IExtHostInitDataService } from './extHostInitDataService.js';
2325

2426
export const IExtHostMpcService = createDecorator<IExtHostMpcService>('IExtHostMpcService');
2527

@@ -38,6 +40,7 @@ export class ExtHostMcpService extends Disposable implements IExtHostMpcService
3840

3941
constructor(
4042
@IExtHostRpcService extHostRpc: IExtHostRpcService,
43+
@IExtHostInitDataService private readonly _extHostInitData: IExtHostInitDataService
4144
) {
4245
super();
4346
this._proxy = extHostRpc.getProxy(MainContext.MainThreadMcp);
@@ -105,6 +108,7 @@ export class ExtHostMcpService extends Disposable implements IExtHostMpcService
105108
scope: StorageScope.WORKSPACE,
106109
canResolveLaunch: typeof provider.resolveMcpServerDefinition === 'function',
107110
extensionId: extension.identifier.value,
111+
configTarget: this._extHostInitData.remote.isRemote ? ConfigurationTarget.USER_REMOTE : ConfigurationTarget.USER,
108112
};
109113

110114
const update = async () => {
@@ -124,7 +128,7 @@ export class ExtHostMcpService extends Disposable implements IExtHostMpcService
124128
id,
125129
label: item.label,
126130
cacheNonce: item.version,
127-
launch: Convert.McpServerDefinition.from(item)
131+
launch: Convert.McpServerDefinition.from(item),
128132
});
129133
}
130134

‎src/vs/workbench/api/node/extHostMcpNode.ts

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,12 +15,14 @@ import { McpConnectionState, McpServerLaunch, McpServerTransportStdio, McpServer
1515
import { ExtHostMcpService } from '../common/extHostMcp.js';
1616
import { IExtHostRpcService } from '../common/extHostRpcService.js';
1717
import * as path from '../../../base/common/path.js';
18+
import { IExtHostInitDataService } from '../common/extHostInitDataService.js';
1819

1920
export class NodeExtHostMpcService extends ExtHostMcpService {
2021
constructor(
2122
@IExtHostRpcService extHostRpc: IExtHostRpcService,
23+
@IExtHostInitDataService initDataService: IExtHostInitDataService,
2224
) {
23-
super(extHostRpc);
25+
super(extHostRpc, initDataService);
2426
}
2527

2628
private nodeServers = new Map<number, {

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

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ import { EditorExtensions, IEditorFactoryRegistry } from '../../../common/editor
3030
import { IWorkbenchAssignmentService } from '../../../services/assignment/common/assignmentService.js';
3131
import { mcpSchemaId } from '../../../services/configuration/common/configuration.js';
3232
import { IEditorResolverService, RegisteredEditorPriority } from '../../../services/editor/common/editorResolverService.js';
33-
import { allDiscoverySources, discoverySourceLabel, mcpConfigurationSection, mcpDiscoverySection, mcpEnabledSection, mcpSchemaExampleServers } from '../../mcp/common/mcpConfiguration.js';
33+
import { allDiscoverySources, discoverySourceLabel, mcpConfigurationSection, mcpDiscoverySection, mcpEnabledSection, mcpSchemaExampleServers, mcpServerSamplingSection } from '../../mcp/common/mcpConfiguration.js';
3434
import { ChatAgentNameService, ChatAgentService, IChatAgentNameService, IChatAgentService } from '../common/chatAgents.js';
3535
import { CodeMapperService, ICodeMapperService } from '../common/chatCodeMapperService.js';
3636
import '../common/chatColors.js';
@@ -250,6 +250,33 @@ configurationRegistry.registerConfiguration({
250250
defaultValue: false
251251
}
252252
},
253+
[mcpServerSamplingSection]: {
254+
type: 'object',
255+
description: nls.localize('chat.mcp.serverSampling', "Configures which models are exposed to MCP servers for sampling (making model requests in the background). This setting can be edited in a graphical way under the `{0}` command.", 'MCP: ' + nls.localize('mcp.list', 'List Servers')),
256+
scope: ConfigurationScope.RESOURCE,
257+
additionalProperties: {
258+
type: 'object',
259+
properties: {
260+
allowedDuringChat: {
261+
type: 'boolean',
262+
description: nls.localize('chat.mcp.serverSampling.allowedDuringChat', "Whether this server is make sampling requests during its tool calls in a chat session."),
263+
default: true,
264+
},
265+
allowedOutsideChat: {
266+
type: 'boolean',
267+
description: nls.localize('chat.mcp.serverSampling.allowedOutsideChat', "Whether this server is allowed to make sampling requests outside of a chat session."),
268+
default: false,
269+
},
270+
allowedModels: {
271+
type: 'array',
272+
items: {
273+
type: 'string',
274+
description: nls.localize('chat.mcp.serverSampling.model', "A model the MCP server has access to."),
275+
},
276+
}
277+
}
278+
},
279+
},
253280
[mcpConfigurationSection]: {
254281
type: 'object',
255282
default: {

‎src/vs/workbench/contrib/mcp/browser/mcp.contribution.ts

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,10 +31,11 @@ import { IMcpDevModeDebugging, McpDevModeDebugging } from '../common/mcpDevMode.
3131
import { McpRegistry } from '../common/mcpRegistry.js';
3232
import { IMcpRegistry } from '../common/mcpRegistryTypes.js';
3333
import { McpResourceFilesystem } from '../common/mcpResourceFilesystem.js';
34+
import { McpSamplingService } from '../common/mcpSamplingService.js';
3435
import { McpService } from '../common/mcpService.js';
35-
import { HasInstalledMcpServersContext, IMcpService, IMcpWorkbenchService, InstalledMcpServersViewId, McpServersGalleryEnabledContext } from '../common/mcpTypes.js';
36+
import { HasInstalledMcpServersContext, IMcpSamplingService, IMcpService, IMcpWorkbenchService, InstalledMcpServersViewId, McpServersGalleryEnabledContext } from '../common/mcpTypes.js';
3637
import { McpAddContextContribution } from './mcpAddContextContribution.js';
37-
import { AddConfigurationAction, EditStoredInput, InstallFromActivation, ListMcpServerCommand, McpBrowseCommand, McpBrowseResourcesCommand, MCPServerActionRendering, McpServerOptionsCommand, RemoveStoredInput, ResetMcpCachedTools, ResetMcpTrustCommand, RestartServer, ShowConfiguration, ShowOutput, StartServer, StopServer } from './mcpCommands.js';
38+
import { AddConfigurationAction, EditStoredInput, InstallFromActivation, ListMcpServerCommand, McpBrowseCommand, McpBrowseResourcesCommand, McpConfigureSamplingModels, MCPServerActionRendering, McpServerOptionsCommand, RemoveStoredInput, ResetMcpCachedTools, ResetMcpTrustCommand, RestartServer, ShowConfiguration, ShowOutput, StartServer, StopServer } from './mcpCommands.js';
3839
import { McpDiscovery } from './mcpDiscovery.js';
3940
import { McpLanguageFeatures } from './mcpLanguageFeatures.js';
4041
import { McpResourceQuickAccess } from './mcpResourceQuickAccess.js';
@@ -49,6 +50,7 @@ registerSingleton(IMcpService, McpService, InstantiationType.Delayed);
4950
registerSingleton(IMcpWorkbenchService, McpWorkbenchService, InstantiationType.Eager);
5051
registerSingleton(IMcpConfigPathsService, McpConfigPathsService, InstantiationType.Delayed);
5152
registerSingleton(IMcpDevModeDebugging, McpDevModeDebugging, InstantiationType.Delayed);
53+
registerSingleton(IMcpSamplingService, McpSamplingService, InstantiationType.Delayed);
5254

5355
mcpDiscoveryRegistry.register(new SyncDescriptor(RemoteNativeMpcDiscovery));
5456
mcpDiscoveryRegistry.register(new SyncDescriptor(ConfigMcpDiscovery));
@@ -76,6 +78,7 @@ registerAction2(RestartServer);
7678
registerAction2(ShowConfiguration);
7779
registerAction2(McpBrowseCommand);
7880
registerAction2(McpBrowseResourcesCommand);
81+
registerAction2(McpConfigureSamplingModels);
7982

8083
registerWorkbenchContribution2('mcpActionRendering', MCPServerActionRendering, WorkbenchPhase.BlockRestore);
8184
registerWorkbenchContribution2('mcpAddContext', McpAddContextContribution, WorkbenchPhase.Eventually);

‎src/vs/workbench/contrib/mcp/browser/mcpCommands.ts

Lines changed: 49 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import { Event } from '../../../../base/common/event.js';
1111
import { Disposable, DisposableStore } from '../../../../base/common/lifecycle.js';
1212
import { autorun, derived } from '../../../../base/common/observable.js';
1313
import { ThemeIcon } from '../../../../base/common/themables.js';
14+
import { isDefined } from '../../../../base/common/types.js';
1415
import { URI } from '../../../../base/common/uri.js';
1516
import { ILocalizedString, localize, localize2 } from '../../../../nls.js';
1617
import { IActionViewItemService } from '../../../../platform/actions/browser/actionViewItemService.js';
@@ -32,14 +33,15 @@ import { IEditorService } from '../../../services/editor/common/editorService.js
3233
import { IViewsService } from '../../../services/views/common/viewsService.js';
3334
import { ChatContextKeys } from '../../chat/common/chatContextKeys.js';
3435
import { ChatMode } from '../../chat/common/constants.js';
36+
import { ILanguageModelsService } from '../../chat/common/languageModels.js';
3537
import { extensionsFilterSubMenu, IExtensionsWorkbenchService } from '../../extensions/common/extensions.js';
3638
import { TEXT_FILE_EDITOR_ID } from '../../files/common/files.js';
3739
import { McpCommandIds } from '../common/mcpCommandIds.js';
3840
import { McpContextKeys } from '../common/mcpContextKeys.js';
3941
import { IMcpRegistry } from '../common/mcpRegistryTypes.js';
40-
import { McpResourceQuickAccess } from './mcpResourceQuickAccess.js';
41-
import { IMcpServer, IMcpServerStartOpts, IMcpService, IMcpWorkbenchService, InstalledMcpServersViewId, LazyCollectionState, McpConnectionState, McpServersGalleryEnabledContext, McpServerCacheState } from '../common/mcpTypes.js';
42+
import { IMcpSamplingService, IMcpServer, IMcpServerStartOpts, IMcpService, IMcpWorkbenchService, InstalledMcpServersViewId, LazyCollectionState, McpConnectionState, McpServerCacheState, McpServersGalleryEnabledContext } from '../common/mcpTypes.js';
4243
import { McpAddConfigurationCommand } from './mcpCommandsAddConfiguration.js';
44+
import { McpResourceQuickAccess } from './mcpResourceQuickAccess.js';
4345
import { McpUrlHandler } from './mcpUrlHandler.js';
4446

4547
// acroynms do not get localized
@@ -152,17 +154,17 @@ export class McpServerOptionsCommand extends Action2 {
152154
const quickInputService = accessor.get(IQuickInputService);
153155
const mcpRegistry = accessor.get(IMcpRegistry);
154156
const editorService = accessor.get(IEditorService);
157+
const commandService = accessor.get(ICommandService);
155158
const server = mcpService.servers.get().find(s => s.definition.id === id);
156159
if (!server) {
157160
return;
158161
}
159162

160-
161163
const collection = mcpRegistry.collections.get().find(c => c.id === server.collection.id);
162164
const serverDefinition = collection?.serverDefinitions.get().find(s => s.id === server.definition.id);
163165

164166
interface ActionItem extends IQuickPickItem {
165-
action: 'start' | 'stop' | 'restart' | 'showOutput' | 'config';
167+
action: 'start' | 'stop' | 'restart' | 'showOutput' | 'config' | 'configSampling';
166168
}
167169

168170
const items: ActionItem[] = [];
@@ -188,6 +190,10 @@ export class McpServerOptionsCommand extends Action2 {
188190
items.push({
189191
label: localize('mcp.showOutput', 'Show Output'),
190192
action: 'showOutput'
193+
}, {
194+
label: localize('mcp.configAccess', 'Configure Model Access'),
195+
description: localize('mcp.showOutput.description', 'Set the models the server can use via MCP sampling'),
196+
action: 'configSampling'
191197
});
192198

193199
const configTarget = serverDefinition?.presentation?.origin || collection?.presentation?.origin;
@@ -228,6 +234,8 @@ export class McpServerOptionsCommand extends Action2 {
228234
options: { selection: URI.isUri(configTarget) ? undefined : configTarget!.range }
229235
});
230236
break;
237+
case 'configSampling':
238+
return commandService.executeCommand(McpCommandIds.ConfigureSamplingModels, server);
231239
default:
232240
assertNever(pick.action);
233241
}
@@ -620,3 +628,40 @@ export class McpBrowseResourcesCommand extends Action2 {
620628
accessor.get(IQuickInputService).quickAccess.show(McpResourceQuickAccess.PREFIX);
621629
}
622630
}
631+
632+
633+
export class McpConfigureSamplingModels extends Action2 {
634+
constructor() {
635+
super({
636+
id: McpCommandIds.ConfigureSamplingModels,
637+
title: localize2('mcp.configureSamplingModels', "Configure SamplingModel"),
638+
category,
639+
});
640+
}
641+
642+
async run(accessor: ServicesAccessor, server: IMcpServer): Promise<number> {
643+
const quickInputService = accessor.get(IQuickInputService);
644+
const lmService = accessor.get(ILanguageModelsService);
645+
const mcpSampling = accessor.get(IMcpSamplingService);
646+
647+
const existingIds = new Set(mcpSampling.getConfig(server).allowedModels);
648+
const allItems: IQuickPickItem[] = lmService.getLanguageModelIds().map(id => {
649+
const model = lmService.lookupLanguageModel(id)!;
650+
return model.isUserSelectable ? ({ label: model.name, description: model.description, id, picked: existingIds.has(id) }) : undefined;
651+
}).filter(isDefined);
652+
653+
allItems.sort((a, b) => (b.picked ? 1 : 0) - (a.picked ? 1 : 0) || a.label.localeCompare(b.label));
654+
655+
// do the quickpick selection
656+
const picked = await quickInputService.pick(allItems, {
657+
placeHolder: localize('mcp.configureSamplingModels.ph', 'Pick the models {0} can access via MCP sampling', server.definition.label),
658+
canPickMany: true,
659+
});
660+
661+
if (picked) {
662+
await mcpSampling.updateConfig(server, c => c.allowedModels = picked.map(p => p.id!));
663+
}
664+
665+
return picked?.length || 0;
666+
}
667+
}

‎src/vs/workbench/contrib/mcp/common/discovery/configMcpDiscovery.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -164,6 +164,7 @@ export class ConfigMcpDiscovery extends Disposable implements IMcpDiscovery {
164164
remoteAuthority: src.path.remoteAuthority || null,
165165
serverDefinitions: src.serverDefinitions,
166166
isTrustedByDefault: true,
167+
configTarget: src.path.target,
167168
scope: src.path.scope,
168169
});
169170
}

‎src/vs/workbench/contrib/mcp/common/discovery/extensionMcpDiscovery.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@ import { Disposable, DisposableMap } from '../../../../../base/common/lifecycle.
77
import { observableValue } from '../../../../../base/common/observable.js';
88
import { isFalsyOrWhitespace } from '../../../../../base/common/strings.js';
99
import { localize } from '../../../../../nls.js';
10+
import { ConfigurationTarget } from '../../../../../platform/configuration/common/configuration.js';
1011
import { IMcpCollectionContribution } from '../../../../../platform/extensions/common/extensions.js';
1112
import { IStorageService, StorageScope, StorageTarget } from '../../../../../platform/storage/common/storage.js';
1213
import { IExtensionService } from '../../../../services/extensions/common/extensions.js';
@@ -85,6 +86,7 @@ export class ExtensionMcpDiscovery extends Disposable implements IMcpDiscovery {
8586
remoteAuthority: null,
8687
isTrustedByDefault: true,
8788
scope: StorageScope.WORKSPACE,
89+
configTarget: ConfigurationTarget.USER,
8890
serverDefinitions: observableValue<McpServerDefinition[]>(this, serverDefs?.map(McpServerDefinition.fromSerialized) || []),
8991
lazy: {
9092
isCached: !!serverDefs,

‎src/vs/workbench/contrib/mcp/common/discovery/nativeMcpDiscoveryAbstract.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@ import { Schemas } from '../../../../../base/common/network.js';
1010
import { autorunWithStore, IObservable, IReader, ISettableObservable, observableValue } from '../../../../../base/common/observable.js';
1111
import { URI } from '../../../../../base/common/uri.js';
1212
import { localize } from '../../../../../nls.js';
13-
import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js';
13+
import { ConfigurationTarget, IConfigurationService } from '../../../../../platform/configuration/common/configuration.js';
1414
import { IFileService } from '../../../../../platform/files/common/files.js';
1515
import { IInstantiationService } from '../../../../../platform/instantiation/common/instantiation.js';
1616
import { ILabelService } from '../../../../../platform/label/common/label.js';
@@ -144,6 +144,7 @@ export abstract class NativeFilesystemMcpDiscovery extends FilesystemMcpDiscover
144144
id: adapter.id,
145145
label: discoverySourceLabel[adapter.discoverySource] + this.suffix,
146146
remoteAuthority: adapter.remoteAuthority,
147+
configTarget: ConfigurationTarget.USER,
147148
scope: StorageScope.PROFILE,
148149
isTrustedByDefault: false,
149150
serverDefinitions: observableValue<readonly McpServerDefinition[]>(this, []),

‎src/vs/workbench/contrib/mcp/common/discovery/workspaceMcpDiscoveryAdapter.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@ import { DisposableMap, IDisposable } from '../../../../../base/common/lifecycle
77
import { observableValue } from '../../../../../base/common/observable.js';
88
import { joinPath } from '../../../../../base/common/resources.js';
99
import { URI } from '../../../../../base/common/uri.js';
10-
import { IConfigurationService } from '../../../../../platform/configuration/common/configuration.js';
10+
import { ConfigurationTarget, IConfigurationService } from '../../../../../platform/configuration/common/configuration.js';
1111
import { IFileService } from '../../../../../platform/files/common/files.js';
1212
import { StorageScope } from '../../../../../platform/storage/common/storage.js';
1313
import { IWorkspaceContextService, IWorkspaceFolder } from '../../../../../platform/workspace/common/workspace.js';
@@ -56,6 +56,7 @@ export class CursorWorkspaceMcpDiscoveryAdapter extends FilesystemMcpDiscovery i
5656
scope: StorageScope.WORKSPACE,
5757
isTrustedByDefault: false,
5858
serverDefinitions: observableValue(this, []),
59+
configTarget: ConfigurationTarget.WORKSPACE_FOLDER,
5960
presentation: {
6061
origin: configFile,
6162
order: McpCollectionSortOrder.WorkspaceFolder + 1,

‎src/vs/workbench/contrib/mcp/common/mcpCommandIds.ts

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -7,19 +7,20 @@
77
* Contains all MCP command IDs used in the workbench.
88
*/
99
export const enum McpCommandIds {
10-
ListServer = 'workbench.mcp.listServer',
11-
ServerOptions = 'workbench.mcp.serverOptions',
12-
ResetTrust = 'workbench.mcp.resetTrust',
13-
ResetCachedTools = 'workbench.mcp.resetCachedTools',
1410
AddConfiguration = 'workbench.mcp.addConfiguration',
15-
RemoveStoredInput = 'workbench.mcp.removeStoredInput',
16-
EditStoredInput = 'workbench.mcp.editStoredInput',
11+
Browse = 'workbench.mcp.browseServers',
1712
BrowseResources = 'workbench.mcp.browseResources',
13+
ConfigureSamplingModels = 'workbench.mcp.configureSamplingModels',
14+
EditStoredInput = 'workbench.mcp.editStoredInput',
15+
InstallFromActivation = 'workbench.mcp.installFromActivation',
16+
ListServer = 'workbench.mcp.listServer',
17+
RemoveStoredInput = 'workbench.mcp.removeStoredInput',
18+
ResetCachedTools = 'workbench.mcp.resetCachedTools',
19+
ResetTrust = 'workbench.mcp.resetTrust',
20+
RestartServer = 'workbench.mcp.restartServer',
21+
ServerOptions = 'workbench.mcp.serverOptions',
1822
ShowConfiguration = 'workbench.mcp.showConfiguration',
1923
ShowOutput = 'workbench.mcp.showOutput',
20-
RestartServer = 'workbench.mcp.restartServer',
2124
StartServer = 'workbench.mcp.startServer',
2225
StopServer = 'workbench.mcp.stopServer',
23-
InstallFromActivation = 'workbench.mcp.installFromActivation',
24-
Browse = 'workbench.mcp.browseServers'
2526
}

‎src/vs/workbench/contrib/mcp/common/mcpConfiguration.ts

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -45,6 +45,13 @@ export const discoverySourceLabel: Record<DiscoverySource, string> = {
4545
export const mcpConfigurationSection = 'mcp';
4646
export const mcpDiscoverySection = 'chat.mcp.discovery.enabled';
4747
export const mcpEnabledSection = 'chat.mcp.enabled';
48+
export const mcpServerSamplingSection = 'chat.mcp.serverSampling';
49+
50+
export interface IMcpServerSamplingConfiguration {
51+
allowedDuringChat?: boolean;
52+
allowedOutsideChat?: boolean;
53+
allowedModels?: string[];
54+
}
4855

4956
export const mcpSchemaExampleServers = {
5057
'mcp-server-time': {

0 commit comments

Comments
 (0)