Refactor RunSubagentTool to enforce multiplier constraints on model selection and update related tests

hediet · hediet · commit b947d0b27f2e · 2026-04-09T17:40:56.000+02:00
diff --git a/src/vs/workbench/contrib/chat/common/tools/builtinTools/runSubagentTool.ts b/src/vs/workbench/contrib/chat/common/tools/builtinTools/runSubagentTool.ts
@@ -121,7 +121,7 @@ export class RunSubagentTool extends Disposable implements IToolImpl {
 
 		properties.model = {
 			type: 'string',
-			description: 'Optional model for the subagent. Use a fast model for simple lookups and searches, or a reasoning model for complex analysis. Format: "Model Name (Vendor)" - vendor is usually "copilot". If not provided, uses the current model.',
+			description: 'Optional model for the subagent. Format: "Model Name (Vendor)", vendor is usually "copilot". Only use to enforce a specific model.',
 		};
 
 		const required: string[] = ['prompt', 'description'];
@@ -424,25 +424,77 @@ export class RunSubagentTool extends Disposable implements IToolImpl {
 	}
 
 	/**
-	 * Returns the list of available model qualified names suitable for agent mode.
+	 * Checks if a model exceeds the main model's cost tier based on multiplier.
+	 * @returns An object with `exceeds: true` and a reason string if blocked, or `exceeds: false` if allowed.
 	 */
-	private getAvailableModelNames(): string[] {
-		return this.languageModelsService.getLanguageModelIds()
-			.map(id => this.languageModelsService.lookupLanguageModel(id))
-			.filter((m): m is ILanguageModelChatMetadata =>
-				!!m
-				&& ILanguageModelChatMetadata.suitableForAgentMode(m)
-				&& m.isUserSelectable !== false
-				&& !m.targetChatSessionType
-			)
-			.map(m => ILanguageModelChatMetadata.asQualifiedName(m));
+	private checkMultiplierConstraint(modelId: string, mainModelId: string | undefined): { exceeds: false } | { exceeds: true; reason: string } {
+		if (!mainModelId || modelId === mainModelId) {
+			return { exceeds: false };
+		}
+
+		const mainModelMetadata = this.languageModelsService.lookupLanguageModel(mainModelId);
+		const modelMetadata = this.languageModelsService.lookupLanguageModel(modelId);
+		const mainMultiplier = mainModelMetadata?.multiplierNumeric;
+		const modelMultiplier = modelMetadata?.multiplierNumeric;
+
+		if (mainMultiplier !== undefined && modelMultiplier !== undefined && modelMultiplier > mainMultiplier) {
+			return {
+				exceeds: true,
+				reason: `exceeds the current model's cost tier (${modelMultiplier}x vs ${mainMultiplier}x)`
+			};
+		}
+
+		return { exceeds: false };
+	}
+
+	/**
+	 * Returns information about available models for error messages.
+	 * Includes which models are unavailable due to multiplier restrictions.
+	 */
+	private getAvailableModelsInfo(mainModelId: string | undefined): string {
+		const models = this.languageModelsService.getLanguageModelIds()
+			.map(id => ({ id, metadata: this.languageModelsService.lookupLanguageModel(id) }))
+			.filter((m): m is { id: string; metadata: ILanguageModelChatMetadata } =>
+				!!m.metadata
+				&& ILanguageModelChatMetadata.suitableForAgentMode(m.metadata)
+				&& m.metadata.isUserSelectable !== false
+				&& !m.metadata.targetChatSessionType
+			);
+
+		if (models.length === 0) {
+			return 'No models available.';
+		}
+
+		const available: string[] = [];
+		const unavailableDueToMultiplier: string[] = [];
+
+		for (const { id, metadata } of models) {
+			const qualifiedName = ILanguageModelChatMetadata.asQualifiedName(metadata);
+			const check = this.checkMultiplierConstraint(id, mainModelId);
+
+			if (check.exceeds) {
+				unavailableDueToMultiplier.push(qualifiedName);
+			} else {
+				available.push(qualifiedName);
+			}
+		}
+
+		const parts: string[] = [];
+		if (available.length > 0) {
+			parts.push(`Available models: ${available.join(', ')}`);
+		}
+		if (unavailableDueToMultiplier.length > 0) {
+			parts.push(`Unavailable (exceeds current model's cost tier): ${unavailableDueToMultiplier.join(', ')}`);
+		}
+
+		return parts.join('. ') || 'No models available.';
 	}
 
 	/**
-	 * Resolves the model to be used by a subagent, applying multiplier-based
-	 * fallback to avoid using a more expensive model than the main agent.
+	 * Resolves the model to be used by a subagent.
 	 * @param explicitModelQualifiedName Optional explicit model specified by the caller.
-	 *        If provided and not found, throws an error with the list of available models.
+	 *        If provided and not found or not allowed, throws an error with available models.
+	 * @throws Error if the requested model is not found or exceeds the main model's cost tier.
 	 */
 	private resolveSubagentModel(subagent: ICustomAgent | undefined, mainModelId: string | undefined, explicitModelQualifiedName?: string): { modeModelId: string | undefined; resolvedModelName: string | undefined } {
 		let modeModelId = mainModelId;
@@ -456,9 +508,7 @@ export class RunSubagentTool extends Disposable implements IToolImpl {
 				explicitModelResolved = true;
 			} else {
 				// Model not found - throw error with available models
-				const available = this.getAvailableModelNames();
-				const availableList = available.length > 0 ? `Available models: ${available.join(', ')}` : 'No models available.';
-				throw new Error(`Requested model '${explicitModelQualifiedName}' not found. ${availableList}`);
+				throw new Error(`Requested model '${explicitModelQualifiedName}' not found. ${this.getAvailableModelsInfo(mainModelId)}`);
 			}
 		}
 
@@ -476,16 +526,12 @@ export class RunSubagentTool extends Disposable implements IToolImpl {
 			}
 		}
 
-		// If the resolved model has a larger multiplier than the main agent's model,
-		// fall back to the main agent's model to avoid using a more expensive model.
-		if (modeModelId && modeModelId !== mainModelId) {
-			const mainModelMetadata = mainModelId ? this.languageModelsService.lookupLanguageModel(mainModelId) : undefined;
-			const subagentModelMetadata = this.languageModelsService.lookupLanguageModel(modeModelId);
-			const mainMultiplier = mainModelMetadata?.multiplierNumeric;
-			const subagentMultiplier = subagentModelMetadata?.multiplierNumeric;
-			if (mainMultiplier !== undefined && subagentMultiplier !== undefined && subagentMultiplier > mainMultiplier) {
-				this.logService.warn(`[RunSubagentTool] Requested model '${subagentModelMetadata?.name}' (multiplier: ${subagentMultiplier}) has a larger multiplier than the main agent model '${mainModelMetadata?.name}' (multiplier: ${mainMultiplier}). Falling back to the main agent model.`);
-				modeModelId = mainModelId;
+		// Check multiplier constraint - throw error if requested model exceeds main model's cost tier
+		if (modeModelId) {
+			const check = this.checkMultiplierConstraint(modeModelId, mainModelId);
+			if (check.exceeds) {
+				const modelMetadata = this.languageModelsService.lookupLanguageModel(modeModelId);
+				throw new Error(`Requested model '${modelMetadata?.name}' ${check.reason}. ${this.getAvailableModelsInfo(mainModelId)}`);
 			}
 		}
 
diff --git a/src/vs/workbench/contrib/chat/test/common/tools/builtinTools/runSubagentTool.test.ts b/src/vs/workbench/contrib/chat/test/common/tools/builtinTools/runSubagentTool.test.ts
@@ -397,7 +397,7 @@ suite('RunSubagentTool', () => {
 			};
 		}
 
-		test('falls back to main model when subagent model has higher multiplier', async () => {
+		test('throws error when subagent model has higher multiplier', async () => {
 			const mainMeta = createMetadata('GPT-4o', 1);
 			const expensiveMeta = createMetadata('O3 Pro', 50);
 			const models = new Map([
@@ -411,22 +411,21 @@ suite('RunSubagentTool', () => {
 			const agent = createAgent('ExpensiveAgent', ['O3 Pro (TestVendor)']);
 			const tool = createTool({ models, qualifiedNameMap, customAgents: [agent] });
 
-			const result = await tool.prepareToolInvocation({
-				parameters: { prompt: 'test', description: 'test task', agentName: 'ExpensiveAgent' },
-				toolCallId: 'call-1',
-				modelId: 'main-model-id',
-				chatSessionResource: URI.parse('test://session'),
-			}, CancellationToken.None);
-
-			assert.ok(result);
-			// Should fall back to the main model's name, not the expensive model
-			assert.deepStrictEqual(result.toolSpecificData, {
-				kind: 'subagent',
-				description: 'test task',
-				agentName: 'ExpensiveAgent',
-				prompt: 'test',
-				modelName: 'GPT-4o',
-			});
+			await assert.rejects(
+				() => tool.prepareToolInvocation({
+					parameters: { prompt: 'test', description: 'test task', agentName: 'ExpensiveAgent' },
+					toolCallId: 'call-1',
+					modelId: 'main-model-id',
+					chatSessionResource: URI.parse('test://session'),
+				}, CancellationToken.None),
+				(err: Error) => {
+					assert.ok(err.message.includes('O3 Pro'));
+					assert.ok(err.message.includes('exceeds'));
+					assert.ok(err.message.includes('cost tier'));
+					assert.ok(err.message.includes('Unavailable'));
+					return true;
+				}
+			);
 		});
 
 		test('uses subagent model when it has equal multiplier', async () => {
@@ -752,7 +751,7 @@ suite('RunSubagentTool', () => {
 			});
 		});
 
-		test('falls back to main model when explicit model has higher multiplier', async () => {
+		test('throws error when explicit model has higher multiplier', async () => {
 			const mainMeta = createMetadata('GPT-4o', 1);
 			const expensiveMeta = createMetadata('O3 Pro', 50);
 			const models = new Map([
@@ -765,21 +764,21 @@ suite('RunSubagentTool', () => {
 
 			const tool = createTool({ models, qualifiedNameMap });
 
-			const result = await tool.prepareToolInvocation({
-				parameters: { prompt: 'test', description: 'test task', model: 'O3 Pro (TestVendor)' },
-				toolCallId: 'model-call-3',
-				modelId: 'main-model-id',
-				chatSessionResource: URI.parse('test://session'),
-			}, CancellationToken.None);
-
-			assert.ok(result);
-			assert.deepStrictEqual(result.toolSpecificData, {
-				kind: 'subagent',
-				description: 'test task',
-				agentName: undefined,
-				prompt: 'test',
-				modelName: 'GPT-4o',
-			});
+			await assert.rejects(
+				() => tool.prepareToolInvocation({
+					parameters: { prompt: 'test', description: 'test task', model: 'O3 Pro (TestVendor)' },
+					toolCallId: 'model-call-3',
+					modelId: 'main-model-id',
+					chatSessionResource: URI.parse('test://session'),
+				}, CancellationToken.None),
+				(err: Error) => {
+					assert.ok(err.message.includes('O3 Pro'));
+					assert.ok(err.message.includes('exceeds'));
+					assert.ok(err.message.includes('cost tier'));
+					assert.ok(err.message.includes('Unavailable'));
+					return true;
+				}
+			);
 		});
 
 		test('throws error with available models when explicit model is not found', async () => {
