Skip to content

Commit a39f1be

Browse files
aeschliaeschli
committed
clean up tool picker
1 parent dd2c32a commit a39f1be

File tree

2 files changed

+66
-102
lines changed

2 files changed

+66
-102
lines changed

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

Lines changed: 1 addition & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,6 @@ import { ChatContextKeys } from '../../common/chatContextKeys.js';
2323
import { IChatToolInvocation } from '../../common/chatService.js';
2424
import { isResponseVM } from '../../common/chatViewModel.js';
2525
import { ChatModeKind } from '../../common/constants.js';
26-
import { IToolData, ToolSet } from '../../common/languageModelToolsService.js';
2726
import { IChatWidget, IChatWidgetService } from '../chat.js';
2827
import { ToolsScope } from '../chatSelectedTools.js';
2928
import { CHAT_CATEGORY } from './chatActions.js';
@@ -137,19 +136,7 @@ class ConfigureToolsAction extends Action2 {
137136
break;
138137
}
139138

140-
const result = await instaService.invokeFunction(showToolsPicker, placeholder, description, entriesMap.get(), newEntriesMap => {
141-
const disableToolSets: ToolSet[] = [];
142-
const disableTools: IToolData[] = [];
143-
for (const [item, enabled] of newEntriesMap) {
144-
if (!enabled) {
145-
if (item instanceof ToolSet) {
146-
disableToolSets.push(item);
147-
} else {
148-
disableTools.push(item);
149-
}
150-
}
151-
}
152-
});
139+
const result = await instaService.invokeFunction(showToolsPicker, placeholder, description, entriesMap.get());
153140
if (result) {
154141
widget.input.selectedToolsModel.set(result, false);
155142
}

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

Lines changed: 65 additions & 88 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@ import { Codicon } from '../../../../../base/common/codicons.js';
77
import { Event } from '../../../../../base/common/event.js';
88
import { DisposableStore } from '../../../../../base/common/lifecycle.js';
99
import { ThemeIcon } from '../../../../../base/common/themables.js';
10-
import { assertType } from '../../../../../base/common/types.js';
1110
import { localize } from '../../../../../nls.js';
1211
import { CommandsRegistry, ICommandService } from '../../../../../platform/commands/common/commands.js';
1312
import { ExtensionIdentifier } from '../../../../../platform/extensions/common/extensions.js';
@@ -55,6 +54,8 @@ interface IBucketTreeItem extends IToolTreeItem {
5554
readonly ordinal: BucketOrdinal;
5655
toolset?: ToolSet; // For MCP servers where the bucket represents the ToolSet - mutable
5756
readonly status?: string;
57+
readonly children: AnyTreeItem[];
58+
checked: boolean | 'partial';
5859
}
5960

