Refactor agent prompt system to use centralized customization resolution (#2202)

* Refactor agent prompt system to use centralized customization resolution

* Update tests

Author: bhavyaus

src/extension/intents/node/agentIntent.ts

@@ -38,6 +38,7 @@ import { IDefaultIntentRequestHandlerOptions } from '../../prompt/node/defaultIn
 import { IDocumentContext } from '../../prompt/node/documentContext';
 import { IBuildPromptResult, IIntent, IIntentInvocation } from '../../prompt/node/intents';
 import { AgentPrompt, AgentPromptProps } from '../../prompts/node/agent/agentPrompt';
+import { AgentPromptCustomizations, PromptRegistry, ToolAllowlist } from '../../prompts/node/agent/promptRegistry';
 import { PromptRenderer } from '../../prompts/node/base/promptRenderer';
 import { ICodeMapperService } from '../../prompts/node/codeMapper/codeMapperService';
 import { EditCodePrompt2 } from '../../prompts/node/panel/editCodePrompt2';
@@ -52,7 +53,7 @@ import { applyPatch5Description } from '../../tools/node/applyPatchTool';
 import { addCacheBreakpoints } from './cacheBreakpoints';
 import { EditCodeIntent, EditCodeIntentInvocation, EditCodeIntentInvocationOptions, mergeMetadata, toNewChatReferences } from './editCodeIntent';
 
-export const getAgentTools = (instaService: IInstantiationService, request: vscode.ChatRequest) =>
+export const getAgentTools = (instaService: IInstantiationService, request: vscode.ChatRequest, registryToolAllowlist?: ToolAllowlist) =>
 	instaService.invokeFunction(async accessor => {
 		const toolsService = accessor.get<IToolsService>(IToolsService);
 		const testService = accessor.get<ITestProvider>(ITestProvider);
@@ -65,6 +66,14 @@ export const getAgentTools = (instaService: IInstantiationService, request: vsco
 
 		const allowTools: Record<string, boolean> = {};
 
+		if (registryToolAllowlist) {
+			for (const [toolName, allowed] of Object.entries(registryToolAllowlist)) {
+				if (typeof allowed === 'boolean') {
+					allowTools[toolName] = allowed;
+				}
+			}
+		}
+
 		const learned = editToolLearningService.getPreferredEndpointEditTool(model);
 		if (learned) { // a learning-enabled (BYOK) model, we should go with what it prefers
 			allowTools[ToolName.EditFile] = learned.includes(ToolName.EditFile);
@@ -209,6 +218,8 @@ export class AgentIntentInvocation extends EditCodeIntentInvocation implements I
 
 	protected extraPromptProps: Partial<AgentPromptProps> | undefined;
 
+	private _resolvedCustomizations: AgentPromptCustomizations | undefined;
+
 	constructor(
 		intent: IIntent,
 		location: ChatLocation,
@@ -233,14 +244,15 @@ export class AgentIntentInvocation extends EditCodeIntentInvocation implements I
 	}
 
 	public override getAvailableTools(): Promise<vscode.LanguageModelToolInformation[]> {
-		return getAgentTools(this.instantiationService, this.request);
+		return getAgentTools(this.instantiationService, this.request, this._resolvedCustomizations?.toolAllowlist);
 	}
 
 	override async buildPrompt(
 		promptContext: IBuildPromptContext,
 		progress: vscode.Progress<vscode.ChatResponseReferencePart | vscode.ChatResponseProgressPart>,
 		token: vscode.CancellationToken
 	): Promise<IBuildPromptResult> {
+		this._resolvedCustomizations = await PromptRegistry.resolveAllCustomizations(this.instantiationService, this.endpoint);
 		// Add any references from the codebase invocation to the request
 		const codebase = await this._getCodebaseReferences(promptContext, token);
 
@@ -283,7 +295,8 @@ export class AgentIntentInvocation extends EditCodeIntentInvocation implements I
 			},
 			location: this.location,
 			enableCacheBreakpoints: summarizationEnabled,
-			...this.extraPromptProps
+			...this.extraPromptProps,
+			customizations: this._resolvedCustomizations
 		};
 		try {
 			const renderer = PromptRenderer.create(this.instantiationService, endpoint, this.prompt, props);

src/extension/prompts/node/agent/agentPrompt.tsx

@@ -11,9 +11,8 @@ import { ResponseTranslationRules } from '../base/responseTranslationRules';
 import { Tag } from '../base/tag';
 import { EXISTING_CODE_MARKER } from '../panel/codeBlockFormattingRules';
 import { MathIntegrationRules } from '../panel/editorIntegrationRules';
-import { KeepGoingReminder } from './agentPrompt';
-import { CodesearchModeInstructions, DefaultAgentPromptProps, detectToolCapabilities, GenericEditingTips, McpToolInstructions, NotebookInstructions } from './defaultAgentInstructions';
-import { IAgentPrompt, PromptConstructor, PromptRegistry } from './promptRegistry';
+import { CodesearchModeInstructions, DefaultAgentPromptProps, detectToolCapabilities, GenericEditingTips, getEditingReminder, McpToolInstructions, NotebookInstructions, ReminderInstructionsProps } from './defaultAgentInstructions';
+import { IAgentPrompt, PromptRegistry, ReminderInstructionsConstructor, SystemPrompt } from './promptRegistry';
 
 class DefaultAnthropicAgentPrompt extends PromptElement<DefaultAgentPromptProps> {
 	async render(state: void, sizing: PromptSizing) {
@@ -23,7 +22,6 @@ class DefaultAnthropicAgentPrompt extends PromptElement<DefaultAgentPromptProps>
 			<Tag name='instructions'>
 				You are a highly sophisticated automated coding agent with expert-level knowledge across many different programming languages and frameworks.<br />
 				The user will ask a question, or ask you to perform a task, and it may require lots of research to answer correctly. There is a selection of tools that let you perform actions or retrieve helpful context to answer the user's question.<br />
-				<KeepGoingReminder modelFamily={this.props.modelFamily} />
 				You will be given some context and attachments along with the user prompt. You can use them if they are relevant to the task, and ignore them if not.{tools[ToolName.ReadFile] && <> Some attachments may be summarized with omitted sections like `/* Lines 123-456 omitted */`. You can use the {ToolName.ReadFile} tool to read more context if needed. Never pass this omitted line marker to an edit tool.</>}<br />
 				If you can infer the project type (languages, frameworks, and libraries) from the user's query or the context that you have, make sure to keep them in mind when making changes.<br />
 				{!this.props.codesearchMode && <>If the user wants you to implement a feature and they have not specified the files to edit, first break down the user's request into smaller concepts and think about the kinds of files you need to grasp each concept.<br /></>}
@@ -210,14 +208,27 @@ class Claude45DefaultPrompt extends PromptElement<DefaultAgentPromptProps> {
 class AnthropicPromptResolver implements IAgentPrompt {
 	static readonly familyPrefixes = ['claude', 'Anthropic'];
 
-	resolvePrompt(endpoint: IChatEndpoint): PromptConstructor | undefined {
+	resolveSystemPrompt(endpoint: IChatEndpoint): SystemPrompt | undefined {
 		const normalizedModel = endpoint.model?.replace(/\./g, '-');
 		if (normalizedModel?.startsWith('claude-sonnet-4-5') ||
 			normalizedModel?.startsWith('claude-haiku-4-5')) {
 			return Claude45DefaultPrompt;
 		}
 		return DefaultAnthropicAgentPrompt;
 	}
+
+	resolveReminderInstructions(endpoint: IChatEndpoint): ReminderInstructionsConstructor | undefined {
+		return AnthropicReminderInstructions;
+	}
+}
+
+class AnthropicReminderInstructions extends PromptElement<ReminderInstructionsProps> {
+	async render(state: void, sizing: PromptSizing) {
+		return <>
+			{getEditingReminder(this.props.hasEditFileTool, this.props.hasReplaceStringTool, false /* useStrongReplaceStringHint */, this.props.hasMultiReplaceStringTool)}
+			Do NOT create a new markdown file to document each change or summarize your work unless specifically requested by the user.<br />
+		</>;
+	}
 }
 
 PromptRegistry.registerPrompt(AnthropicPromptResolver);

src/extension/prompts/node/agent/anthropicPrompts.tsx

@@ -7,6 +7,7 @@ import { BasePromptElementProps, PromptElement, PromptSizing } from '@vscode/pro
 import type { LanguageModelToolInformation } from 'vscode';
 import { ConfigKey, IConfigurationService } from '../../../../platform/configuration/common/configurationService';
 import { isGpt5PlusFamily } from '../../../../platform/endpoint/common/chatModelCapabilities';
+import { IChatEndpoint } from '../../../../platform/networking/common/networking';
 import { IExperimentationService } from '../../../../platform/telemetry/common/nullExperimentationService';
 import { LanguageModelToolMCPSource } from '../../../../vscodeTypes';
 import { ToolName } from '../../../tools/common/toolNames';
@@ -16,10 +17,9 @@ import { ResponseTranslationRules } from '../base/responseTranslationRules';
 import { Tag } from '../base/tag';
 import { CodeBlockFormattingRules, EXISTING_CODE_MARKER } from '../panel/codeBlockFormattingRules';
 import { MathIntegrationRules } from '../panel/editorIntegrationRules';
-import { KeepGoingReminder } from './agentPrompt';
 
 // Types and interfaces for reusable components
-interface ToolCapabilities extends Partial<Record<ToolName, boolean>> {
+export interface ToolCapabilities extends Partial<Record<ToolName, boolean>> {
 	readonly hasSomeEditTool: boolean;
 }
 
@@ -44,6 +44,65 @@ export interface DefaultAgentPromptProps extends BasePromptElementProps {
 	readonly codesearchMode: boolean | undefined;
 }
 
+export interface ToolReferencesHintProps extends BasePromptElementProps {
+	readonly toolReferences: readonly { name: string }[];
+}
+
+export class DefaultToolReferencesHint extends PromptElement<ToolReferencesHintProps> {
+	async render() {
+		if (!this.props.toolReferences.length) {
+			return;
+		}
+
+		return <>
+			<Tag name='toolReferences'>
+				The user attached the following tools to this message. The userRequest may refer to them using the tool name with "#". These tools are likely relevant to the user's query:<br />
+				{this.props.toolReferences.map(tool => `- ${tool.name}`).join('\n')}
+			</Tag>
+		</>;
+	}
+}
+
+export interface ReminderInstructionsProps extends BasePromptElementProps {
+	readonly endpoint: IChatEndpoint;
+	readonly hasTodoTool: boolean;
+	readonly hasEditFileTool: boolean;
+	readonly hasReplaceStringTool: boolean;
+	readonly hasMultiReplaceStringTool: boolean;
+}
+
+export function getEditingReminder(hasEditFileTool: boolean, hasReplaceStringTool: boolean, useStrongReplaceStringHint: boolean, hasMultiStringReplace: boolean) {
+	const lines = [];
+	if (hasEditFileTool) {
+		lines.push(<>When using the {ToolName.EditFile} tool, avoid repeating existing code, instead use a line comment with \`{EXISTING_CODE_MARKER}\` to represent regions of unchanged code.<br /></>);
+	}
+	if (hasReplaceStringTool) {
+		lines.push(<>
+			When using the {ToolName.ReplaceString} tool, include 3-5 lines of unchanged code before and after the string you want to replace, to make it unambiguous which part of the file should be edited.<br />
+			{hasMultiStringReplace && <>For maximum efficiency, whenever you plan to perform multiple independent edit operations, invoke them simultaneously using {ToolName.MultiReplaceString} tool rather than sequentially. This will greatly improve user's cost and time efficiency leading to a better user experience. Do not announce which tool you're using (for example, avoid saying "I'll implement all the changes using multi_replace_string_in_file").<br /></>}
+		</>);
+	}
+	if (hasEditFileTool && hasReplaceStringTool) {
+		const eitherOr = hasMultiStringReplace ? `${ToolName.ReplaceString} or ${ToolName.MultiReplaceString} tools` : `${ToolName.ReplaceString} tool`;
+		if (useStrongReplaceStringHint) {
+			lines.push(<>You must always try making file edits using the {eitherOr}. NEVER use {ToolName.EditFile} unless told to by the user or by a tool.</>);
+		} else {
+			lines.push(<>It is much faster to edit using the {eitherOr}. Prefer the {eitherOr} for making edits and only fall back to {ToolName.EditFile} if it fails.</>);
+		}
+	}
+
+	return lines;
+}
+
+export class DefaultReminderInstructions extends PromptElement<ReminderInstructionsProps> {
+	async render(state: void, sizing: PromptSizing) {
+		return <>
+			{/* Tool-dependent editing reminders that apply to all models */}
+			{getEditingReminder(this.props.hasEditFileTool, this.props.hasReplaceStringTool, false /* useStrongReplaceStringHint */, this.props.hasMultiReplaceStringTool)}
+		</>;
+	}
+}
+
 /**
  * Base system prompt for agent mode
  */
@@ -55,7 +114,6 @@ export class DefaultAgentPrompt extends PromptElement<DefaultAgentPromptProps> {
 			<Tag name='instructions'>
 				You are a highly sophisticated automated coding agent with expert-level knowledge across many different programming languages and frameworks.<br />
 				The user will ask a question, or ask you to perform a task, and it may require lots of research to answer correctly. There is a selection of tools that let you perform actions or retrieve helpful context to answer the user's question.<br />
-				<KeepGoingReminder modelFamily={this.props.modelFamily} />
 				You will be given some context and attachments along with the user prompt. You can use them if they are relevant to the task, and ignore them if not.{tools[ToolName.ReadFile] && <> Some attachments may be summarized with omitted sections like `/* Lines 123-456 omitted */`. You can use the {ToolName.ReadFile} tool to read more context if needed. Never pass this omitted line marker to an edit tool.</>}<br />
 				If you can infer the project type (languages, frameworks, and libraries) from the user's query or the context that you have, make sure to keep them in mind when making changes.<br />
 				{!this.props.codesearchMode && <>If the user wants you to implement a feature and they have not specified the files to edit, first break down the user's request into smaller concepts and think about the kinds of files you need to grasp each concept.<br /></>}
@@ -156,7 +214,6 @@ export class AlternateGPTPrompt extends PromptElement<DefaultAgentPromptProps> {
 		return <InstructionMessage>
 			<Tag name='gptAgentInstructions'>
 				You are a highly sophisticated coding agent with expert-level knowledge across programming languages and frameworks.<br />
-				<KeepGoingReminder modelFamily={this.props.modelFamily} />
 				You will be given some context and attachments along with the user prompt. You can use them if they are relevant to the task, and ignore them if not.{tools[ToolName.ReadFile] && <> Some attachments may be summarized. You can use the {ToolName.ReadFile} tool to read more context, but only do this if the attached file is incomplete.</>}<br />
 				If you can infer the project type (languages, frameworks, and libraries) from the user's query or the context that you have, make sure to keep them in mind when making changes.<br />
 				Use multiple tools as needed, and do not give up until the task is complete or impossible.<br />

src/extension/prompts/node/agent/defaultAgentInstructions.tsx

@@ -11,9 +11,8 @@ import { ResponseTranslationRules } from '../base/responseTranslationRules';
 import { Tag } from '../base/tag';
 import { EXISTING_CODE_MARKER } from '../panel/codeBlockFormattingRules';
 import { MathIntegrationRules } from '../panel/editorIntegrationRules';
-import { KeepGoingReminder } from './agentPrompt';
-import { CodesearchModeInstructions, DefaultAgentPromptProps, detectToolCapabilities, GenericEditingTips, McpToolInstructions, NotebookInstructions } from './defaultAgentInstructions';
-import { IAgentPrompt, PromptConstructor, PromptRegistry } from './promptRegistry';
+import { CodesearchModeInstructions, DefaultAgentPromptProps, detectToolCapabilities, GenericEditingTips, getEditingReminder, McpToolInstructions, NotebookInstructions, ReminderInstructionsProps } from './defaultAgentInstructions';
+import { IAgentPrompt, PromptRegistry, ReminderInstructionsConstructor, SystemPrompt } from './promptRegistry';
 
 /**
  * Base system prompt for agent mode
@@ -26,7 +25,6 @@ export class DefaultGeminiAgentPrompt extends PromptElement<DefaultAgentPromptPr
 			<Tag name='instructions'>
 				You are a highly sophisticated automated coding agent with expert-level knowledge across many different programming languages and frameworks.<br />
 				The user will ask a question, or ask you to perform a task, and it may require lots of research to answer correctly. There is a selection of tools that let you perform actions or retrieve helpful context to answer the user's question.<br />
-				<KeepGoingReminder modelFamily={this.props.modelFamily} />
 				You will be given some context and attachments along with the user prompt. You can use them if they are relevant to the task, and ignore them if not.{tools[ToolName.ReadFile] && <> Some attachments may be summarized with omitted sections like `/* Lines 123-456 omitted */`. You can use the {ToolName.ReadFile} tool to read more context if needed. Never pass this omitted line marker to an edit tool.</>}<br />
 				If you can infer the project type (languages, frameworks, and libraries) from the user's query or the context that you have, make sure to keep them in mind when making changes.<br />
 				{!this.props.codesearchMode && <>If the user wants you to implement a feature and they have not specified the files to edit, first break down the user's request into smaller concepts and think about the kinds of files you need to grasp each concept.<br /></>}
@@ -119,9 +117,22 @@ class GeminiPromptResolver implements IAgentPrompt {
 
 	static readonly familyPrefixes = ['gemini'];
 
-	resolvePrompt(endpoint: IChatEndpoint): PromptConstructor | undefined {
+	resolveSystemPrompt(endpoint: IChatEndpoint): SystemPrompt | undefined {
 		return DefaultGeminiAgentPrompt;
 	}
+
+	resolveReminderInstructions(endpoint: IChatEndpoint): ReminderInstructionsConstructor | undefined {
+		return GeminiReminderInstructions;
+	}
+}
+
+class GeminiReminderInstructions extends PromptElement<ReminderInstructionsProps> {
+	async render(state: void, sizing: PromptSizing) {
+		// Gemini models need the strong replace string hint
+		return <>
+			{getEditingReminder(this.props.hasEditFileTool, this.props.hasReplaceStringTool, true /* useStrongReplaceStringHint */, this.props.hasMultiReplaceStringTool)}
+		</>;
+	}
 }