mcp: rework how virtual tool embeddings get expanded (#1164)

Our virtual-tool-group-ranking approach to expansion wasn't working as
well as we expected it to. Previously we auto-expanded groups containing
relevant tools to reach the budget. In this PR we instead put the top
10 most relevant tools at the top level of the tools list and leave the
groups unchanged.

I also added code to recalculat the embeddings when summarization happens.

Author: connor4312

src/extension/tools/common/virtualTools/toolGrouping.ts

@@ -7,13 +7,13 @@ import type { LanguageModelToolInformation } from 'vscode';
 import { ConfigKey, HARD_TOOL_LIMIT, IConfigurationService } from '../../../../platform/configuration/common/configurationService';
 import { IExperimentationService } from '../../../../platform/telemetry/common/nullExperimentationService';
 import { ITelemetryService } from '../../../../platform/telemetry/common/telemetry';
-import { equals as arraysEqual } from '../../../../util/vs/base/common/arrays';
+import { equals as arraysEqual, uniqueFilter } from '../../../../util/vs/base/common/arrays';
 import { CancellationToken } from '../../../../util/vs/base/common/cancellation';
 import { Iterable } from '../../../../util/vs/base/common/iterator';
 import { IObservable } from '../../../../util/vs/base/common/observableInternal';
 import { IInstantiationService } from '../../../../util/vs/platform/instantiation/common/instantiation';
 import { LanguageModelTextPart, LanguageModelToolResult } from '../../../../vscodeTypes';
-import { VIRTUAL_TOOL_NAME_PREFIX, VirtualTool } from './virtualTool';
+import { EMBEDDINGS_GROUP_NAME, VIRTUAL_TOOL_NAME_PREFIX, VirtualTool } from './virtualTool';
 import { VirtualToolGrouper } from './virtualToolGrouper';
 import * as Constant from './virtualToolsConstants';
 import { IToolCategorization, IToolGrouping } from './virtualToolTypes';
@@ -27,7 +27,7 @@ export function computeToolGroupingMinThreshold(experimentationService: IExperim
 
 export class ToolGrouping implements IToolGrouping {
 
-	private readonly _root = new VirtualTool(VIRTUAL_TOOL_NAME_PREFIX, '', Infinity, { groups: [], toolsetKey: '', preExpanded: true });
+	private readonly _root = new VirtualTool(VIRTUAL_TOOL_NAME_PREFIX, '', Infinity, { groups: [], toolsetKey: '', wasExpandedByDefault: true });
 	protected _grouper: IToolCategorization = this._instantiationService.createInstance(VirtualToolGrouper);
 	private _didToolsChange = true;
 	private _turnNo = 0;
@@ -80,7 +80,9 @@ export class ToolGrouping implements IToolGrouping {
 					"isVirtual": { "classification": "SystemMetaData", "purpose": "PerformanceAndHealth", "comment": "Whether this called a virtual tool", "isMeasurement": true },
 					"turnNo": { "classification": "SystemMetaData", "purpose": "PerformanceAndHealth", "comment": "Number of turns into the loop when this expansion was made", "isMeasurement": true },
 					"depth": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Nesting depth of the tool", "isMeasurement": true },
-					"preExpanded": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Whether the tool was pre-expanded or expanded on demand", "isMeasurement": true }
+					"preExpanded": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Whether the tool was pre-expanded or expanded on demand", "isMeasurement": true },
+					"wasEmbedding": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Whether the tool was pre-expanded due to an embedding", "isMeasurement": true },
+					"totalTools": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Total number of tools available when this tool was called", "isMeasurement": true }
 				}
 			*/
 			this._telemetryService.sendMSFTTelemetryEvent('virtualTools.called', {
@@ -89,7 +91,9 @@ export class ToolGrouping implements IToolGrouping {
 				turnNo: localTurnNumber,
 				isVirtual: tool instanceof VirtualTool ? 1 : 0,
 				depth: path.length - 1,
-				preExpanded: path.every(p => p.metadata.preExpanded) ? 1 : 0,
+				preExpanded: path.every(p => p.metadata.wasExpandedByDefault) ? 1 : 0,
+				wasEmbedding: path.some(p => p.name === EMBEDDINGS_GROUP_NAME) ? 1 : 0,
+				totalTools: this._tools.length,
 			});
 		}
 
@@ -124,7 +128,7 @@ export class ToolGrouping implements IToolGrouping {
 
 	async compute(query: string, token: CancellationToken): Promise<LanguageModelToolInformation[]> {
 		await this._doCompute(query, token);
-		return [...this._root.tools()];
+		return [...this._root.tools()].filter(uniqueFilter(t => t.name));
 	}
 
 	async computeAll(query: string, token: CancellationToken): Promise<(LanguageModelToolInformation | VirtualTool)[]> {
@@ -151,6 +155,7 @@ export class ToolGrouping implements IToolGrouping {
 		let trimDownTo = HARD_TOOL_LIMIT;
 
 		if (this._trimOnNextCompute) {
+			await this._grouper.recomputeEmbeddingRankings(query, this._root, token);
 			trimDownTo = Constant.TRIM_THRESHOLD;
 			this._trimOnNextCompute = false;
 		}
@@ -159,12 +164,16 @@ export class ToolGrouping implements IToolGrouping {
 
 		while (Iterable.length(this._root.tools()) > trimDownTo) {
 			const lowest = this._root.getLowestExpandedTool();
-			if (!lowest || lowest === this._root) {
+			if (!lowest || !isFinite(lowest.lastUsedOnTurn)) {
 				break; // No more tools to trim.
 			}
+			if (lowest.metadata.canBeCollapsed === false) {
+				lowest.lastUsedOnTurn = Infinity;
+				continue;
+			}
 
 			lowest.isExpanded = false;
-			lowest.metadata.preExpanded = false;
+			lowest.metadata.wasExpandedByDefault = false;
 		}
 		this._trimOnNextCompute = false;
 	}

src/extension/tools/common/virtualTools/virtualTool.ts

@@ -7,12 +7,14 @@ import type { LanguageModelToolInformation } from 'vscode';
 import { ISummarizedToolCategory } from './virtualToolTypes';
 
 export const VIRTUAL_TOOL_NAME_PREFIX = 'activate_';
+export const EMBEDDINGS_GROUP_NAME = VIRTUAL_TOOL_NAME_PREFIX + 'embeddings';
 
 export interface IVirtualToolMetadata {
 	toolsetKey: string;
 	possiblePrefix?: string;
 	groups: ISummarizedToolCategory[];
-	preExpanded?: boolean;
+	wasExpandedByDefault?: boolean;
+	canBeCollapsed?: boolean;
 }
 
 export class VirtualTool {

src/extension/tools/common/virtualTools/virtualToolGrouper.ts

@@ -18,7 +18,7 @@ import { StopWatch } from '../../../../util/vs/base/common/stopwatch';
 import { IInstantiationService } from '../../../../util/vs/platform/instantiation/common/instantiation';
 import { LanguageModelToolExtensionSource, LanguageModelToolMCPSource } from '../../../../vscodeTypes';
 import { EMBEDDING_TYPE_FOR_TOOL_GROUPING, ToolEmbeddingsComputer } from './toolEmbeddingsCache';
-import { VIRTUAL_TOOL_NAME_PREFIX, VirtualTool } from './virtualTool';
+import { EMBEDDINGS_GROUP_NAME, VIRTUAL_TOOL_NAME_PREFIX, VirtualTool } from './virtualTool';
 import { divideToolsIntoExistingGroups, divideToolsIntoGroups, summarizeToolGroup } from './virtualToolSummarizer';
 import { ISummarizedToolCategory, IToolCategorization, IToolGroupingCache } from './virtualToolTypes';
 import * as Constant from './virtualToolsConstants';
@@ -44,6 +44,10 @@ export class VirtualToolGrouper implements IToolCategorization {
 		this.toolEmbeddingsComputer = _instantiationService.createInstance(ToolEmbeddingsComputer);
 	}
 
+	private get virtualToolEmbeddingRankingEnabled() {
+		return this._configurationService.getExperimentBasedConfig(ConfigKey.Internal.VirtualToolEmbeddingRanking, this._expService);
+	}
+
 	async addGroups(query: string, root: VirtualTool, tools: LanguageModelToolInformation[], token: CancellationToken): Promise<void> {
 		// If there's no need to group tools, just add them all directly;
 		if (tools.length < Constant.START_GROUPING_AFTER_TOOL_COUNT) {
@@ -72,9 +76,8 @@ export class VirtualToolGrouper implements IToolCategorization {
 			}
 		}
 
-		const virtualToolEmbeddingRankingEnabled = this._configurationService.getExperimentBasedConfig(ConfigKey.Internal.VirtualToolEmbeddingRanking, this._expService);
 		const predictedToolsSw = new StopWatch();
-		const predictedToolsPromise = virtualToolEmbeddingRankingEnabled && this._getPredictedTools(query, tools, token).then(tools => ({ tools, durationMs: predictedToolsSw.elapsed() }));
+		const predictedToolsPromise = this.virtualToolEmbeddingRankingEnabled && this._getPredictedTools(query, tools, token).then(tools => ({ tools, durationMs: predictedToolsSw.elapsed() }));
 
 		const grouped = await Promise.all(Object.entries(byToolset).map(([key, tools]) => {
 			if (key === BUILT_IN_GROUP) {
@@ -92,7 +95,7 @@ export class VirtualToolGrouper implements IToolCategorization {
 				const prev = previousGroups.get(tool.name);
 				if (prev) {
 					tool.isExpanded = prev.isExpanded;
-					tool.metadata.preExpanded = prev.metadata.preExpanded;
+					tool.metadata.wasExpandedByDefault = prev.metadata.wasExpandedByDefault;
 					tool.lastUsedOnTurn = prev.lastUsedOnTurn;
 				}
 			}
@@ -101,6 +104,40 @@ export class VirtualToolGrouper implements IToolCategorization {
 		await this._reExpandTools(root, predictedToolsPromise);
 	}
 
+	async recomputeEmbeddingRankings(query: string, root: VirtualTool, token: CancellationToken): Promise<void> {
+		if (!this.virtualToolEmbeddingRankingEnabled) {
+			return;
+		}
+
+		const predictedToolsSw = new StopWatch();
+
+		await this._reExpandTools(root, this._getPredictedTools(query, [...root.tools()], token).then(tools => ({
+			tools,
+			durationMs: predictedToolsSw.elapsed()
+		})));
+	}
+
+	private _addPredictedToolsGroup(root: VirtualTool, predictedTools: LanguageModelToolInformation[]): void {
+		const newGroup = new VirtualTool(EMBEDDINGS_GROUP_NAME, 'Tools with high predicted relevancy for this query', Infinity, {
+			toolsetKey: EMBEDDINGS_GROUP_NAME,
+			wasExpandedByDefault: true,
+			canBeCollapsed: false,
+			groups: [],
+		});
+
+		newGroup.isExpanded = true;
+		for (const tool of predictedTools) {
+			newGroup.contents.push(tool);
+		}
+
+		const idx = root.contents.findIndex(t => t.name === EMBEDDINGS_GROUP_NAME);
+		if (idx >= 0) {
+			root.contents[idx] = newGroup;
+		} else {
+			root.contents.unshift(newGroup);
+		}
+	}
+
 	private async _reExpandTools(root: VirtualTool, predictedToolsPromise: Promise<{ tools: LanguageModelToolInformation[]; durationMs: number }> | false): Promise<void> {
 		if (predictedToolsPromise) {
 			// Aggressively expand groups with predicted tools up to hard limit
@@ -110,8 +147,7 @@ export class VirtualToolGrouper implements IToolCategorization {
 			try {
 				const { tools, durationMs } = await predictedToolsPromise;
 				computeMs = durationMs;
-				this._reExpandToolsToHitBudget(root, g => this._getGroupPredictedRelevancy(g, tools), HARD_TOOL_LIMIT);
-				return;
+				this._addPredictedToolsGroup(root, tools);
 			} catch (e) {
 				error = e;
 			} finally {
@@ -161,30 +197,6 @@ export class VirtualToolGrouper implements IToolCategorization {
 		return [...seen.values()];
 	}
 
-	/**
-	 * Gets the predicted relevancy score for a group based on the highest priority predicted tool it contains.
-	 * Lower scores indicate higher relevancy (earlier in the predictedTools array).
-	 */
-	private _getGroupPredictedRelevancy(group: VirtualTool, predictedTools: LanguageModelToolInformation[]): number {
-		// Create a set of predicted tool names for fast lookup
-		const predictedToolNames = new Set(predictedTools.map(tool => tool.name));
-
-		// Create a map of tool name to its priority (index in predictedTools array)
-		const toolPriority = new Map<string, number>();
-		predictedTools.forEach((tool, index) => {
-			toolPriority.set(tool.name, index);
-		});
-
-		// Find the highest priority (lowest index) predicted tool in this group
-		const priorities = group.contents
-			.filter(tool => 'name' in tool && predictedToolNames.has(tool.name))
-			.map(tool => toolPriority.get(tool.name) ?? Infinity)
-			.filter(index => index !== Infinity);
-
-		// Return the highest priority (lowest index), or Infinity if no predicted tools
-		return priorities.length > 0 ? Math.min(...priorities) : Infinity;
-	}
-
 	/**
 	 * Eagerly expand groups when possible just to reduce the number of indirections.
 	 * Uses the provided ranker function to determine expansion priority.
@@ -215,7 +227,7 @@ export class VirtualToolGrouper implements IToolCategorization {
 			}
 
 			vtool.isExpanded = true;
-			vtool.metadata.preExpanded = true;
+			vtool.metadata.wasExpandedByDefault = true;
 			toolCount = nextCount;
 
 			if (toolCount > targetLimit) {

src/extension/tools/common/virtualTools/virtualToolTypes.ts

@@ -104,6 +104,12 @@ export interface IToolCategorization {
 	 * the appropriate virtual tool or top-level tool in the `root`.
 	 */
 	addGroups(query: string, root: VirtualTool, tools: LanguageModelToolInformation[], token: CancellationToken): Promise<void>;
+
+	/**
+	 * Recalculates the "embeddings" group, when enabled, so relevant tools
+	 * for the query are shown at the top level.
+	 */
+	recomputeEmbeddingRankings(query: string, root: VirtualTool, token: CancellationToken): Promise<void>;
 }
 
 export interface ISummarizedToolCategory {