Skip to content

Commit 71b706f

Browse files
jriekenjrieken
authored
Bring back MCP tool set but keep its tools (microsoft#250273)
* Iterator#every * register tool set for MCP serer * don't show MCP toolset in tools picker (because they are a bucket already) * remove the filtering of homogenous tool sets
1 parent d1e632d commit 71b706f

File tree

5 files changed

+98
-40
lines changed

5 files changed

+98
-40
lines changed

‎src/vs/base/common/iterator.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,16 @@ export namespace Iterable {
5656
return false;
5757
}
5858

59+
export function every<T>(iterable: Iterable<T>, predicate: (t: T, i: number) => unknown): boolean {
60+
let i = 0;
61+
for (const element of iterable) {
62+
if (!predicate(element, i++)) {
63+
return false;
64+
}
65+
}
66+
return true;
67+
}
68+
5969
export function find<T, R extends T>(iterable: Iterable<T>, predicate: (t: T) => t is R): R | undefined;
6070
export function find<T>(iterable: Iterable<T>, predicate: (t: T) => boolean): T | undefined;
6171
export function find<T>(iterable: Iterable<T>, predicate: (t: T) => boolean): T | undefined {

‎src/vs/base/test/common/iterator.test.ts

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,4 +32,43 @@ suite('Iterable', function () {
3232
assert.deepStrictEqual([...Iterable.wrap(1)], [1]);
3333
assert.deepStrictEqual([...Iterable.wrap([1, 2, 3])], [1, 2, 3]);
3434
});
35+
36+
test('every', function () {
37+
// Empty iterable should return true (vacuous truth)
38+
assert.strictEqual(Iterable.every([], () => false), true);
39+
40+
// All elements match predicate
41+
assert.strictEqual(Iterable.every([2, 4, 6, 8], x => x % 2 === 0), true);
42+
assert.strictEqual(Iterable.every([1, 2, 3], x => x > 0), true);
43+
44+
// Not all elements match predicate
45+
assert.strictEqual(Iterable.every([1, 2, 3, 4], x => x % 2 === 0), false);
46+
assert.strictEqual(Iterable.every([1, 2, 3], x => x > 2), false);
47+
48+
// Single element - matches
49+
assert.strictEqual(Iterable.every([5], x => x === 5), true);
50+
51+
// Single element - doesn't match
52+
assert.strictEqual(Iterable.every([5], x => x === 6), false);
53+
54+
// Test index parameter in predicate
55+
const numbers = [10, 11, 12, 13];
56+
assert.strictEqual(Iterable.every(numbers, (x, i) => x === 10 + i), true);
57+
assert.strictEqual(Iterable.every(numbers, (x, i) => i < 2), false);
58+
59+
// Test early termination - predicate should not be called for all elements
60+
let callCount = 0;
61+
const result = Iterable.every([1, 2, 3, 4, 5], x => {
62+
callCount++;
63+
return x < 3;
64+
});
65+
assert.strictEqual(result, false);
66+
assert.strictEqual(callCount, 3); // Should stop at the third element
67+
68+
// Test with truthy/falsy values
69+
assert.strictEqual(Iterable.every([1, 2, 3], x => x), true);
70+
assert.strictEqual(Iterable.every([1, 0, 3], x => x), false);
71+
assert.strictEqual(Iterable.every(['a', 'b', 'c'], x => x), true);
72+
assert.strictEqual(Iterable.every(['a', '', 'c'], x => x), false);
73+
});
3574
});

‎src/vs/workbench/contrib/chat/browser/actions/chatToolPicker.ts

Lines changed: 30 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { assertNever } from '../../../../../base/common/assert.js';
66
import { Codicon } from '../../../../../base/common/codicons.js';
77
import { diffSets } from '../../../../../base/common/collections.js';
88
import { Event } from '../../../../../base/common/event.js';
9+
import { Iterable } from '../../../../../base/common/iterator.js';
910
import { DisposableStore } from '../../../../../base/common/lifecycle.js';
1011
import { ThemeIcon } from '../../../../../base/common/themables.js';
1112
import { assertType } from '../../../../../base/common/types.js';
@@ -187,17 +188,19 @@ export async function showToolsPicker(
187188
}
188189

