Instant Apply -> respect proxy's /models

src/extension/prompts/node/codeMapper/codeMapper.ts

@@ -292,8 +292,6 @@ interface ICompletedRequest {
 export class CodeMapper {
 
 	static closingXmlTag = 'copilot-edited-file';
-	private gpt4oProxyEndpoint: Promise<Proxy4oEndpoint>;
-	private shortIAEndpoint: Promise<ProxyInstantApplyShortEndpoint>;
 	private shortContextLimit: number;
 
 	constructor(
@@ -311,12 +309,19 @@ export class CodeMapper {
 		@INotebookService private readonly notebookService: INotebookService,
 		@IConfigurationService configurationService: IConfigurationService,
 	) {
-		this.gpt4oProxyEndpoint = this.experimentationService.hasTreatments().then(() => this.instantiationService.createInstance(Proxy4oEndpoint));
-		this.shortIAEndpoint = this.experimentationService.hasTreatments().then(() => this.instantiationService.createInstance(ProxyInstantApplyShortEndpoint));
-
 		this.shortContextLimit = configurationService.getExperimentBasedConfig<number>(ConfigKey.Advanced.InstantApplyShortContextLimit, experimentationService) ?? 8000;
 	}
 
+	private async getGpt4oProxyEndpoint(): Promise<Proxy4oEndpoint> {
+		await this.experimentationService.hasTreatments();
+		return this.instantiationService.createInstance(Proxy4oEndpoint);
+	}
+
+	private async getShortIAEndpoint(): Promise<ProxyInstantApplyShortEndpoint> {
+		await this.experimentationService.hasTreatments();
+		return this.instantiationService.createInstance(ProxyInstantApplyShortEndpoint);
+	}
+
 	public async mapCode(request: ICodeMapperRequestInput, resultStream: MappedEditsResponseStream, telemetryInfo: ICodeMapperTelemetryInfo | undefined, token: CancellationToken): Promise<CodeMapperOutcome | undefined> {
 
 		const fastEdit = await this.mapCodeUsingFastEdit(request, resultStream, telemetryInfo, token);
@@ -403,7 +408,7 @@ export class CodeMapper {
 	//#region Full file rewrite with speculation / predicted outputs
 
 	private async buildPrompt(request: ICodeMapperRequestInput, token: CancellationToken): Promise<IFullRewritePrompt> {
-		let endpoint: ChatEndpoint = await this.gpt4oProxyEndpoint;
+		let endpoint: ChatEndpoint = await this.getGpt4oProxyEndpoint();
 		const tokenizer = this.tokenizerProvider.acquireTokenizer(endpoint);
 		const requestId = generateUuid();
 
@@ -440,7 +445,7 @@ export class CodeMapper {
 		}, '').trimEnd() + `\n\n\nThe resulting document:\n<${CodeMapper.closingXmlTag}>\n${fence}${languageIdToMDCodeBlockLang(languageId)}\n`;
 
 		if (prompt.length < this.shortContextLimit) {
-			endpoint = await this.shortIAEndpoint;
+			endpoint = await this.getShortIAEndpoint();
 		}
 
 		const promptTokenCount = await tokenizer.tokenLength(prompt);

src/extension/tools/node/userPreferencesTool.tsx

@@ -57,7 +57,6 @@ class UserPreferenceUpdatePrompt extends PromptElement<IUserPreferenceUpdateProm
 class UpdateUserPreferencesTool implements ICopilotTool<IUpdateUserPreferencesToolParams> {
 
 	public static readonly toolName = ToolName.UpdateUserPreferences;
-	private readonly endpoint = this.instantiationService.createInstance(Proxy4oEndpoint);
 
 	constructor(
 		@IVSCodeExtensionContext private readonly extensionContext: IVSCodeExtensionContext,
@@ -66,6 +65,10 @@ class UpdateUserPreferencesTool implements ICopilotTool<IUpdateUserPreferencesTo
 	) {
 	}
 
+	private getEndpoint(): Proxy4oEndpoint {
+		return this.instantiationService.createInstance(Proxy4oEndpoint);
+	}
+
 	private get userPreferenceFile(): URI {
 		return URI.joinPath(this.extensionContext.globalStorageUri, 'copilotUserPreferences.md');
 	}
@@ -87,9 +90,9 @@ class UpdateUserPreferencesTool implements ICopilotTool<IUpdateUserPreferencesTo
 	}
 
 	private async generateNewContent(currentContent: string, facts: string[], token: CancellationToken): Promise<string> {
-
-		const { messages } = await renderPromptElement(this.instantiationService, this.endpoint, UserPreferenceUpdatePrompt, { facts: facts, currentContent, userPreferenceFile: this.userPreferenceFile }, undefined, token);
-		return this.doFetch(messages, this.endpoint, currentContent, token);
+		const endpoint = this.getEndpoint();
+		const { messages } = await renderPromptElement(this.instantiationService, endpoint, UserPreferenceUpdatePrompt, { facts: facts, currentContent, userPreferenceFile: this.userPreferenceFile }, undefined, token);
+		return this.doFetch(messages, endpoint, currentContent, token);
 	}
 
 	private async doFetch(promptMessages: Raw.ChatMessage[], endpoint: IChatEndpoint, speculation: string, token: CancellationToken) {

src/platform/configuration/common/configurationService.ts

@@ -788,6 +788,7 @@ export namespace ConfigKey {
 		export const InlineEditsJointCompletionsProviderStrategy = defineTeamInternalSetting<JointCompletionsProviderStrategy>('chat.advanced.inlineEdits.jointCompletionsProvider.strategy', ConfigType.ExperimentBased, JointCompletionsProviderStrategy.Regular);
 		export const InlineEditsJointCompletionsProviderTriggerChangeStrategy = defineTeamInternalSetting<JointCompletionsProviderTriggerChangeStrategy>('chat.advanced.inlineEdits.jointCompletionsProvider.triggerChangeStrategy', ConfigType.ExperimentBased, JointCompletionsProviderTriggerChangeStrategy.NoTriggerOnCompletionsRequestInFlight);
 		export const InstantApplyModelName = defineTeamInternalSetting<string>('chat.advanced.instantApply.modelName', ConfigType.ExperimentBased, 'gpt-4o-instant-apply-full-ft-v66');
+		export const UseProxyModelsServiceForInstantApply = defineTeamInternalSetting<boolean>('chat.advanced.instantApply.useProxyModelsService', ConfigType.ExperimentBased, false);
 		export const VerifyTextDocumentChanges = defineTeamInternalSetting<boolean>('chat.advanced.inlineEdits.verifyTextDocumentChanges', ConfigType.ExperimentBased, false);
 
 		// TODO: @sandy081 - These should be moved away from this namespace

src/platform/endpoint/node/proxy4oEndpoint.ts

@@ -12,12 +12,14 @@ import { IChatMLFetcher } from '../../chat/common/chatMLFetcher';
 import { CHAT_MODEL, ConfigKey, IConfigurationService } from '../../configuration/common/configurationService';
 import { ILogService } from '../../log/common/logService';
 import { IFetcherService } from '../../networking/common/fetcherService';
+import { IProxyModelsService } from '../../proxyModels/common/proxyModelsService';
 import { IExperimentationService } from '../../telemetry/common/nullExperimentationService';
 import { ITokenizerProvider } from '../../tokenizer/node/tokenizer';
 import { ICAPIClientService } from '../common/capiClient';
 import { IDomainService } from '../common/domainService';
 import { IChatModelInformation } from '../common/endpointProvider';
 import { ChatEndpoint } from './chatEndpoint';
+import { getInstantApplyModel } from './proxyModelHelper';
 
 export class Proxy4oEndpoint extends ChatEndpoint {
 
@@ -35,8 +37,15 @@ export class Proxy4oEndpoint extends ChatEndpoint {
 		@IConfigurationService configurationService: IConfigurationService,
 		@IExperimentationService experimentationService: IExperimentationService,
 		@ILogService logService: ILogService,
+		@IProxyModelsService proxyModelsService: IProxyModelsService,
 	) {
-		const model = configurationService.getExperimentBasedConfig<string>(ConfigKey.TeamInternal.InstantApplyModelName, experimentationService) ?? CHAT_MODEL.GPT4OPROXY;
+		const model = getInstantApplyModel(
+			configurationService,
+			experimentationService,
+			proxyModelsService,
+			ConfigKey.TeamInternal.InstantApplyModelName,
+			CHAT_MODEL.GPT4OPROXY
+		);
 
 		const modelInfo: IChatModelInformation = {
 			id: model,

src/platform/endpoint/node/proxyInstantApplyShortEndpoint.ts

@@ -11,13 +11,15 @@ import { IChatMLFetcher } from '../../chat/common/chatMLFetcher';
 import { CHAT_MODEL, ConfigKey, IConfigurationService } from '../../configuration/common/configurationService';
 import { ILogService } from '../../log/common/logService';
 import { IFetcherService } from '../../networking/common/fetcherService';
+import { IProxyModelsService } from '../../proxyModels/common/proxyModelsService';
 import { IExperimentationService } from '../../telemetry/common/nullExperimentationService';
 import { ITelemetryService } from '../../telemetry/common/telemetry';
 import { ITokenizerProvider } from '../../tokenizer/node/tokenizer';
 import { ICAPIClientService } from '../common/capiClient';
 import { IDomainService } from '../common/domainService';
 import { IChatModelInformation } from '../common/endpointProvider';
 import { ChatEndpoint } from './chatEndpoint';
+import { getInstantApplyModel } from './proxyModelHelper';
 
 export class ProxyInstantApplyShortEndpoint extends ChatEndpoint {
 
@@ -33,8 +35,15 @@ export class ProxyInstantApplyShortEndpoint extends ChatEndpoint {
 		@IConfigurationService configurationService: IConfigurationService,
 		@IExperimentationService experimentationService: IExperimentationService,
 		@ILogService logService: ILogService,
+		@IProxyModelsService proxyModelsService: IProxyModelsService,
 	) {
-		const model = configurationService.getExperimentBasedConfig<string>(ConfigKey.Advanced.InstantApplyShortModelName, experimentationService) ?? CHAT_MODEL.SHORT_INSTANT_APPLY;
+		const model = getInstantApplyModel(
+			configurationService,
+			experimentationService,
+			proxyModelsService,
+			ConfigKey.Advanced.InstantApplyShortModelName,
+			CHAT_MODEL.SHORT_INSTANT_APPLY
+		);
 		const modelInfo: IChatModelInformation = {
 			id: model,
 			name: model,

src/platform/endpoint/node/proxyModelHelper.ts

@@ -0,0 +1,39 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { ConfigKey, IConfigurationService } from '../../configuration/common/configurationService';
+import { IProxyModelsService } from '../../proxyModels/common/proxyModelsService';
+import { IExperimentationService } from '../../telemetry/common/nullExperimentationService';
+
+/**
+ * Determines which model to use for instant apply endpoints based on proxy models service availability
+ * and configuration settings.
+ *
+ * @param configurationService - Service for accessing configuration values
+ * @param experimentationService - Service for accessing experiment flags
+ * @param proxyModelsService - Service providing proxy model information
+ * @param fallbackConfigKey - Configuration key to use if proxy models service is not enabled
+ * @param defaultModel - Default model to use if no configuration is found
+ * @returns The model name to use for the endpoint
+ */
+export function getInstantApplyModel(
+	configurationService: IConfigurationService,
+	experimentationService: IExperimentationService,
+	proxyModelsService: IProxyModelsService,
+	fallbackConfigKey: typeof ConfigKey.Advanced.InstantApplyShortModelName | typeof ConfigKey.TeamInternal.InstantApplyModelName,
+	defaultModel: string
+): string {
+	// Check experimental flag to determine if we should use proxy models service
+	const useProxyModelsService = configurationService.getExperimentBasedConfig<boolean>(
+		ConfigKey.TeamInternal.UseProxyModelsServiceForInstantApply,
+		experimentationService
+	);
+
+	const instantApplyModels = useProxyModelsService ? proxyModelsService.instantApplyModels : undefined;
+
+	return (instantApplyModels && instantApplyModels.length > 0)
+		? instantApplyModels[0].name
+		: configurationService.getExperimentBasedConfig<string>(fallbackConfigKey, experimentationService) ?? defaultModel;
+}