Skip to content

Commit 25d003b

Browse files
connor4312connor4312
authored
tools: surface tool source in chat additions api (#256558)
For MCP I can already kind of get this from the tool name, but for extension tools we previously had no way to know which extension provided the tool. This PR adds a `LanguageModelToolInformation.source` containing details about where the tool came from. I also use this to include the MCP server instructions which is needed for #250017. cc @aeschli @roblourens
1 parent dac120d commit 25d003b

12 files changed

+172
-49
lines changed

‎src/vs/workbench/api/browser/mainThreadLanguageModelTools.ts

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ export class MainThreadLanguageModelTools extends Disposable implements MainThre
4141
userDescription: tool.userDescription,
4242
modelDescription: tool.modelDescription,
4343
inputSchema: tool.inputSchema,
44+
source: tool.source,
4445
} satisfies IToolDataDto));
4546
}
4647

‎src/vs/workbench/api/common/extHost.api.impl.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1855,6 +1855,8 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I
18551855
LanguageModelToolResult: extHostTypes.LanguageModelToolResult,
18561856
LanguageModelToolResult2: extHostTypes.LanguageModelToolResult2,
18571857
LanguageModelDataPart: extHostTypes.LanguageModelDataPart,
1858+
LanguageModelToolExtensionSource: extHostTypes.LanguageModelToolExtensionSource,
1859+
LanguageModelToolMCPSource: extHostTypes.LanguageModelToolMCPSource,
18581860
ExtendedLanguageModelToolResult: extHostTypes.ExtendedLanguageModelToolResult,
18591861
PreparedTerminalToolInvocation: extHostTypes.PreparedTerminalToolInvocation,
18601862
LanguageModelChatToolMode: extHostTypes.LanguageModelChatToolMode,