189190
if (toolSetOrTool instanceof ToolSet) {
190-
bucket.children.push({
191-
parent: bucket,
192-
type: 'item',
193-
picked,
194-
toolset: toolSetOrTool,
195-
label: toolSetOrTool.toolReferenceName,
196-
description: toolSetOrTool.description,
197-
indented: true,
198-
buttons
191+
if (toolSetOrTool.source.type !== 'mcp') { // don't show the MCP toolset
192+
bucket.children.push({
193+
parent: bucket,
194+
type: 'item',
195+
picked,
196+
toolset: toolSetOrTool,
197+
label: toolSetOrTool.toolReferenceName,
198+
description: toolSetOrTool.description,
199+
indented: true,
200+
buttons
201+
});
202+
}
199203

200-
});
201204
} else if (toolSetOrTool.canBeReferencedInPrompt) {
202205
bucket.children.push({
203206
parent: bucket,
@@ -374,5 +377,22 @@ export async function showToolsPicker(
374377

375378
store.dispose();
376379

380+
const mcpToolSets = new Set<ToolSet>();
381+
382+
for (const item of toolsService.toolSets.get()) {
383+
if (item.source.type === 'mcp') {
384+
mcpToolSets.add(item);
385+
386+
if (Iterable.every(item.getTools(), tool => result.get(tool))) {
387+
// ALL tools from the MCP tool set are here, replace them with just the toolset
388+
// but only when computing the final result
389+
for (const tool of item.getTools()) {
390+
result.delete(tool);
391+
}
392+
result.set(item, true);
393+
}
394+
}
395+
}
396+
377397
return didAccept ? result : undefined;
378398
}

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

Lines changed: 3 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ import { autorun, derived, IObservable, observableFromEvent, ObservableMap, obse
88
import { ObservableMemento, observableMemento } from '../../../../platform/observable/common/observableMemento.js';
99
import { IStorageService, StorageScope, StorageTarget } from '../../../../platform/storage/common/storage.js';
1010
import { ChatMode } from '../common/constants.js';
11-
import { ILanguageModelToolsService, IToolData, ToolSet, ToolDataSource } from '../common/languageModelToolsService.js';
11+
import { ILanguageModelToolsService, IToolData, ToolSet } from '../common/languageModelToolsService.js';
1212

1313

1414
/**
@@ -93,32 +93,9 @@ export class ChatSelectedTools extends Disposable {
9393

9494
this._store.add(autorun(r => {
9595

96-
const sourceByTool = new Map<IToolData, ToolDataSource>();
97-
98-
for (const tool of this._allTools.read(r)) {
99-
if (!tool.canBeReferencedInPrompt) {
100-
continue;
101-
}
102-
sourceByTool.set(tool, tool.source);
103-
}
104-
96+
const tools = this._allTools.read(r).filter(t => t.canBeReferencedInPrompt);
10597
const toolSets = _toolsService.toolSets.read(r);
10698

107-
for (const toolSet of toolSets) {
108-
109-
if (!toolSet.isHomogenous.read(r)) {
110-
// only homogenous tool sets can swallow tools
111-
continue;
112-
}
113-
114-
for (const toolInSet of toolSet.getTools(r)) {
115-
const source = sourceByTool.get(toolInSet);
116-
if (source && ToolDataSource.equals(source, toolInSet.source)) {
117-
sourceByTool.delete(toolInSet);
118-
}
119-
}
120-
}
121-
12299
const oldItems = new Set(this.entriesMap.keys());
123100

124101
const disabledData = mode.read(r) === ChatMode.Agent
@@ -127,7 +104,7 @@ export class ChatSelectedTools extends Disposable {
127104

128105
transaction(tx => {
129106

130-
for (const tool of sourceByTool.keys()) {
107+
for (const tool of tools) {
131108
const enabled = !disabledData || !disabledData.toolIds.has(tool.id);
132109
this.entriesMap.set(tool, enabled, tx);
133110
oldItems.delete(tool);

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

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@ import { ILogService } from '../../../../platform/log/common/log.js';
1818
import { IProductService } from '../../../../platform/product/common/productService.js';
1919
import { StorageScope } from '../../../../platform/storage/common/storage.js';
2020
import { getAttachableImageExtension } from '../../chat/common/chatModel.js';
21-
import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, IToolResultInputOutputDetails, ToolProgress } from '../../chat/common/languageModelToolsService.js';
21+
import { CountTokensCallback, ILanguageModelToolsService, IPreparedToolInvocation, IToolData, IToolImpl, IToolInvocation, IToolResult, IToolResultInputOutputDetails, ToolDataSource, ToolProgress, ToolSet } from '../../chat/common/languageModelToolsService.js';
2222
import { McpCommandIds } from './mcpCommandIds.js';
2323
import { IMcpRegistry } from './mcpRegistryTypes.js';
2424
import { McpServer, McpServerMetadataCache } from './mcpServer.js';
@@ -91,7 +91,7 @@ export class McpService extends Disposable implements IMcpService {
9191
await Promise.all(todo);
9292
}
9393

94-
private _syncTools(server: McpServer, store: DisposableStore) {
94+
private _syncTools(server: McpServer, toolSet: ToolSet, source: ToolDataSource, store: DisposableStore) {
9595
const tools = new Map</* tool ID */string, ISyncedToolData>();
9696

9797
store.add(autorun(reader => {
@@ -103,14 +103,15 @@ export class McpService extends Disposable implements IMcpService {
103103
const registerTool = (tool: IMcpTool, toolData: IToolData, store: DisposableStore) => {
104104
store.add(this._toolsService.registerToolData(toolData));
105105
store.add(this._toolsService.registerToolImplementation(tool.id, this._instantiationService.createInstance(McpToolImplementation, tool, server)));
106+
store.add(toolSet.addTool(toolData));
106107
};
107108

108109
for (const tool of server.tools.read(reader)) {
109110
const existing = tools.get(tool.id);
110111
const collection = this._mcpRegistry.collections.get().find(c => c.id === server.collection.id);
111112
const toolData: IToolData = {
112113
id: tool.id,
113-
source: { type: 'mcp', label: server.definition.label, collectionId: server.collection.id, definitionId: server.definition.id },
114+
source,
114115
icon: Codicon.tools,
115116
displayName: tool.definition.annotations?.title || tool.definition.name,
116117
toolReferenceName: tool.id,
@@ -205,8 +206,19 @@ export class McpService extends Disposable implements IMcpService {
205206
def.toolPrefix,
206207
);
207208

209+
const source: ToolDataSource = { type: 'mcp', label: object.definition.label, collectionId: object.collection.id, definitionId: object.definition.id };
210+
const toolSet = this._toolsService.createToolSet(
211+
source,
212+
def.serverDefinition.id, def.serverDefinition.label,
213+
{
214+
icon: Codicon.mcp,
215+
toolReferenceName: def.serverDefinition.label,
216+
description: localize('mcp.toolset', "{0}: All Tools", def.serverDefinition.label)
217+
}
218+
);
219+
store.add(toolSet);
208220
store.add(object);
209-
this._syncTools(object, store);
221+
this._syncTools(object, toolSet, source, store);
210222

211223
nextServers.push({ object, dispose: () => store.dispose(), toolPrefix: def.toolPrefix });
212224
}

0 commit comments

Comments
 (0)