edits: allow oai config and byok api config for edit tools (#1161)

* edits: allow oai config and byok api config for edit tools

Adds on to #1154 and builds on microsoft/vscode#268506

Allows 3p providers to configure preferred edit tools for endpoints
(once the API is finalized) and uses that to allow the edit tools to
be configured in `github.copilot.chat.customOAIModels` for power users.

* comment

* fix gdpr comment

Author: connor4312

package.json

@@ -2997,6 +2997,19 @@
 									"description": "Whether the model requires an API key for authentication",
 									"default": true
 								},
+								"editTools": {
+									"type": "array",
+									"description": "List of edit tools supported by the model. If this is not configured, the editor will try multiple edit tools and pick the best one.\n\n- 'find-replace': Find and replace text in a document.\n- 'multi-find-replace': Find and replace text in a document.\n- 'apply-patch': A file-oriented diff format used by some OpenAI models\n- 'code-rewrite': A general but slower editing tool that allows the model to rewrite and code snippet and provide only the replacement to the editor.",
+									"items": {
+										"type": "string",
+										"enum": [
+											"find-replace",
+											"multi-find-replace",
+											"apply-patch",
+											"code-rewrite"
+										]
+									}
+								},
 								"thinking": {
 									"type": "boolean",
 									"default": false,

src/extension/byok/common/byokProvider.ts

@@ -5,7 +5,7 @@
 import type { Disposable, LanguageModelChatInformation, LanguageModelChatProvider, LanguageModelDataPart, LanguageModelTextPart, LanguageModelThinkingPart, LanguageModelToolCallPart } from 'vscode';
 import { CopilotToken } from '../../../platform/authentication/common/copilotToken';
 import { ICAPIClientService } from '../../../platform/endpoint/common/capiClient';
-import { IChatModelInformation } from '../../../platform/endpoint/common/endpointProvider';
+import { EndpointEditToolName, IChatModelInformation } from '../../../platform/endpoint/common/endpointProvider';
 import { TokenizerType } from '../../../util/common/tokenizer';
 import { localize } from '../../../util/vs/nls';
 
@@ -54,6 +54,7 @@ export interface BYOKModelCapabilities {
 	toolCalling: boolean;
 	vision: boolean;
 	thinking?: boolean;
+	editTools?: EndpointEditToolName[];
 }
 
 export interface BYOKModelRegistry {

src/extension/byok/vscode-node/customOAIModelConfigurator.ts

@@ -5,6 +5,7 @@
 
 import { InputBoxOptions, LanguageModelChatInformation, QuickInputButtons, QuickPickItem, window } from 'vscode';
 import { Config, ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService';
+import { EndpointEditToolName } from '../../../platform/endpoint/common/endpointProvider';
 import { DisposableStore } from '../../../util/vs/base/common/lifecycle';
 import { BYOKModelProvider } from '../common/byokProvider';
 
@@ -15,6 +16,7 @@ interface ModelConfig {
 	vision: boolean;
 	maxInputTokens: number;
 	maxOutputTokens: number;
+	editTools?: EndpointEditToolName[];
 	requiresAPIKey?: boolean;
 	thinking?: boolean;
 }
@@ -414,6 +416,7 @@ export class CustomOAIModelConfigurator {
 		}
 
 		return {
+			...currentConfig,
 			name: modelName.trim(),
 			url: url.trim(),
 			toolCalling: capabilities.toolCalling,

src/extension/byok/vscode-node/customOAIProvider.ts

@@ -5,6 +5,7 @@
 
 import { CancellationToken, LanguageModelChatInformation, LanguageModelChatMessage, LanguageModelChatMessage2, LanguageModelResponsePart2, Progress, ProvideLanguageModelChatResponseOptions, QuickPickItem, window } from 'vscode';
 import { ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService';
+import { EndpointEditToolName, isEndpointEditToolName } from '../../../platform/endpoint/common/endpointProvider';
 import { ILogService } from '../../../platform/log/common/logService';
 import { IExperimentationService } from '../../../platform/telemetry/common/nullExperimentationService';
 import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation';
@@ -66,8 +67,8 @@ export class CustomOAIBYOKModelProvider implements BYOKModelProvider<CustomOAIMo
 		return resolveCustomOAIUrl(modelId, url);
 	}
 
-	private getUserModelConfig(): Record<string, { name: string; url: string; toolCalling: boolean; vision: boolean; maxInputTokens: number; maxOutputTokens: number; requiresAPIKey: boolean; thinking?: boolean }> {
-		const modelConfig = this._configurationService.getConfig(this.getConfigKey()) as Record<string, { name: string; url: string; toolCalling: boolean; vision: boolean; maxInputTokens: number; maxOutputTokens: number; requiresAPIKey: boolean; thinking?: boolean }>;
+	private getUserModelConfig(): Record<string, { name: string; url: string; toolCalling: boolean; vision: boolean; maxInputTokens: number; maxOutputTokens: number; requiresAPIKey: boolean; thinking?: boolean; editTools?: EndpointEditToolName[] }> {
+		const modelConfig = this._configurationService.getConfig(this.getConfigKey()) as Record<string, { name: string; url: string; toolCalling: boolean; vision: boolean; maxInputTokens: number; maxOutputTokens: number; requiresAPIKey: boolean; thinking?: boolean; editTools?: EndpointEditToolName[] }>;
 		return modelConfig;
 	}
 
@@ -88,6 +89,7 @@ export class CustomOAIBYOKModelProvider implements BYOKModelProvider<CustomOAIMo
 				maxInputTokens: modelInfo.maxInputTokens,
 				maxOutputTokens: modelInfo.maxOutputTokens,
 				thinking: modelInfo.thinking,
+				editTools: modelInfo.editTools,
 			};
 		}
 		return models;
@@ -129,7 +131,8 @@ export class CustomOAIBYOKModelProvider implements BYOKModelProvider<CustomOAIMo
 			tooltip: `${capabilities.name} is contributed via the ${this.providerName} provider.`,
 			capabilities: {
 				toolCalling: capabilities.toolCalling,
-				imageInput: capabilities.vision
+				imageInput: capabilities.vision,
+				editTools: capabilities.editTools
 			},
 			thinking: capabilities.thinking || false,
 		};
@@ -168,7 +171,8 @@ export class CustomOAIBYOKModelProvider implements BYOKModelProvider<CustomOAIMo
 			vision: !!model.capabilities?.imageInput || false,
 			name: model.name,
 			url: model.url,
-			thinking: model.thinking
+			thinking: model.thinking,
+			editTools: model.capabilities.editTools?.filter(isEndpointEditToolName),
 		});
 		const openAIChatEndpoint = this._instantiationService.createInstance(OpenAIEndpoint, modelInfo, apiKey ?? '', model.url);
 		return this._lmWrapper.provideLanguageModelResponse(openAIChatEndpoint, messages, options, options.requestInitiator, progress, token);