6061
/**
@@ -64,6 +65,8 @@ interface IBucketTreeItem extends IToolTreeItem {
6465
interface IToolSetTreeItem extends IToolTreeItem {
6566
readonly itemType: 'toolset';
6667
readonly toolset: ToolSet;
68+
readonly children: AnyTreeItem[];
69+
checked: boolean | 'partial';
6770
}
6871

6972
/**
@@ -73,6 +76,7 @@ interface IToolSetTreeItem extends IToolTreeItem {
7376
interface IToolTreeItemData extends IToolTreeItem {
7477
readonly itemType: 'tool';
7578
readonly tool: IToolData;
79+
checked: boolean;
7680
}
7781

7882
/**
@@ -158,8 +162,7 @@ export async function showToolsPicker(
158162
accessor: ServicesAccessor,
159163
placeHolder: string,
160164
description?: string,
161-
toolsEntries?: ReadonlyMap<ToolSet | IToolData, boolean>,
162-
onUpdate?: (toolsEntries: ReadonlyMap<ToolSet | IToolData, boolean>) => void
165+
toolsEntries?: ReadonlyMap<ToolSet | IToolData, boolean>
163166
): Promise<ReadonlyMap<ToolSet | IToolData, boolean> | undefined> {
164167

165168
const quickPickService = accessor.get(IQuickInputService);
@@ -200,18 +203,18 @@ export async function showToolsPicker(
200203
// Process entries and organize into buckets
201204
for (const [toolSetOrTool, picked] of toolsEntries) {
202205
let bucketItem: IBucketTreeItem | undefined;
203-
const buttons: ActionableButton[] = [];
204206

205207
if (toolSetOrTool.source.type === 'mcp') {
206208
const key = ToolDataSource.toKey(toolSetOrTool.source);
207-
const { definitionId } = toolSetOrTool.source;
208-
const mcpServer = mcpService.servers.get().find(candidate => candidate.definition.id === definitionId);
209-
if (!mcpServer) {
210-
continue;
211-
}
212-
213209
bucketItem = bucketMap.get(key);
214210
if (!bucketItem) {
211+
const { definitionId } = toolSetOrTool.source;
212+
const mcpServer = mcpService.servers.get().find(candidate => candidate.definition.id === definitionId);
213+
if (!mcpServer) {
214+
continue;
215+
}
216+
217+
const buttons: ActionableButton[] = [];
215218
const collection = mcpRegistry.collections.get().find(c => c.id === mcpServer.collection.id);
216219
if (collection?.source) {
217220
buttons.push({
@@ -257,7 +260,7 @@ export async function showToolsPicker(
257260
} else if (toolSetOrTool.canBeReferencedInPrompt) {
258261
// Add MCP tools directly as children
259262
const toolTreeItem = createToolTreeItemFromData(toolSetOrTool, picked);
260-
bucketItem.children = [...(bucketItem.children || []), toolTreeItem];
263+
bucketItem.children.push(toolTreeItem);
261264
}
262265

263266
} else {
@@ -282,14 +285,6 @@ export async function showToolsPicker(
282285
label = localize('userBucket', "User Defined Tool Sets");
283286
// Group all user tools under one bucket
284287
key = ordinal.toString();
285-
buttons.push({
286-
iconClass: ThemeIcon.asClassName(Codicon.edit),
287-
tooltip: localize('editUserBucket', "Edit Tool Set"),
288-
action: () => {
289-
assertType(toolSetOrTool.source.type === 'user');
290-
editorService.openEditor({ resource: toolSetOrTool.source.file });
291-
}
292-
});
293288
} else {
294289
assertNever(toolSetOrTool.source);
295290
}
@@ -307,7 +302,7 @@ export async function showToolsPicker(
307302
label,
308303
checked: false,
309304
children: [],
310-
buttons,
305+
buttons: [],
311306
collapsed,
312307
...iconProps
313308
};
@@ -317,6 +312,15 @@ export async function showToolsPicker(
317312
if (toolSetOrTool instanceof ToolSet) {
318313
// Add ToolSet as child with its tools as grandchildren - create directly instead of using legacy pick structure
319314
const iconProps = mapIconToTreeItem(toolSetOrTool.icon);
315+
const buttons = [];
316+
if (toolSetOrTool.source.type === 'user') {
317+
const resource = toolSetOrTool.source.file;
318+
buttons.push({
319+
iconClass: ThemeIcon.asClassName(Codicon.edit),
320+
tooltip: localize('editUserBucket', "Edit Tool Set"),
321+
action: () => editorService.openEditor({ resource })
322+
});
323+
}
320324
const toolSetTreeItem: IToolSetTreeItem = {
321325
itemType: 'toolset',
322326
toolset: toolSetOrTool,
@@ -325,16 +329,17 @@ export async function showToolsPicker(
325329
label: toolSetOrTool.referenceName,
326330
description: toolSetOrTool.description,
327331
checked: picked,
332+
children: [],
328333
collapsed: true,
329334
// TODO: Bring this back when tools in toolsets can be enabled/disabled.
330335
// children: Array.from(toolSetOrTool.getTools()).map(tool => createToolTreeItemFromData(tool, picked)),
331336
...iconProps
332337
};
333-
bucketItem.children = [...(bucketItem.children || []), toolSetTreeItem];
338+
bucketItem.children.push(toolSetTreeItem);
334339
} else if (toolSetOrTool.canBeReferencedInPrompt) {
335340
// Add individual tool as child
336341
const toolTreeItem = createToolTreeItemFromData(toolSetOrTool, picked);
337-
bucketItem.children = [...(bucketItem.children || []), toolTreeItem];
342+
bucketItem.children.push(toolTreeItem);
338343
}
339344
}
340345
}
@@ -344,15 +349,10 @@ export async function showToolsPicker(
344349
treeItems.push(...sortedBuckets);
345350

346351
// Set up checkbox states based on parent-child relationships
347-
for (const bucketItem of treeItems.filter(isBucketTreeItem)) {
348-
if (bucketItem.checked) {
352+
for (const bucketItem of sortedBuckets) {
353+
if (bucketItem.checked === true) { // only set for MCP tool sets
349354
// Check all children if bucket is checked
350-
bucketItem.children?.forEach(child => {
351-
(child as any).checked = true;
352-
});
353-
} else {
354-
// Check bucket if any child is checked
355-
bucketItem.checked = bucketItem.children?.some(child => (child as any).checked) || false;
355+
bucketItem.children.forEach(child => child.checked = true);
356356
}
357357
}
358358

@@ -380,59 +380,54 @@ export async function showToolsPicker(
380380
}
381381
}));
382382

383-
// Result collection
384-
const result = new Map<IToolData | ToolSet, boolean>();
383+
const updateToolLimitMessage = () => {
384+
if (toolLimit) {
385+
let count = 0;
386+
const traverse = (items: readonly AnyTreeItem[]) => {
387+
for (const item of items) {
388+
if (isBucketTreeItem(item) || isToolSetTreeItem(item)) {
389+
traverse(item.children);
390+
} else if (isToolTreeItem(item) && item.checked) {
391+
count++;
392+
}
393+
}
394+
};
395+
traverse(treeItems);
396+
if (count > toolLimit) {
397+
treePicker.severity = Severity.Warning;
398+
treePicker.validationMessage = localize('toolLimitExceeded', "{0} tools are enabled. You may experience degraded tool calling above {1} tools.", count, markdownCommandLink({ title: String(toolLimit), id: '_chat.toolPicker.closeAndOpenVirtualThreshold' }));
399+
} else {
400+
treePicker.severity = Severity.Ignore;
401+
treePicker.validationMessage = undefined;
402+
}
403+
}
404+
};
405+
updateToolLimitMessage();
385406

386407
const collectResults = () => {
387-
result.clear();
388408

389-
let count = 0;
409+
const result = new Map<IToolData | ToolSet, boolean>();
390410
const traverse = (items: readonly AnyTreeItem[]) => {
391411
for (const item of items) {
392412
if (isBucketTreeItem(item)) {
393-
if (item.toolset) {
394-
// MCP server bucket represents a ToolSet
395-
const checked = typeof item.checked === 'boolean' ? item.checked : false;
396-
result.set(item.toolset, checked);
397-
}
398-
if (item.children) {
399-
traverse(item.children as readonly AnyTreeItem[]);
413+
if (item.toolset) { // MCP server
414+
// MCP toolset is enabled only if all tools are enabled
415+
const allChecked = item.checked === true;
416+
result.set(item.toolset, allChecked);
400417
}
418+
traverse(item.children);
401419
} else if (isToolSetTreeItem(item)) {
402-
const checked = typeof item.checked === 'boolean' ? item.checked : false;
403-
result.set(item.toolset, checked);
404-
if (item.children) {
405-
traverse(item.children as readonly AnyTreeItem[]);
406-
}
420+
result.set(item.toolset, item.checked === true);
421+
traverse(item.children);
407422
} else if (isToolTreeItem(item)) {
408-
const checked = typeof item.checked === 'boolean' ? item.checked : false;
409-
if (checked) { count++; }
410-
result.set(item.tool, checked);
423+
result.set(item.tool, item.checked);
411424
}
412425
}
413426
};
414427

415428
traverse(treeItems);
416-
417-
if (toolLimit) {
418-
if (count > toolLimit) {
419-
treePicker.severity = Severity.Warning;
420-
treePicker.validationMessage = localize('toolLimitExceeded', "{0} tools are enabled. You may experience degraded tool calling above {1} tools.", count, markdownCommandLink({ title: String(toolLimit), id: '_chat.toolPicker.closeAndOpenVirtualThreshold' }));
421-
} else {
422-
treePicker.severity = Severity.Ignore;
423-
treePicker.validationMessage = undefined;
424-
}
425-
}
426-
427-
// Special MCP handling: MCP toolset is enabled only if all tools are enabled
428-
for (const item of toolsService.toolSets.get()) {
429-
if (item.source.type === 'mcp') {
430-
const toolsInSet = Array.from(item.getTools());
431-
result.set(item, toolsInSet.every(tool => result.get(tool)));
432-
}
433-
}
429+
return result;
434430
};
435-
collectResults();
436431

437432
// Temporary command to close the picker and open settings, for use in the validation message
438433
store.add(CommandsRegistry.registerCommand({
@@ -444,24 +439,7 @@ export async function showToolsPicker(
444439
}));
445440

446441
// Handle checkbox state changes
447-
store.add(treePicker.onDidChangeCheckedLeafItems(() => {
448-
collectResults();
449-
450-
if (onUpdate) {
451-
// Check if results changed
452-
let didChange = toolsEntries.size !== result.size;
453-
for (const [key, value] of toolsEntries) {
454-
if (didChange) {
455-
break;
456-
}
457-
didChange = result.get(key) !== value;
458-
}
459-
460-
if (didChange) {
461-
onUpdate(result);
462-
}
463-
}
464-
}));
442+
store.add(treePicker.onDidChangeCheckedLeafItems(() => updateToolLimitMessage()));
465443

466444
// Handle acceptance
467445
let didAccept = false;
@@ -507,6 +485,5 @@ export async function showToolsPicker(
507485

508486
store.dispose();
509487

510-
collectResults();
511-
return didAccept ? result : undefined;
488+
return didAccept ? collectResults() : undefined;
512489
}

0 commit comments

Comments
 (0)