tools: trigger tool grouping on 128 limit automatically (#385)

Triggers tool grouping when users hit the 128 limit. Reworks an EXP
setting which allows the limit to be configurable (lowered or raised
arbitrarily high.) Users of the setting keep the existing 64 tool limit.
Will have a corresponding VS Code PR.

Author: connor4312

src/extension/intents/node/agentIntent.ts

@@ -95,8 +95,6 @@ export class AgentIntent extends EditCodeIntent {
 		@IExperimentationService expService: IExperimentationService,
 		@ICodeMapperService codeMapperService: ICodeMapperService,
 		@IWorkspaceService workspaceService: IWorkspaceService,
-		@IConfigurationService private readonly _configurationService: IConfigurationService,
-		@IExperimentationService private readonly _experimentationService: IExperimentationService,
 		@IToolGroupingService private readonly _toolGroupingService: IToolGroupingService,
 	) {
 		super(instantiationService, endpointProvider, configurationService, expService, codeMapperService, workspaceService, { intentInvocation: AgentIntentInvocation, processCodeblocks: false });
@@ -113,7 +111,8 @@ export class AgentIntent extends EditCodeIntent {
 
 	private async listTools(request: vscode.ChatRequest, stream: vscode.ChatResponseStream, token: CancellationToken) {
 		const editingTools = await getTools(this.instantiationService, request);
-		if (!this._configurationService.getExperimentBasedConfig(ConfigKey.VirtualTools, this._experimentationService)) {
+		const grouping = this._toolGroupingService.create(editingTools);
+		if (!grouping.isEnabled) {
 			stream.markdown(`Available tools: \n${editingTools.map(tool => `- ${tool.name}`).join('\n')}\n`);
 			return;
 		}
@@ -137,7 +136,7 @@ export class AgentIntent extends EditCodeIntent {
 			}
 		}
 
-		const tools = await this._toolGroupingService.create(editingTools).computeAll(token);
+		const tools = await grouping.computeAll(token);
 		tools.forEach(t => printTool(t));
 		stream.markdown(str);
 

src/extension/prompt/node/defaultIntentRequestHandler.ts

@@ -11,7 +11,6 @@ import { IAuthenticationChatUpgradeService } from '../../../platform/authenticat
 import { ICopilotTokenStore } from '../../../platform/authentication/common/copilotTokenStore';
 import { CanceledResult, ChatFetchResponseType, ChatLocation, ChatResponse, getErrorDetailsFromChatFetchError } from '../../../platform/chat/common/commonTypes';
 import { IConversationOptions } from '../../../platform/chat/common/conversationOptions';
-import { ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService';
 import { IEditSurvivalTrackerService, IEditSurvivalTrackingSession, NullEditSurvivalTrackingSession } from '../../../platform/editSurvivalTracking/common/editSurvivalTrackerService';
 import { IEndpointProvider } from '../../../platform/endpoint/common/endpointProvider';
 import { HAS_IGNORED_FILES_MESSAGE } from '../../../platform/ignore/common/ignoreService';
@@ -508,7 +507,6 @@ class DefaultToolCallingLoop extends ToolCallingLoop<IDefaultToolLoopOptions> {
 		@IAuthenticationChatUpgradeService authenticationChatUpgradeService: IAuthenticationChatUpgradeService,
 		@ITelemetryService telemetryService: ITelemetryService,
 		@IToolGroupingService private readonly toolGroupingService: IToolGroupingService,
-		@IConfigurationService private readonly _configurationService: IConfigurationService,
 		@IExperimentationService private readonly _experimentationService: IExperimentationService,
 		@ICopilotTokenStore private readonly _copilotTokenStore: ICopilotTokenStore,
 		@IThinkingDataService thinkingDataService: IThinkingDataService,
@@ -557,7 +555,7 @@ class DefaultToolCallingLoop extends ToolCallingLoop<IDefaultToolLoopOptions> {
 	private async _doMirroredCallWithVirtualTools(delta: IResponseDelta, messages: Raw.ChatMessage[], requestOptions: OptionalChatRequestParams) {
 		const shouldDo = !this._didParallelToolCallLoop
 			&& this._copilotTokenStore.copilotToken?.isInternal
-			&& !DefaultToolCallingLoop.toolGrouping;
+			&& !DefaultToolCallingLoop.toolGrouping?.isEnabled;
 		if (!shouldDo) {
 			return;
 		}
@@ -692,16 +690,16 @@ class DefaultToolCallingLoop extends ToolCallingLoop<IDefaultToolLoopOptions> {
 
 	protected override async getAvailableTools(outputStream: ChatResponseStream | undefined, token: CancellationToken): Promise<LanguageModelToolInformation[]> {
 		const tools = await this.options.invocation.getAvailableTools?.() ?? [];
-		if (!this._configurationService.getExperimentBasedConfig(ConfigKey.VirtualTools, this._experimentationService)) {
-			return tools;
-		}
-
 		if (DefaultToolCallingLoop.toolGrouping) {
 			DefaultToolCallingLoop.toolGrouping.tools = tools;
 		} else {
 			DefaultToolCallingLoop.toolGrouping = this.toolGroupingService.create(tools);
 		}
 
+		if (!DefaultToolCallingLoop.toolGrouping.isEnabled) {
+			return tools;
+		}
+
 		const computePromise = DefaultToolCallingLoop.toolGrouping.compute(token);
 
 		// Show progress if this takes a moment...

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

@@ -4,7 +4,8 @@
  *--------------------------------------------------------------------------------------------*/
 
 import type { LanguageModelToolInformation } from 'vscode';
-import { HARD_TOOL_LIMIT } from '../../../../platform/configuration/common/configurationService';
+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 { CancellationToken } from '../../../../util/vs/base/common/cancellation';
@@ -16,6 +17,19 @@ import { VirtualToolGrouper } from './virtualToolGrouper';
 import * as Constant from './virtualToolsConstants';
 import { IToolCategorization, IToolGrouping } from './virtualToolTypes';
 
+export function computeToolGroupingMinThreshold(experimentationService: IExperimentationService, configurationService: IConfigurationService): number {
+	let threshold = HARD_TOOL_LIMIT;
+
+	const override = experimentationService.getTreatmentVariable<number>('vscode', 'copilotchat.virtualToolThreshold');
+	if (override) {
+		threshold = override;
+	} else if (configurationService.getExperimentBasedConfig(ConfigKey.VirtualTools, experimentationService)) {
+		threshold = Constant.START_GROUPING_AFTER_TOOL_COUNT;
+	}
+
+	return threshold;
+}
+
 export class ToolGrouping implements IToolGrouping {
 
 	private readonly _root = new VirtualTool(VIRTUAL_TOOL_NAME_PREFIX, '', Infinity, { groups: [], toolsetKey: '', preExpanded: true });
@@ -36,10 +50,16 @@ export class ToolGrouping implements IToolGrouping {
 		}
 	}
 
+	public get isEnabled() {
+		return this._tools.length >= computeToolGroupingMinThreshold(this._experimentationService, this._configurationService);
+	}
+
 	constructor(
 		private _tools: readonly LanguageModelToolInformation[],
 		@IInstantiationService private readonly _instantiationService: IInstantiationService,
 		@ITelemetryService private readonly _telemetryService: ITelemetryService,
+		@IConfigurationService private readonly _configurationService: IConfigurationService,
+		@IExperimentationService private readonly _experimentationService: IExperimentationService
 	) {
 		this._root.isExpanded = true;
 	}

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

@@ -4,14 +4,24 @@
  *--------------------------------------------------------------------------------------------*/
 
 import type { LanguageModelToolInformation } from 'vscode';
+import { IConfigurationService } from '../../../../platform/configuration/common/configurationService';
+import { IExperimentationService } from '../../../../platform/telemetry/common/nullExperimentationService';
 import { IInstantiationService } from '../../../../util/vs/platform/instantiation/common/instantiation';
-import { ToolGrouping } from './toolGrouping';
+import { computeToolGroupingMinThreshold, ToolGrouping } from './toolGrouping';
 import { IToolGrouping, IToolGroupingService } from './virtualToolTypes';
 
 export class ToolGroupingService implements IToolGroupingService {
 	declare readonly _serviceBrand: undefined;
 
-	constructor(@IInstantiationService private readonly _instantiationService: IInstantiationService) { }
+	constructor(
+		@IInstantiationService private readonly _instantiationService: IInstantiationService,
+		@IConfigurationService private readonly _configurationService: IConfigurationService,
+		@IExperimentationService private readonly _experimentationService: IExperimentationService
+	) { }
+
+	public get threshold() {
+		return computeToolGroupingMinThreshold(this._experimentationService, this._configurationService);
+	}
 
 	create(tools: readonly LanguageModelToolInformation[]): IToolGrouping {
 		return this._instantiationService.createInstance(ToolGrouping, tools);

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

@@ -14,6 +14,11 @@ export interface IToolGrouping {
 	 */
 	tools: readonly LanguageModelToolInformation[];
 
+	/**
+	 * Whether tool grouping logic is enabled at the current tool threshold.
+	 */
+	isEnabled: boolean;
+
 	/**
 	 * Should be called for each model tool call. Returns a tool result if the
 	 * call was a virtual tool call that was expanded.
@@ -52,6 +57,10 @@ export interface IToolGrouping {
 
 export interface IToolGroupingService {
 	_serviceBrand: undefined;
+	/**
+	 * The current tool count threshold for grouping to kick in.
+	 */
+	threshold: number;
 	/**
 	 * Creates a tool grouping for a request, based on its conversation and the
 	 * initial set of tools.

src/extension/tools/test/node/virtualTools/virtualToolGrouping.spec.ts

@@ -5,7 +5,7 @@
 
 import { afterEach, beforeEach, describe, expect, it } from 'vitest';
 import type { LanguageModelTextPart, LanguageModelToolInformation } from 'vscode';
-import { HARD_TOOL_LIMIT } from '../../../../../platform/configuration/common/configurationService';
+import { HARD_TOOL_LIMIT, IConfigurationService } from '../../../../../platform/configuration/common/configurationService';
 import { ITelemetryService } from '../../../../../platform/telemetry/common/telemetry';
 import { ITestingServicesAccessor } from '../../../../../platform/test/node/services';
 import { shuffle } from '../../../../../util/vs/base/common/arrays';
@@ -17,6 +17,7 @@ import { ToolGrouping } from '../../../common/virtualTools/toolGrouping';
 import { VIRTUAL_TOOL_NAME_PREFIX, VirtualTool } from '../../../common/virtualTools/virtualTool';
 import { TRIM_THRESHOLD } from '../../../common/virtualTools/virtualToolsConstants';
 import { IToolCategorization } from '../../../common/virtualTools/virtualToolTypes';
+import { IExperimentationService } from '../../../../../platform/telemetry/common/nullExperimentationService';
 
 describe('Virtual Tools - Grouping', () => {
 	let accessor: ITestingServicesAccessor;
@@ -28,8 +29,10 @@ describe('Virtual Tools - Grouping', () => {
 			_tools: readonly LanguageModelToolInformation[],
 			@IInstantiationService _instantiationService: IInstantiationService,
 			@ITelemetryService _telemetryService: ITelemetryService,
+			@IConfigurationService _configurationService: IConfigurationService,
+			@IExperimentationService _experimentationService: IExperimentationService
 		) {
-			super(_tools, _instantiationService, _telemetryService);
+			super(_tools, _instantiationService, _telemetryService, _configurationService, _experimentationService);
 			this._grouper = mockGrouper;
 		}
 

src/extension/tools/vscode-node/tools.ts

@@ -8,14 +8,17 @@ import { Disposable } from '../../../util/vs/base/common/lifecycle';
 import { getContributedToolName } from '../common/toolNames';
 import { IToolsService } from '../common/toolsService';
 
-import { IToolGroupingCache } from '../common/virtualTools/virtualToolTypes';
+import { IExperimentationService } from '../../../platform/telemetry/common/nullExperimentationService';
+import { IToolGroupingCache, IToolGroupingService } from '../common/virtualTools/virtualToolTypes';
 import '../node/allTools';
 import './allTools';
 
 export class ToolsContribution extends Disposable {
 	constructor(
 		@IToolsService toolsService: IToolsService,
 		@IToolGroupingCache toolGrouping: IToolGroupingCache,
+		@IToolGroupingService toolGroupingService: IToolGroupingService,
+		@IExperimentationService experimentationService: IExperimentationService
 	) {
 		super();
 
@@ -27,5 +30,9 @@ export class ToolsContribution extends Disposable {
 			await toolGrouping.clear();
 			vscode.window.showInformationMessage('Tool groups have been reset. They will be regenerated on the next agent request.');
 		}));
+
+		experimentationService.initializePromise.then(() => {
+			vscode.commands.executeCommand('setContext', 'chat.toolGroupingThreshold', toolGroupingService.threshold);
+		});
 	}
 }

test/base/extHostContext/simulationExtHostToolsService.ts

@@ -9,16 +9,17 @@ import type { CancellationToken, ChatRequest, LanguageModelTool, LanguageModelTo
 import { getToolName, ToolName } from '../../../src/extension/tools/common/toolNames';
 import { ICopilotTool } from '../../../src/extension/tools/common/toolsRegistry';
 import { BaseToolsService, IToolsService } from '../../../src/extension/tools/common/toolsService';
+import { getPackagejsonToolsForTest } from '../../../src/extension/tools/node/test/testToolsService';
 import { ToolsContribution } from '../../../src/extension/tools/vscode-node/tools';
 import { ToolsService } from '../../../src/extension/tools/vscode-node/toolsService';
 import { packageJson } from '../../../src/platform/env/common/packagejson';
 import { ILogService } from '../../../src/platform/log/common/logService';
+import { NullExperimentationService } from '../../../src/platform/telemetry/common/nullExperimentationService';
+import { raceTimeout } from '../../../src/util/vs/base/common/async';
 import { CancellationError } from '../../../src/util/vs/base/common/errors';
 import { Iterable } from '../../../src/util/vs/base/common/iterator';
 import { IInstantiationService } from '../../../src/util/vs/platform/instantiation/common/instantiation';
 import { logger } from '../../simulationLogger';
-import { raceTimeout } from '../../../src/util/vs/base/common/async';
-import { getPackagejsonToolsForTest } from '../../../src/extension/tools/node/test/testToolsService';
 
 export class SimulationExtHostToolsService extends BaseToolsService implements IToolsService {
 	declare readonly _serviceBrand: undefined;
@@ -63,7 +64,7 @@ export class SimulationExtHostToolsService extends BaseToolsService implements I
 	}
 
 	private ensureToolsRegistered() {
-		this._lmToolRegistration ??= new ToolsContribution(this, {} as any);
+		this._lmToolRegistration ??= new ToolsContribution(this, {} as any, { threshold: Infinity } as any, new NullExperimentationService());
 	}
 
 	getCopilotTool(name: string): ICopilotTool<any> | undefined {