src/extension/tools/common/editToolLearningService.ts

@@ -7,11 +7,13 @@ import type { LanguageModelChat } from 'vscode';
 import { IEndpointProvider } from '../../../platform/endpoint/common/endpointProvider';
 import { IVSCodeExtensionContext } from '../../../platform/extContext/common/extensionContext';
 import { IChatEndpoint } from '../../../platform/networking/common/networking';
+import { ITelemetryService } from '../../../platform/telemetry/common/telemetry';
 import { createServiceIdentifier } from '../../../util/common/services';
 import { LRUCache } from '../../../util/vs/base/common/map';
 import { mapValues } from '../../../util/vs/base/common/objects';
+import { isDefined } from '../../../util/vs/base/common/types';
 import { EditTools as _EditTools, EDIT_TOOL_LEARNING_STATES, IEditToolLearningData, LearningConfig, State } from './editToolLearningStates';
-import { ToolName } from './toolNames';
+import { byokEditToolNamesToToolNames, ToolName } from './toolNames';
 
 export type EditTools = _EditTools;
 
@@ -49,6 +51,7 @@ export class EditToolLearningService implements IEditToolLearningService {
 	constructor(
 		@IVSCodeExtensionContext private readonly _context: IVSCodeExtensionContext,
 		@IEndpointProvider private readonly _endpointProvider: IEndpointProvider,
+		@ITelemetryService private readonly _telemetryService: ITelemetryService,
 	) { }
 
 	async getPreferredEditTool(model: LanguageModelChat): Promise<EditTools[] | undefined> {
@@ -61,7 +64,16 @@ export class EditToolLearningService implements IEditToolLearningService {
 			return undefined;
 		}
 
-		const hardcoded = this._getHardcodedPreferences(endpoint.model);
+		const fromEndpoint = endpoint.supportedEditTools
+			?.map(e => byokEditToolNamesToToolNames.hasOwnProperty(e) ? byokEditToolNamesToToolNames[e] : undefined)
+			.filter(isDefined);
+		if (fromEndpoint?.length) {
+			return fromEndpoint;
+		}
+
+		// Note: looking at the 'name' rather than 'model' is intentional, 'model' is the user-
+		// provided model ID whereas the 'name' is the name of the model on the BYOK provider.
+		const hardcoded = this._getHardcodedPreferences(endpoint.name);
 		if (hardcoded) {
 			return hardcoded;
 		}
@@ -78,7 +90,7 @@ export class EditToolLearningService implements IEditToolLearningService {
 		}
 
 		const learningData = this._getModelLearningData(model.id);
-		this._recordEdit(learningData, tool, success);
+		this._recordEdit(model.id, learningData, tool, success);
 		await this._saveModelLearningData(model.id, learningData);
 	}
 
@@ -100,25 +112,45 @@ export class EditToolLearningService implements IEditToolLearningService {
 		return EDIT_TOOL_LEARNING_STATES[data.state].allowedTools;
 	}
 
-	private _checkStateTransitions(data: IEditToolLearningData): State {
+	private _checkStateTransitions(modelId: string, data: IEditToolLearningData): State {
 		const currentConfig = EDIT_TOOL_LEARNING_STATES[data.state];
 
+		if (!currentConfig.transitions) {
+			return data.state;
+		}
+
 		for (const [targetState, condition] of Object.entries(currentConfig.transitions)) {
-			if (condition(data)) {
-				return Number(targetState) as State;
+			if (!condition(data)) {
+				continue;
 			}
+
+			const target = Number(targetState) as State;
+
+			/* __GDPR__
+				"editToolLearning.transition" : {
+					"owner": "connor4312",
+					"comment": "Tracks state transitions in the edit tool learning system.",
+					"modelId": { "classification": "SystemMetaData", "purpose": "PerformanceAndHealth", "comment": "Model ID" },
+					"state": { "classification": "SystemMetaData", "purpose": "PerformanceAndHealth", "comment": "State the model transitioned to", "isMeasurement": true }
+				}
+			*/
+			this._telemetryService.sendMSFTTelemetryEvent('editToolLearning.transition', { modelId }, {
+				state: target,
+			});
+
+			return target;
 		}
 
 		return data.state; // No transition
 	}
 
-	private _recordEdit(data: IEditToolLearningData, tool: EditTools, success: boolean): void {
+	private _recordEdit(modelId: string, data: IEditToolLearningData, tool: EditTools, success: boolean): void {
 		const successBit = success ? 1n : 0n;
 		const toolData = (data.tools[tool] ??= { successBitset: 0n, attempts: 0 });
 		toolData.successBitset = addToWindow(toolData.successBitset, successBit);
 		toolData.attempts++;
 
-		const newState = this._checkStateTransitions(data);
+		const newState = this._checkStateTransitions(modelId, data);
 		if (newState !== data.state) {
 			data.state = newState;
 			data.tools = {};

src/extension/tools/common/editToolLearningStates.ts

@@ -42,7 +42,7 @@ export const enum State {
 
 interface StateConfig {
 	allowedTools: EditTools[];
-	transitions: { [K in State]?: (data: IEditToolLearningData) => boolean };
+	transitions?: { [K in State]?: (data: IEditToolLearningData) => boolean };
 }
 
 
@@ -126,14 +126,11 @@ export const EDIT_TOOL_LEARNING_STATES: Record<State, StateConfig> = {
 	// Terminal states have no transitions
 	[State.EditFileOnly]: {
 		allowedTools: [ToolName.EditFile],
-		transitions: {},
 	},
 	[State.ReplaceStringOnly]: {
 		allowedTools: [ToolName.ReplaceString],
-		transitions: {},
 	},
 	[State.ReplaceStringWithMulti]: {
 		allowedTools: [ToolName.ReplaceString, ToolName.MultiReplaceString],
-		transitions: {},
 	},
 };

src/extension/tools/common/toolNames.ts

@@ -103,6 +103,13 @@ export enum ContributedToolName {
 	ExecuteTask = 'execute_task',
 }
 
+export const byokEditToolNamesToToolNames = {
+	'find-replace': ToolName.ReplaceString,
+	'multi-find-replace': ToolName.MultiReplaceString,
+	'apply-patch': ToolName.ApplyPatch,
+	'code-rewrite': ToolName.EditFile,
+} as const;
+
 const toolNameToContributedToolNames = new Map<ToolName, ContributedToolName>();
 const contributedToolNameToToolNames = new Map<ContributedToolName, ToolName>();
 for (const [contributedNameKey, contributedName] of Object.entries(ContributedToolName)) {

src/extension/tools/node/test/editToolLearningService.spec.ts

@@ -7,6 +7,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest';
 import type { LanguageModelChat } from 'vscode';
 import { IEndpointProvider } from '../../../../platform/endpoint/common/endpointProvider';
 import { IChatEndpoint } from '../../../../platform/networking/common/networking';
+import { NullTelemetryService } from '../../../../platform/telemetry/common/nullTelemetryService';
 import { MockExtensionContext } from '../../../../platform/test/node/extensionContext';
 import { EditToolLearningService, EditTools } from '../../common/editToolLearningService';
 import { LearningConfig } from '../../common/editToolLearningStates';
@@ -48,7 +49,7 @@ describe('EditToolLearningService', () => {
 		policy: 'enabled',
 		urlOrRequestMetadata: 'test-url',
 		modelMaxPromptTokens: 4000,
-		name: 'test-model',
+		name: model.id,
 		version: '1.0',
 		tokenizer: 'gpt',
 		acceptChatPolicy: vi.fn().mockResolvedValue(true),
@@ -76,7 +77,7 @@ describe('EditToolLearningService', () => {
 		// Set up proper spies for global state methods
 		mockContext.globalState.get = vi.fn().mockReturnValue(undefined);
 		mockContext.globalState.update = vi.fn().mockResolvedValue(undefined);
-		service = new EditToolLearningService(mockContext as any, mockEndpointProvider);
+		service = new EditToolLearningService(mockContext as any, mockEndpointProvider, new NullTelemetryService());
 	});
 
 	describe('getPreferredEditTool', () => {
@@ -271,7 +272,7 @@ describe('EditToolLearningService', () => {
 			await simulateEdits(model, ToolName.ReplaceString, 10, 5);
 
 			// Create a new service instance with the same context
-			const newService = new EditToolLearningService(mockContext as any, mockEndpointProvider);
+			const newService = new EditToolLearningService(mockContext as any, mockEndpointProvider, new NullTelemetryService());
 
 			// The new service should have access to the persisted data
 			const result = await newService.getPreferredEditTool(model);

src/extension/vscode.d.ts

@@ -20454,20 +20454,24 @@ declare module 'vscode' {
 		/**
 		 * Various features that the model supports such as tool calling or image input.
 		 */
-		readonly capabilities: {
+		readonly capabilities: LanguageModelChatCapabilities;
+	}
 
-			/**
-			 * Whether image input is supported by the model.
-			 * Common supported images are jpg and png, but each model will vary in supported mimetypes.
-			 */
-			readonly imageInput?: boolean;
+	/**
+	 * Various features that the {@link LanguageModelChatInformation} supports such as tool calling or image input.
+	 */
+	export interface LanguageModelChatCapabilities {
+		/**
+		 * Whether image input is supported by the model.
+		 * Common supported images are jpg and png, but each model will vary in supported mimetypes.
+		 */
+		readonly imageInput?: boolean;
 
-			/**
-			 * Whether tool calling is supported by the model.
-			 * If a number is provided, that is the maximum number of tools that can be provided in a request to the model.
-			 */
-			readonly toolCalling?: boolean | number;
-		};
+		/**
+		 * Whether tool calling is supported by the model.
+		 * If a number is provided, that is the maximum number of tools that can be provided in a request to the model.
+		 */
+		readonly toolCalling?: boolean | number;
 	}
 
 	/**

src/extension/vscode.proposed.chatProvider.d.ts

@@ -56,6 +56,25 @@ declare module 'vscode' {
 		readonly statusIcon?: ThemeIcon;
 	}
 
+	export interface LanguageModelChatCapabilities {
+		/**
+		 * The tools the model prefers for making file edits. If not provided or if none of the tools,
+		 * are recognized, the editor will try multiple edit tools and pick the best one. The available
+		 * edit tools WILL change over time and this capability only serves as a hint to the editor.
+		 *
+		 * Edit tools currently recognized include:
+		 * - 'find-replace': Find and replace text in a document.
+		 * - 'multi-find-replace': Find and replace text in a document.
+		 * - 'apply-patch': A file-oriented diff format used by some OpenAI models
+		 * - 'code-rewrite': A general but slower editing tool that allows the model
+		 *   to rewrite and code snippet and provide only the replacement to the editor.
+		 *
+		 * The order of edit tools in this array has no significance; all of the recognized edit
+		 * tools will be made available to the model.
+		 */
+		readonly editTools?: string[];
+	}
+
 	export type LanguageModelResponsePart2 = LanguageModelResponsePart | LanguageModelDataPart | LanguageModelThinkingPart;
 
 	export interface LanguageModelChatProvider<T extends LanguageModelChatInformation = LanguageModelChatInformation> {