‎src/vs/workbench/api/common/extHost.protocol.ts

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ import { IChatContentInlineReference, IChatFollowup, IChatNotebookEdit, IChatPro
6161
import { IChatRequestVariableValue } from '../../contrib/chat/common/chatVariables.js';
6262
import { ChatAgentLocation } from '../../contrib/chat/common/constants.js';
6363
import { IChatMessage, IChatResponseFragment, ILanguageModelChatMetadata, ILanguageModelChatSelector, ILanguageModelsChangeEvent } from '../../contrib/chat/common/languageModels.js';
64-
import { IPreparedToolInvocation, IToolInvocation, IToolInvocationPreparationContext, IToolProgressStep, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
64+
import { IPreparedToolInvocation, IToolInvocation, IToolInvocationPreparationContext, IToolProgressStep, IToolResult, ToolDataSource } from '../../contrib/chat/common/languageModelToolsService.js';
6565
import { DebugConfigurationProviderTriggerKind, IAdapterDescriptor, IConfig, IDebugSessionReplMode, IDebugTestRunReference, IDebugVisualization, IDebugVisualizationContext, IDebugVisualizationTreeItem, MainThreadDebugVisualization } from '../../contrib/debug/common/debug.js';
6666
import { McpCollectionDefinition, McpConnectionState, McpServerDefinition, McpServerLaunch } from '../../contrib/mcp/common/mcpTypes.js';
6767
import * as notebookCommon from '../../contrib/notebook/common/notebookCommon.js';
@@ -1391,6 +1391,7 @@ export interface IToolDataDto {
13911391
displayName: string;
13921392
userDescription?: string;
13931393
modelDescription: string;
1394+
source: Dto<ToolDataSource>;
13941395
inputSchema?: IJSONSchema;
13951396
}
13961397

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

Lines changed: 31 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,34 @@ import { ExtHostLanguageModelToolsShape, IMainContext, IToolDataDto, MainContext
2020
import { ExtHostLanguageModels } from './extHostLanguageModels.js';
2121
import * as typeConvert from './extHostTypeConverters.js';
2222
import { SearchExtensionsToolId } from '../../contrib/extensions/common/searchExtensionsTool.js';
23+
import { Lazy } from '../../../base/common/lazy.js';
2324

2425
class Tool {
2526

2627
private _data: IToolDataDto;
27-
private _apiObject: vscode.LanguageModelToolInformation | undefined;
28+
private _apiObject = new Lazy<vscode.LanguageModelToolInformation>(() => {
29+
const that = this;
30+
return Object.freeze({
31+
get name() { return that._data.id; },
32+
get description() { return that._data.modelDescription; },
33+
get inputSchema() { return that._data.inputSchema; },
34+
get tags() { return that._data.tags ?? []; },
35+
get source() { return undefined; }
36+
});
37+
});
38+
39+
private _apiObjectWithChatParticipantAdditions = new Lazy<vscode.LanguageModelToolInformation>(() => {
40+
const that = this;
41+
const source = typeConvert.LanguageModelToolSource.to(that._data.source);
42+
43+
return Object.freeze({
44+
get name() { return that._data.id; },
45+
get description() { return that._data.modelDescription; },
46+
get inputSchema() { return that._data.inputSchema; },
47+
get tags() { return that._data.tags ?? []; },
48+
get source() { return source; }
49+
});
50+
});
2851

2952
constructor(data: IToolDataDto) {
3053
this._data = data;
@@ -39,16 +62,11 @@ class Tool {
3962
}
4063

4164
get apiObject(): vscode.LanguageModelToolInformation {
42-
if (!this._apiObject) {
43-
const that = this;
44-
this._apiObject = Object.freeze({
45-
get name() { return that._data.id; },
46-
get description() { return that._data.modelDescription; },
47-
get inputSchema() { return that._data.inputSchema; },
48-
get tags() { return that._data.tags ?? []; },
49-
});
50-
}
51-
return this._apiObject;
65+
return this._apiObject.value;
66+
}
67+
68+
get apiObjectWithChatParticipantAdditions() {
69+
return this._apiObjectWithChatParticipantAdditions.value;
5270
}
5371
}
5472

@@ -136,8 +154,9 @@ export class ExtHostLanguageModelTools implements ExtHostLanguageModelToolsShape
136154
}
137155

138156
getTools(extension: IExtensionDescription): vscode.LanguageModelToolInformation[] {
157+
const hasParticipantAdditions = isProposedApiEnabled(extension, 'chatParticipantPrivate');
139158
return Array.from(this._allTools.values())
140-
.map(tool => tool.apiObject)
159+
.map(tool => hasParticipantAdditions ? tool.apiObjectWithChatParticipantAdditions : tool.apiObject)
141160
.filter(tool => {
142161
switch (tool.name) {
143162
case InternalEditToolId:

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

Lines changed: 10 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ import { IChatAgentRequest, IChatAgentResult } from '../../contrib/chat/common/c
4242
import { IChatRequestDraft } from '../../contrib/chat/common/chatEditingService.js';
4343
import { IChatRequestVariableEntry, isImageVariableEntry } from '../../contrib/chat/common/chatVariableEntries.js';
4444
import { IChatAgentMarkdownContentWithVulnerability, IChatCodeCitation, IChatCommandButton, IChatConfirmation, IChatContentInlineReference, IChatContentReference, IChatExtensionsContent, IChatFollowup, IChatMarkdownContent, IChatMoveMessage, IChatPrepareToolInvocationPart, IChatProgressMessage, IChatResponseCodeblockUriPart, IChatTaskDto, IChatTaskResult, IChatTextEdit, IChatTreeData, IChatUserActionEvent, IChatWarningMessage } from '../../contrib/chat/common/chatService.js';
45-
import { IToolData, IToolResult } from '../../contrib/chat/common/languageModelToolsService.js';
45+
import { IToolResult, ToolDataSource } from '../../contrib/chat/common/languageModelToolsService.js';
4646
import * as chatProvider from '../../contrib/chat/common/languageModels.js';
4747
import { IChatMessageDataPart, IChatResponseDataPart, IChatResponsePromptTsxPart, IChatResponseTextPart } from '../../contrib/chat/common/languageModels.js';
4848
import { DebugTreeItemCollapsibleState, IDebugVisualizationTreeItem } from '../../contrib/debug/common/debug.js';
@@ -3279,15 +3279,15 @@ export namespace DebugTreeItem {
32793279
}
32803280
}
32813281

3282-
export namespace LanguageModelToolDescription {
3283-
export function to(item: IToolData): vscode.LanguageModelToolInformation {
3284-
return {
3285-
// Note- the reason this is a unique 'name' is just to avoid confusion with the toolCallId
3286-
name: item.id,
3287-
description: item.modelDescription,
3288-
inputSchema: item.inputSchema,
3289-
tags: item.tags ?? [],
3290-
};
3282+
export namespace LanguageModelToolSource {
3283+
export function to(source: Dto<ToolDataSource>): vscode.LanguageModelToolInformation['source'] {
3284+
if (source.type === 'mcp') {
3285+
return new types.LanguageModelToolMCPSource(source.label, source.serverLabel || source.label, source.instructions);
3286+
} else if (source.type === 'extension') {
3287+
return new types.LanguageModelToolExtensionSource(source.extensionId.value, source.label);
3288+
} else {
3289+
return undefined;
3290+
}
32913291
}
32923292
}
32933293

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

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5133,6 +5133,14 @@ export enum LanguageModelChatToolMode {
51335133
Required = 2
51345134
}
51355135

5136+
export class LanguageModelToolExtensionSource implements vscode.LanguageModelToolExtensionSource {
5137+
constructor(public readonly id: string, public readonly label: string) { }
5138+
}
5139+
5140+
export class LanguageModelToolMCPSource implements vscode.LanguageModelToolMCPSource {
5141+
constructor(public readonly label: string, public readonly name: string, public readonly instructions: string | undefined) { }
5142+
}
5143+
51365144
//#endregion
51375145

51385146
//#region ai
@@ -5175,7 +5183,7 @@ export enum KeywordRecognitionStatus {
51755183

51765184
//#endregion
51775185

5178-
//#region MC
5186+
//#region MCP
51795187
export class McpStdioServerDefinition implements vscode.McpStdioServerDefinition {
51805188
cwd?: URI;
51815189

‎src/vs/workbench/contrib/chat/common/languageModelToolsService.ts

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,8 @@ export type ToolDataSource =
6060
| {
6161
type: 'mcp';
6262
label: string;
63+
serverLabel: string | undefined;
64+
instructions: string | undefined;
6365
collectionId: string;
6466
definitionId: string;
6567
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ suite('ChatSelectedTools', () => {
4949

5050
ensureNoDisposablesAreLeakedInTestSuite();
5151

52-
const mcpSource: ToolDataSource = { type: 'mcp', label: 'MCP', collectionId: '', definitionId: '' };
52+
const mcpSource: ToolDataSource = { type: 'mcp', label: 'MCP', collectionId: '', definitionId: '', instructions: '', serverLabel: '' };
5353
test('Can\'t enable/disable MCP tools directly #18161', () => {
5454

5555
return runWithFakedTimers({}, async () => {

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

Lines changed: 46 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,9 @@ type ServerBootStateClassification = {
7979
};
8080

8181
interface IToolCacheEntry {
82+
readonly serverName: string | undefined;
83+
readonly serverInstructions: string | undefined;
84+
8285
readonly nonce: string | undefined;
8386
/** Cached tools so we can show what's available before it's started */
8487
readonly tools: readonly IValidatedMcpTool[];
@@ -173,29 +176,35 @@ interface IValidatedMcpTool extends MCP.Tool {
173176
serverToolName: string;
174177
}
175178

179+
interface ServerMetadata {
180+
readonly serverName: string | undefined;
181+
readonly serverInstructions: string | undefined;
182+
}
183+
176184
class CachedPrimitive<T, C> {
177185
constructor(
178186
private readonly _definitionId: string,
179187
private readonly _cache: McpServerMetadataCache,
180-
private readonly _fromCache: (entry: IToolCacheEntry) => readonly C[],
181-
private readonly _toT: (values: readonly C[], reader: IDerivedReader<void>) => T[],
188+
private readonly _fromCache: (entry: IToolCacheEntry) => C,
189+
private readonly _toT: (values: C, reader: IDerivedReader<void>) => T,
190+
private readonly defaultValue: C,
182191
) { }
183192

184-
public get fromCache(): { nonce: string | undefined; data: readonly C[] } | undefined {
193+
public get fromCache(): { nonce: string | undefined; data: C } | undefined {
185194
const c = this._cache.get(this._definitionId);
186195
return c ? { data: this._fromCache(c), nonce: c.nonce } : undefined;
187196
}
188197

189198
public readonly fromServerPromise = observableValue<ObservablePromise<{
190-
readonly data: C[];
199+
readonly data: C;
191200
readonly nonce: string | undefined;
192201
}> | undefined>(this, undefined);
193202

194203
private readonly fromServer = derived(reader => this.fromServerPromise.read(reader)?.promiseResult.read(reader)?.data);
195204

196-
public readonly value: IObservable<readonly T[]> = derived(reader => {
205+
public readonly value: IObservable<T> = derived(reader => {
197206
const serverTools = this.fromServer.read(reader);
198-
const definitions = serverTools?.data ?? this.fromCache?.data ?? [];
207+
const definitions = serverTools?.data ?? this.fromCache?.data ?? this.defaultValue;
199208
return this._toT(definitions, reader);
200209
});
201210
}
@@ -254,16 +263,21 @@ export class McpServer extends Disposable implements IMcpServer {
254263
return this._capabilities;
255264
}
256265

257-
private readonly _tools: CachedPrimitive<IMcpTool, IValidatedMcpTool>;
266+
private readonly _tools: CachedPrimitive<readonly IMcpTool[], readonly IValidatedMcpTool[]>;
258267
public get tools() {
259268
return this._tools.value;
260269
}
261270

262-
private readonly _prompts: CachedPrimitive<IMcpPrompt, MCP.Prompt>;
271+
private readonly _prompts: CachedPrimitive<readonly IMcpPrompt[], readonly MCP.Prompt[]>;
263272
public get prompts() {
264273
return this._prompts.value;
265274
}
266275

276+
private readonly _serverMetadata: CachedPrimitive<ServerMetadata, ServerMetadata | undefined>;
277+
public get serverMetadata() {
278+
return this._serverMetadata.value;
279+
}
280+
267281
private readonly _fullDefinitions: IObservable<{
268282
server: McpServerDefinition | undefined;
269283
collection: McpCollectionDefinition | undefined;
@@ -388,19 +402,29 @@ export class McpServer extends Disposable implements IMcpServer {
388402
}));
389403

390404
// 3. Publish tools
391-
this._tools = new CachedPrimitive<IMcpTool, IValidatedMcpTool>(
405+
this._tools = new CachedPrimitive<readonly IMcpTool[], readonly IValidatedMcpTool[]>(
392406
this.definition.id,
393407
this._primitiveCache,
394408
(entry) => entry.tools,
395409
(entry) => entry.map(def => new McpTool(this, toolPrefix, def)).sort((a, b) => a.compare(b)),
410+
[],
396411
);
397412

398413
// 4. Publish promtps
399-
this._prompts = new CachedPrimitive<IMcpPrompt, MCP.Prompt>(
414+
this._prompts = new CachedPrimitive<readonly IMcpPrompt[], readonly MCP.Prompt[]>(
400415
this.definition.id,
401416
this._primitiveCache,
402417
(entry) => entry.prompts || [],
403418
(entry) => entry.map(e => new McpPrompt(this, e)),
419+
[],
420+
);
421+
422+
this._serverMetadata = new CachedPrimitive<ServerMetadata, ServerMetadata | undefined>(
423+
this.definition.id,
424+
this._primitiveCache,
425+
(entry) => ({ serverName: entry.serverName, serverInstructions: entry.serverInstructions }),
426+
(entry) => ({ serverName: entry?.serverName, serverInstructions: entry?.serverInstructions }),
427+
undefined,
404428
);
405429

406430
this._capabilities.set(this._primitiveCache.get(this.definition.id)?.capabilities, undefined);
@@ -664,13 +688,25 @@ export class McpServer extends Disposable implements IMcpServer {
664688
updatePrompts(undefined);
665689
}));
666690

691+
const metadataPromise = new ObservablePromise(Promise.resolve({
692+
nonce: cacheNonce,
693+
data: {
694+
serverName: handler.serverInfo.title || handler.serverInfo.name,
695+
serverInstructions: handler.serverInstructions,
696+
},
697+
}));
698+
667699
transaction(tx => {
668700
// note: all update* methods must use tx synchronously
669701
const capabilities = encodeCapabilities(handler.capabilities);
670702
this._capabilities.set(capabilities, tx);
671703

704+
this._serverMetadata.fromServerPromise.set(metadataPromise, tx);
705+
672706
Promise.all([updateTools(tx), updatePrompts(tx)]).then(([{ data: tools }, { data: prompts }]) => {
673707
this._primitiveCache.store(this.definition.id, {
708+
serverName: handler.serverInfo.title || handler.serverInfo.name,
709+
serverInstructions: handler.serverInstructions,
674710
nonce: cacheNonce,
675711
tools,
676712
prompts,

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,10 @@ export class McpServerRequestHandler extends Disposable {
7272
return this._serverInit.serverInfo;
7373
}
7474

75+
public get serverInstructions(): string | undefined {
76+
return this._serverInit.instructions;
77+
}
78+
7579
// Event emitters for server notifications
7680
private readonly _onDidReceiveCancelledNotification = this._register(new Emitter<MCP.CancelledNotification>());
7781
readonly onDidReceiveCancelledNotification = this._onDidReceiveCancelledNotification.event;

0 commit comments

Comments
 (0)