Add EntraId authentication support for Azure specific endpoints (#2070)

* Add EntraId authentication support for Azure specific endpoints

* fix from copilot review

* changes from feedback

* update error msg

* fixes based on comments

* fix typing for tests

* fixing tests

* comment feedback

extensions/copilot/package.json

@@ -3212,6 +3212,22 @@
 						],
 						"markdownDescription": "%github.copilot.config.virtualTools.threshold%"
 					},
+					"github.copilot.chat.azureAuthType": {
+						"type": "string",
+						"enum": [
+							"entraId",
+							"apiKey"
+						],
+						"enumDescriptions": [
+							"%github.copilot.config.azureAuthType.entraId%",
+							"%github.copilot.config.azureAuthType.apiKey%"
+						],
+						"default": "entraId",
+						"tags": [
+							"experimental"
+						],
+						"markdownDescription": "%github.copilot.config.azureAuthType%"
+					},
 					"github.copilot.chat.azureModels": {
 						"type": "object",
 						"default": {},

extensions/copilot/package.nls.json

@@ -167,6 +167,9 @@
 	"github.copilot.config.codeGeneration.useInstructionFiles": "Controls whether code instructions from `.github/copilot-instructions.md` are added to Copilot requests.\n\nNote: Keep your instructions short and precise. Poor instructions can degrade Copilot's quality and performance. [Learn more](https://aka.ms/github-copilot-custom-instructions) about customizing Copilot.",
 	"github.copilot.config.codeGeneration.instruction.text": "A text instruction that will be added to Copilot requests that generate code. Optionally, you can specify a language for the instruction.",
 	"github.copilot.config.codeGeneration.instruction.file": "A path to a file that will be added to Copilot requests that generate code. Optionally, you can specify a language for the instruction.",
+	"github.copilot.config.azureAuthType": "Authentication method for Azure OpenAI models. Entra ID is recommended for enterprise security and uses your Azure credentials.",
+	"github.copilot.config.azureAuthType.entraId": "Use Entra ID (Azure AD) authentication with your Microsoft account credentials",
+	"github.copilot.config.azureAuthType.apiKey": "Use API key authentication. Not recommended as this is less secure than token based authentication.",
 	"github.copilot.config.testGeneration.instructions": "A set of instructions that will be added to Copilot requests that generate tests.\nInstructions can come from: \n- a file in the workspace: `{ \"file\": \"fileName\" }`\n- text in natural language: `{ \"text\": \"Use underscore for field names.\" }`\n\nNote: Keep your instructions short and precise. Poor instructions can degrade Copilot's quality and performance.",
 	"github.copilot.config.testGeneration.instructions.deprecated": "Use instructions files instead. See https://aka.ms/vscode-ghcp-custom-instructions for more information.",
 	"github.copilot.config.experimental.testGeneration.instruction.text": "A text instruction that will be added to Copilot requests that generate tests. Optionally, you can specify a language for the instruction.",

extensions/copilot/src/extension/byok/node/azureOpenAIEndpoint.ts

@@ -0,0 +1,23 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+import { OpenAIEndpoint } from './openAIEndpoint';
+
+/**
+ * Azure-specific OpenAI endpoint that supports Entra ID authentication.
+ * Extends OpenAIEndpoint to override header generation for Azure-specific auth methods.
+ * Note: Authentication token refresh is handled at the provider level (azureProvider.ts).
+ */
+export class AzureOpenAIEndpoint extends OpenAIEndpoint {
+	/**
+	 * Override to use Entra ID authentication headers instead of API key.
+	 */
+	public override getExtraHeaders(): Record<string, string> {
+		const headers = super.getExtraHeaders();
+		headers['Authorization'] = `Bearer ${this._apiKey}`;
+		// Defensive: Ensure 'api-key' header is never sent for Azure endpoints, even if parent class changes.
+		delete headers['api-key'];
+		return headers;
+	}
+}

extensions/copilot/src/extension/byok/node/test/azureOpenAIEndpoint.spec.ts

@@ -0,0 +1,160 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { afterEach, beforeEach, describe, expect, it } from 'vitest';
+import { IChatModelInformation, ModelSupportedEndpoint } from '../../../../platform/endpoint/common/endpointProvider';
+import { ITestingServicesAccessor } from '../../../../platform/test/node/services';
+import { TokenizerType } from '../../../../util/common/tokenizer';
+import { DisposableStore } from '../../../../util/vs/base/common/lifecycle';
+import { IInstantiationService } from '../../../../util/vs/platform/instantiation/common/instantiation';
+import { createExtensionUnitTestingServices } from '../../../test/node/services';
+import { AzureOpenAIEndpoint } from '../azureOpenAIEndpoint';
+
+describe('AzureOpenAIEndpoint', () => {
+	let modelMetadata: IChatModelInformation;
+	const disposables = new DisposableStore();
+	let accessor: ITestingServicesAccessor;
+	let instaService: IInstantiationService;
+
+	beforeEach(() => {
+		modelMetadata = {
+			id: 'test-azure-model',
+			name: 'Test Azure Model',
+			version: '1.0',
+			model_picker_enabled: true,
+			is_chat_default: false,
+			is_chat_fallback: false,
+			supported_endpoints: [ModelSupportedEndpoint.ChatCompletions],
+			capabilities: {
+				type: 'chat',
+				family: 'openai',
+				tokenizer: TokenizerType.O200K,
+				supports: {
+					parallel_tool_calls: false,
+					streaming: true,
+					tool_calls: false,
+					vision: false,
+					prediction: false,
+					thinking: false
+				},
+				limits: {
+					max_prompt_tokens: 128000,
+					max_output_tokens: 4096,
+					max_context_window_tokens: 132096
+				}
+			}
+		};
+
+		const testingServiceCollection = createExtensionUnitTestingServices();
+		accessor = disposables.add(testingServiceCollection.createTestingAccessor());
+		instaService = accessor.get(IInstantiationService);
+	});
+
+	afterEach(() => {
+		disposables.clear();
+	});
+
+	describe('getExtraHeaders', () => {
+		it('should use Authorization header with Bearer token for Entra ID authentication', () => {
+			const entraToken = 'test-entra-token-abc123';
+			const endpoint = instaService.createInstance(
+				AzureOpenAIEndpoint,
+				modelMetadata,
+				entraToken,
+				'https://example-endpoint.example.com/v1/chat/completions'
+			);
+			const headers = endpoint.getExtraHeaders();
+
+			// Should have Authorization header with Bearer token
+			expect(headers['Authorization']).toBe(`Bearer ${entraToken}`);
+
+			// Should NOT have api-key header (Azure API key auth)
+			expect(headers['api-key']).toBeUndefined();
+
+			// Should have standard headers
+			expect(headers['Content-Type']).toBe('application/json');
+		});
+
+		it('should override parent class headers to replace api-key with Authorization', () => {
+			const entraToken = 'test-entra-token-xyz789';
+			const endpoint = instaService.createInstance(
+				AzureOpenAIEndpoint,
+				modelMetadata,
+				entraToken,
+				'https://example-endpoint.example.com/v1/chat/completions'
+			);
+			const headers = endpoint.getExtraHeaders();
+
+			// Verify the override worked correctly
+			expect(headers['Authorization']).toBe(`Bearer ${entraToken}`);
+			expect(headers['api-key']).toBeUndefined();
+			expect(Object.keys(headers)).not.toContain('api-key');
+		});
+
+		it('should work with different Azure OpenAI endpoint URLs', () => {
+			const entraToken = 'test-token-456';
+
+			// Test with different endpoint formats
+			const urls = [
+				'https://example-endpoint-1.example.com/v1/chat/completions',
+				'https://example-endpoint-2.example.com/v1/chat/completions',
+				'https://example-endpoint-3.example.com/v1/chat/completions'
+			];
+
+			for (const url of urls) {
+				const endpoint = instaService.createInstance(
+					AzureOpenAIEndpoint,
+					modelMetadata,
+					entraToken,
+					url
+				);
+
+				const headers = endpoint.getExtraHeaders();
+
+				expect(headers['Authorization']).toBe(`Bearer ${entraToken}`);
+				expect(headers['api-key']).toBeUndefined();
+			}
+		});
+
+		it('should preserve other headers from parent class', () => {
+			const entraToken = 'test-token-789';
+			const endpoint = instaService.createInstance(
+				AzureOpenAIEndpoint,
+				modelMetadata,
+				entraToken,
+				'https://example-endpoint.example.com/v1/chat/completions'
+			);
+
+			const headers = endpoint.getExtraHeaders();
+
+			// Should preserve Content-Type from parent
+			expect(headers['Content-Type']).toBe('application/json');
+
+			// Should have Authorization header
+			expect(headers['Authorization']).toBeDefined();
+			expect(headers['Authorization']).toContain('Bearer');
+		});
+	});
+
+	describe('inheritance', () => {
+		it('should inherit from OpenAIEndpoint and maintain same constructor signature', () => {
+			const entraToken = 'test-token-inheritance';
+
+			// Should be able to instantiate with same parameters as OpenAIEndpoint
+			const endpoint = instaService.createInstance(
+				AzureOpenAIEndpoint,
+				modelMetadata,
+				entraToken,
+				'https://example-endpoint.example.com/v1/chat/completions'
+			);
+
+			// Should be an instance of AzureOpenAIEndpoint
+			expect(endpoint).toBeInstanceOf(AzureOpenAIEndpoint);
+
+			// Should have getExtraHeaders method
+			expect(typeof endpoint.getExtraHeaders).toBe('function');
+		});
+	});
+});

extensions/copilot/src/extension/byok/vscode-node/azureProvider.ts

@@ -3,12 +3,17 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService';
+import * as vscode from 'vscode';
+import { CancellationToken, LanguageModelChatMessage, LanguageModelChatMessage2, LanguageModelResponsePart2, Progress, ProvideLanguageModelChatResponseOptions } from 'vscode';
+import { AzureAuthMode, ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService';
+import { 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';
+import { BYOKKnownModels } from '../common/byokProvider';
+import { AzureOpenAIEndpoint } from '../node/azureOpenAIEndpoint';
 import { IBYOKStorageService } from './byokStorageService';
-import { CustomOAIBYOKModelProvider, hasExplicitApiPath } from './customOAIProvider';
+import { CustomOAIBYOKModelProvider, CustomOAIModelInfo, hasExplicitApiPath } from './customOAIProvider';
 
 export function resolveAzureUrl(modelId: string, url: string): string {
 	// The fully resolved url was already passed in
@@ -65,4 +70,86 @@ export class AzureBYOKModelProvider extends CustomOAIBYOKModelProvider {
 	protected override resolveUrl(modelId: string, url: string): string {
 		return resolveAzureUrl(modelId, url);
 	}
+
+	protected override async getModelsWithCredentials(silent: boolean): Promise<BYOKKnownModels> {
+		// Check user's authentication preference from settings github.copilot.chat.azureAuthType (default: AzureAuthMode.EntraId)
+		const authType = this._configurationService.getConfig(ConfigKey.AzureAuthType);
+
+		if (authType === AzureAuthMode.EntraId) {
+			// Pre-authenticate during model enumeration (not when sending message)
+			// This mirrors API key behavior where user is prompted during enumeration
+			if (!silent) {
+				try {
+					await vscode.authentication.getSession(
+						AzureAuthMode.MICROSOFT_AUTH_PROVIDER,
+						[AzureAuthMode.COGNITIVE_SERVICES_SCOPE],
+						{ createIfNone: true }
+					);
+				} catch (error) {
+					// If sign-in fails, don't show models in picker
+					this._logService.error('[AzureBYOKModelProvider] Authentication failed during Entra ID sign-in:', error);
+					return {};
+				}
+			}
+
+			// Return all configured models (no API key check needed for Entra ID)
+			return this.getAllModels();
+		} else {
+			// API KEY MODE: Use traditional API key authentication
+			return super.getModelsWithCredentials(silent);
+		}
+	}
+
+	override async provideLanguageModelChatResponse(
+		model: CustomOAIModelInfo,
+		messages: Array<LanguageModelChatMessage | LanguageModelChatMessage2>,
+		options: ProvideLanguageModelChatResponseOptions,
+		progress: Progress<LanguageModelResponsePart2>,
+		token: CancellationToken
+	): Promise<void> {
+		const authType = this._configurationService.getConfig(ConfigKey.AzureAuthType);
+
+		if (authType === AzureAuthMode.EntraId) {
+			// Session is guaranteed to be defined when createIfNone: true
+			const session: vscode.AuthenticationSession = await vscode.authentication.getSession(
+				AzureAuthMode.MICROSOFT_AUTH_PROVIDER,
+				[AzureAuthMode.COGNITIVE_SERVICES_SCOPE],
+				{
+					createIfNone: true,
+					silent: false
+				}
+			);
+
+			const modelInfo = await this.getModelInfo(model.id, undefined, {
+				maxInputTokens: model.maxInputTokens,
+				maxOutputTokens: model.maxOutputTokens,
+				toolCalling: !!model.capabilities?.toolCalling,
+				vision: !!model.capabilities?.imageInput,
+				name: model.name,
+				url: model.url,
+				thinking: model.thinking,
+				editTools: model.capabilities?.editTools?.filter(isEndpointEditToolName),
+				requestHeaders: model.requestHeaders,
+			});
+
+			const openAIChatEndpoint = this._instantiationService.createInstance(
+				AzureOpenAIEndpoint,
+				modelInfo,
+				session.accessToken,  // Pass Entra ID token
+				model.url
+			);
+
+			return this._lmWrapper.provideLanguageModelResponse(
+				openAIChatEndpoint,
+				messages,
+				options,
+				options.requestInitiator,
+				progress,
+				token
+			);
+		} else {
+			// API KEY AUTHENTICATION FLOW using parent logic
+			return super.provideLanguageModelChatResponse(model, messages, options, progress, token);
+		}
+	}
 }

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

@@ -44,7 +44,7 @@ export function hasExplicitApiPath(url: string): boolean {
 	return url.includes('/responses') || url.includes('/chat/completions');
 }
 
-interface CustomOAIModelInfo extends LanguageModelChatInformation {
+export interface CustomOAIModelInfo extends LanguageModelChatInformation {
 	url: string;
 	thinking: boolean;
 	requestHeaders?: Record<string, string>;
@@ -96,7 +96,7 @@ export class CustomOAIBYOKModelProvider implements BYOKModelProvider<CustomOAIMo
 		return userModelConfig[modelId]?.requiresAPIKey !== false;
 	}
 
-	private async getAllModels(): Promise<BYOKKnownModels> {
+	protected async getAllModels(): Promise<BYOKKnownModels> {
 		const modelConfig = this.getUserModelConfig();
 		const models: BYOKKnownModels = {};
 
@@ -119,7 +119,7 @@ export class CustomOAIBYOKModelProvider implements BYOKModelProvider<CustomOAIMo
 		return models;
 	}
 
-	private async getModelsWithAPIKeys(silent: boolean): Promise<BYOKKnownModels> {
+	protected async getModelsWithCredentials(silent: boolean): Promise<BYOKKnownModels> {
 		const models = await this.getAllModels();
 		const modelsWithApiKeys: BYOKKnownModels = {};
 		for (const [modelId, modelInfo] of Object.entries(models)) {
@@ -166,10 +166,10 @@ export class CustomOAIBYOKModelProvider implements BYOKModelProvider<CustomOAIMo
 
 	async provideLanguageModelChatInformation(options: { silent: boolean }, token: CancellationToken): Promise<CustomOAIModelInfo[]> {
 		try {
-			let knownModels = await this.getModelsWithAPIKeys(options.silent);
+			let knownModels = await this.getModelsWithCredentials(options.silent);
 			if (Object.keys(knownModels).length === 0 && !options.silent) {
 				await new CustomOAIModelConfigurator(this._configurationService, this.providerName.toLowerCase(), this).configure(true);
-				knownModels = await this.getModelsWithAPIKeys(options.silent);
+				knownModels = await this.getModelsWithCredentials(options.silent);
 			}
 			return Object.entries(knownModels).map(([id, capabilities]) => {
 				return this.createModelInfo(id, capabilities);

extensions/copilot/src/extension/byok/vscode-node/test/azureProvider.spec.ts

@@ -0,0 +1,418 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+import * as vscode from 'vscode';
+import { BlockedExtensionService, IBlockedExtensionService } from '../../../../platform/chat/common/blockedExtensionService';
+import { AzureAuthMode, ConfigKey, IConfigurationService } from '../../../../platform/configuration/common/configurationService';
+import { ITestingServicesAccessor } from '../../../../platform/test/node/services';
+import { DisposableStore } from '../../../../util/vs/base/common/lifecycle';
+import { SyncDescriptor } from '../../../../util/vs/platform/instantiation/common/descriptors';
+import { IInstantiationService } from '../../../../util/vs/platform/instantiation/common/instantiation';
+import { createExtensionUnitTestingServices } from '../../../test/node/services';
+import { BYOKKnownModels } from '../../common/byokProvider';
+import { AzureBYOKModelProvider, resolveAzureUrl } from '../azureProvider';
+import { IBYOKStorageService } from '../byokStorageService';
+import { CustomOAIModelInfo } from '../customOAIProvider';
+
+describe('AzureBYOKModelProvider', () => {
+	const disposables = new DisposableStore();
+	let accessor: ITestingServicesAccessor;
+	let instaService: IInstantiationService;
+	let provider: AzureBYOKModelProvider;
+	let mockByokStorageService: IBYOKStorageService;
+
+	beforeEach(() => {
+		const testingServiceCollection = createExtensionUnitTestingServices();
+
+		// Add IBlockedExtensionService which is required by CopilotLanguageModelWrapper
+		testingServiceCollection.define(IBlockedExtensionService, new SyncDescriptor(BlockedExtensionService));
+
+		accessor = disposables.add(testingServiceCollection.createTestingAccessor());
+		instaService = accessor.get(IInstantiationService);
+
+		// Create mock storage service
+		mockByokStorageService = {
+			getAPIKey: vi.fn().mockResolvedValue(undefined),
+			storeAPIKey: vi.fn().mockResolvedValue(undefined),
+			deleteAPIKey: vi.fn().mockResolvedValue(undefined),
+			getStoredModelConfigs: vi.fn().mockResolvedValue({}),
+			saveModelConfig: vi.fn().mockResolvedValue(undefined),
+			removeModelConfig: vi.fn().mockResolvedValue(undefined)
+		};
+	});
+
+	afterEach(() => {
+		disposables.clear();
+		vi.restoreAllMocks();
+	});
+
+	describe('resolveAzureUrl', () => {
+		it('should handle Azure AI Foundry (models.ai.azure.com) URLs', () => {
+			const url = 'https://my-endpoint.models.ai.azure.com';
+			const result = resolveAzureUrl('gpt-4', url);
+			expect(result).toBe('https://my-endpoint.models.ai.azure.com/v1/chat/completions');
+		});
+
+		it('should handle Azure ML (inference.ml.azure.com) URLs', () => {
+			const url = 'https://my-endpoint.inference.ml.azure.com';
+			const result = resolveAzureUrl('gpt-4', url);
+			expect(result).toBe('https://my-endpoint.inference.ml.azure.com/v1/chat/completions');
+		});
+
+		it('should handle Azure OpenAI (openai.azure.com) URLs with deployment name', () => {
+			const url = 'https://my-resource.openai.azure.com';
+			const result = resolveAzureUrl('gpt-4-deployment', url);
+			expect(result).toBe('https://my-resource.openai.azure.com/openai/deployments/gpt-4-deployment/chat/completions?api-version=2025-01-01-preview');
+		});
+
+		it('should return URL unchanged if it already has explicit API path', () => {
+			const url = 'https://my-endpoint.example.com/v1/chat/completions';
+			const result = resolveAzureUrl('gpt-4', url);
+			expect(result).toBe(url);
+		});
+
+		it('should remove trailing slash before processing', () => {
+			const url = 'https://my-endpoint.models.ai.azure.com/';
+			const result = resolveAzureUrl('gpt-4', url);
+			expect(result).toBe('https://my-endpoint.models.ai.azure.com/v1/chat/completions');
+		});
+
+		it('should remove /v1 suffix before processing', () => {
+			const url = 'https://my-endpoint.models.ai.azure.com/v1';
+			const result = resolveAzureUrl('gpt-4', url);
+			expect(result).toBe('https://my-endpoint.models.ai.azure.com/v1/chat/completions');
+		});
+
+		it('should throw error for unrecognized Azure URL', () => {
+			const url = 'https://unknown.example.com';
+			expect(() => resolveAzureUrl('gpt-4', url)).toThrow('Unrecognized Azure deployment URL');
+		});
+	});
+
+	describe('getModelsWithCredentials - Entra ID mode', () => {
+		beforeEach(() => {
+			const configService = accessor.get(IConfigurationService);
+			vi.spyOn(configService, 'getConfig').mockImplementation((key: any) => {
+				if (key === ConfigKey.AzureAuthType) {
+					return AzureAuthMode.EntraId;
+				}
+				if (key === ConfigKey.AzureModels) {
+					return {
+						'gpt-4': {
+							name: 'GPT-4',
+							url: 'https://test.models.ai.azure.com',
+							toolCalling: true,
+							vision: false,
+							maxInputTokens: 128000,
+							maxOutputTokens: 4096
+						},
+						'gpt-35-turbo': {
+							name: 'GPT-3.5 Turbo',
+							url: 'https://test.openai.azure.com',
+							toolCalling: true,
+							vision: false,
+							maxInputTokens: 16000,
+							maxOutputTokens: 4096
+						}
+					};
+				}
+				return undefined;
+			});
+
+			provider = instaService.createInstance(AzureBYOKModelProvider, mockByokStorageService);
+		});
+
+		it('should return all models without prompting authentication in silent mode', async () => {
+			const getSessionSpy = vi.spyOn(vscode.authentication, 'getSession');
+
+			const models = await provider['getModelsWithCredentials'](true);
+
+			expect(getSessionSpy).not.toHaveBeenCalled();
+			expect(Object.keys(models)).toHaveLength(2);
+			expect(models['gpt-4']).toBeDefined();
+			expect(models['gpt-35-turbo']).toBeDefined();
+		});
+
+		it('should return empty object when authentication fails in non-silent mode', async () => {
+			const authError = new Error('User canceled authentication');
+			vi.spyOn(vscode.authentication, 'getSession').mockRejectedValue(authError);
+
+			const models = await provider['getModelsWithCredentials'](false);
+
+			expect(models).toEqual({});
+			expect(vscode.authentication.getSession).toHaveBeenCalledWith(
+				AzureAuthMode.MICROSOFT_AUTH_PROVIDER,
+				[AzureAuthMode.COGNITIVE_SERVICES_SCOPE],
+				{ createIfNone: true }
+			);
+		});
+
+		it('should use enum constants instead of magic strings for auth provider and scope', async () => {
+			const mockSession = { accessToken: 'test-token', account: { id: 'test', label: 'test' }, scopes: [], id: 'test' };
+			const getSessionSpy = vi.spyOn(vscode.authentication, 'getSession').mockResolvedValue(mockSession);
+
+			await provider['getModelsWithCredentials'](false);
+
+			const callArgs = getSessionSpy.mock.calls[0];
+			expect(callArgs[0]).toBe('microsoft'); // AzureAuthMode.MICROSOFT_AUTH_PROVIDER
+			expect(callArgs[1]).toEqual(['https://cognitiveservices.azure.com/.default']); // AzureAuthMode.COGNITIVE_SERVICES_SCOPE
+		});
+	});
+
+	describe('getModelsWithCredentials - API Key mode', () => {
+		beforeEach(() => {
+			const configService = accessor.get(IConfigurationService);
+			vi.spyOn(configService, 'getConfig').mockImplementation((key: any) => {
+				if (key === ConfigKey.AzureAuthType) {
+					return AzureAuthMode.ApiKey;
+				}
+				if (key === ConfigKey.AzureModels) {
+					return {
+						'gpt-4': {
+							name: 'GPT-4',
+							url: 'https://test.openai.azure.com',
+							toolCalling: true,
+							vision: false,
+							maxInputTokens: 128000,
+							maxOutputTokens: 4096,
+							requiresAPIKey: true
+						}
+					};
+				}
+				return undefined;
+			});
+
+			provider = instaService.createInstance(AzureBYOKModelProvider, mockByokStorageService);
+		}); it('should delegate to parent class when in API Key mode', async () => {
+			// Mock the parent's getModelsWithCredentials
+			const parentProto = Object.getPrototypeOf(Object.getPrototypeOf(provider));
+			const parentGetModels = vi.spyOn(parentProto, 'getModelsWithCredentials');
+			const expectedModels: BYOKKnownModels = {
+				'gpt-4': {
+					name: 'GPT-4',
+					url: 'https://test.openai.azure.com/openai/deployments/gpt-4/chat/completions?api-version=2025-01-01-preview',
+					toolCalling: true,
+					vision: false,
+					maxInputTokens: 128000,
+					maxOutputTokens: 4096
+				}
+			};
+			parentGetModels.mockResolvedValue(expectedModels);
+
+			const getSessionSpy = vi.spyOn(vscode.authentication, 'getSession');
+
+			const models = await provider['getModelsWithCredentials'](false);
+
+			expect(getSessionSpy).not.toHaveBeenCalled();
+			expect(parentGetModels).toHaveBeenCalledWith(false);
+			expect(models).toEqual(expectedModels);
+		});
+	});
+
+	describe('provideLanguageModelChatResponse - Entra ID mode', () => {
+		let mockModel: CustomOAIModelInfo;
+		let mockMessages: vscode.LanguageModelChatMessage[];
+		let mockOptions: vscode.ProvideLanguageModelChatResponseOptions;
+		let mockProgress: vscode.Progress<vscode.LanguageModelResponsePart2>;
+		let mockToken: vscode.CancellationToken;
+
+		beforeEach(() => {
+			const configService = accessor.get(IConfigurationService);
+			vi.spyOn(configService, 'getConfig').mockImplementation((key: any) => {
+				if (key === ConfigKey.AzureAuthType) {
+					return AzureAuthMode.EntraId;
+				}
+				return undefined;
+			});
+
+			provider = instaService.createInstance(AzureBYOKModelProvider, mockByokStorageService); mockModel = {
+				id: 'gpt-4',
+				name: 'GPT-4',
+				url: 'https://test.models.ai.azure.com/v1/chat/completions',
+				detail: 'Azure',
+				version: '1.0.0',
+				maxInputTokens: 128000,
+				maxOutputTokens: 4096,
+				family: 'Azure',
+				tooltip: 'GPT-4 via Azure',
+				capabilities: {
+					toolCalling: true,
+					imageInput: false
+				},
+				thinking: false
+			};
+
+			mockMessages = [
+				new vscode.LanguageModelChatMessage(vscode.LanguageModelChatMessageRole.User, 'Hello')
+			];
+
+			mockOptions = {
+				requestInitiator: 'test-extension',
+				tools: [],
+				toolMode: vscode.LanguageModelChatToolMode.Auto
+			};
+
+			mockProgress = { report: vi.fn() };
+			mockToken = new vscode.CancellationTokenSource().token;
+		});
+
+		it('should acquire Entra ID session and use token for authentication', async () => {
+			const mockSession = {
+				accessToken: 'test-entra-token-abc123',
+				account: { id: 'test', label: 'test' },
+				scopes: [AzureAuthMode.COGNITIVE_SERVICES_SCOPE],
+				id: 'test'
+			};
+			vi.spyOn(vscode.authentication, 'getSession').mockResolvedValue(mockSession);
+
+			// Mock the language model wrapper
+			const provideResponseSpy = vi.spyOn(provider['_lmWrapper'], 'provideLanguageModelResponse').mockResolvedValue(undefined);
+
+			await provider.provideLanguageModelChatResponse(mockModel, mockMessages, mockOptions, mockProgress, mockToken);
+
+			expect(vscode.authentication.getSession).toHaveBeenCalledWith(
+				AzureAuthMode.MICROSOFT_AUTH_PROVIDER,
+				[AzureAuthMode.COGNITIVE_SERVICES_SCOPE],
+				{ createIfNone: true, silent: false }
+			);
+			expect(provideResponseSpy).toHaveBeenCalled();
+		});
+
+		it('should throw original error when authentication is rejected', async () => {
+			const authError = new Error('User did not consent to login.');
+			vi.spyOn(vscode.authentication, 'getSession').mockRejectedValue(authError);
+
+			try {
+				await provider.provideLanguageModelChatResponse(mockModel, mockMessages, mockOptions, mockProgress, mockToken);
+				expect.fail('Should have thrown error');
+			} catch (err: any) {
+				expect(err.message).toBe('User did not consent to login.');
+			}
+		});
+
+		it('should pass Entra ID token to AzureOpenAIEndpoint', async () => {
+			const mockSession = {
+				accessToken: 'test-entra-token-xyz789',
+				account: { id: 'test', label: 'test' },
+				scopes: [],
+				id: 'test'
+			};
+			vi.spyOn(vscode.authentication, 'getSession').mockResolvedValue(mockSession);
+
+			const createInstanceSpy = vi.spyOn(provider['_instantiationService'], 'createInstance');
+			vi.spyOn(provider['_lmWrapper'], 'provideLanguageModelResponse').mockResolvedValue(undefined);
+
+			await provider.provideLanguageModelChatResponse(mockModel, mockMessages, mockOptions, mockProgress, mockToken);
+
+			// Verify AzureOpenAIEndpoint was created with the token
+			expect(createInstanceSpy).toHaveBeenCalled();
+			const callArgs = createInstanceSpy.mock.calls[0];
+			expect(callArgs[1]).toBeDefined(); // modelInfo
+			expect(callArgs[2]).toBe('test-entra-token-xyz789'); // Entra ID token passed as apiKey
+			expect(callArgs[3]).toBe(mockModel.url); // URL
+		});
+
+		it('should use enum constants for auth provider and scope', async () => {
+			const mockSession = { accessToken: 'test-token', account: { id: 'test', label: 'test' }, scopes: [], id: 'test' };
+			const getSessionSpy = vi.spyOn(vscode.authentication, 'getSession').mockResolvedValue(mockSession);
+			vi.spyOn(provider['_lmWrapper'], 'provideLanguageModelResponse').mockResolvedValue(undefined);
+
+			await provider.provideLanguageModelChatResponse(mockModel, mockMessages, mockOptions, mockProgress, mockToken);
+
+			const callArgs = getSessionSpy.mock.calls[0];
+			expect(callArgs[0]).toBe('microsoft'); // AzureAuthMode.MICROSOFT_AUTH_PROVIDER
+			expect(callArgs[1]).toEqual(['https://cognitiveservices.azure.com/.default']); // AzureAuthMode.COGNITIVE_SERVICES_SCOPE
+		});
+	});
+
+	describe('provideLanguageModelChatResponse - API Key mode', () => {
+		let mockModel: CustomOAIModelInfo;
+		let mockMessages: vscode.LanguageModelChatMessage[];
+		let mockOptions: vscode.ProvideLanguageModelChatResponseOptions;
+		let mockProgress: vscode.Progress<vscode.LanguageModelResponsePart2>;
+		let mockToken: vscode.CancellationToken;
+
+		beforeEach(() => {
+			const configService = accessor.get(IConfigurationService);
+			vi.spyOn(configService, 'getConfig').mockImplementation((key: any) => {
+				if (key === ConfigKey.AzureAuthType) {
+					return AzureAuthMode.ApiKey;
+				}
+				return undefined;
+			});
+
+			provider = instaService.createInstance(AzureBYOKModelProvider, mockByokStorageService); mockModel = {
+				id: 'gpt-4',
+				name: 'GPT-4',
+				url: 'https://test.openai.azure.com',
+				detail: 'Azure',
+				version: '1.0.0',
+				maxInputTokens: 128000,
+				maxOutputTokens: 4096,
+				family: 'Azure',
+				tooltip: 'GPT-4 via Azure',
+				capabilities: {
+					toolCalling: true,
+					imageInput: false
+				},
+				thinking: false
+			};
+
+			mockMessages = [
+				new vscode.LanguageModelChatMessage(vscode.LanguageModelChatMessageRole.User, 'Hello')
+			];
+
+			mockOptions = {
+				requestInitiator: 'test-extension',
+				tools: [],
+				toolMode: vscode.LanguageModelChatToolMode.Auto
+			};
+
+			mockProgress = { report: vi.fn() };
+			mockToken = new vscode.CancellationTokenSource().token;
+		});
+
+		it('should delegate to parent class when in API Key mode', async () => {
+			const parentProto = Object.getPrototypeOf(Object.getPrototypeOf(provider));
+			const parentProvideResponse = vi.spyOn(parentProto, 'provideLanguageModelChatResponse').mockResolvedValue(undefined);
+			const getSessionSpy = vi.spyOn(vscode.authentication, 'getSession');
+
+			await provider.provideLanguageModelChatResponse(mockModel, mockMessages, mockOptions, mockProgress, mockToken);
+
+			expect(getSessionSpy).not.toHaveBeenCalled();
+			expect(parentProvideResponse).toHaveBeenCalledWith(mockModel, mockMessages, mockOptions, mockProgress, mockToken);
+		});
+
+		it('should not use Entra ID authentication in API Key mode', async () => {
+			const getSessionSpy = vi.spyOn(vscode.authentication, 'getSession');
+			const parentProto = Object.getPrototypeOf(Object.getPrototypeOf(provider));
+			vi.spyOn(parentProto, 'provideLanguageModelChatResponse').mockResolvedValue(undefined);
+
+			await provider.provideLanguageModelChatResponse(mockModel, mockMessages, mockOptions, mockProgress, mockToken);
+
+			expect(getSessionSpy).not.toHaveBeenCalled();
+		});
+	});
+
+	describe('configuration', () => {
+		it('should use AzureAuthMode enum for configuration', () => {
+			const configService = accessor.get(IConfigurationService);
+			const getConfigSpy = vi.spyOn(configService, 'getConfig').mockReturnValue(AzureAuthMode.EntraId);
+
+			provider = instaService.createInstance(AzureBYOKModelProvider, mockByokStorageService);
+
+			configService.getConfig(ConfigKey.AzureAuthType);
+
+			expect(getConfigSpy).toHaveBeenCalledWith(ConfigKey.AzureAuthType);
+		});
+
+		it('should default to Entra ID mode', () => {
+			const defaultValue = ConfigKey.AzureAuthType.defaultValue;
+
+			expect(defaultValue).toBe(AzureAuthMode.EntraId);
+		});
+	});
+});

extensions/copilot/src/platform/configuration/common/configurationService.ts

@@ -577,6 +577,18 @@ export enum AuthPermissionMode {
 	Minimal = 'minimal'
 }
 
+export enum AzureAuthMode {
+	EntraId = 'entraId',
+	ApiKey = 'apiKey'
+}
+
+export namespace AzureAuthMode {
+	/** Microsoft authentication provider ID for VS Code authentication API */
+	export const MICROSOFT_AUTH_PROVIDER = 'microsoft';
+	/** Azure Cognitive Services scope for Entra ID authentication */
+	export const COGNITIVE_SERVICES_SCOPE = 'https://cognitiveservices.azure.com/.default';
+}
+
 export type CodeGenerationImportInstruction = { language?: string; file: string };
 export type CodeGenerationTextInstruction = { language?: string; text: string };
 export type CodeGenerationInstruction = CodeGenerationImportInstruction | CodeGenerationTextInstruction;
@@ -838,6 +850,7 @@ export namespace ConfigKey {
 	export const CurrentEditorAgentContext = defineSetting<boolean>('chat.agent.currentEditorContext.enabled', ConfigType.Simple, true);
 	/** BYOK  */
 	export const OllamaEndpoint = defineSetting<string>('chat.byok.ollamaEndpoint', ConfigType.Simple, 'http://localhost:11434');
+	export const AzureAuthType = defineSetting<AzureAuthMode>('chat.azureAuthType', ConfigType.Simple, AzureAuthMode.EntraId);
 	export const AzureModels = defineSetting<Record<string, { name: string; url: string; toolCalling: boolean; vision: boolean; maxInputTokens: number; maxOutputTokens: number; requiresAPIKey?: boolean; thinking?: boolean }>>('chat.azureModels', ConfigType.Simple, {});
 	export const CustomOAIModels = defineSetting<Record<string, { name: string; url: string; toolCalling: boolean; vision: boolean; maxInputTokens: number; maxOutputTokens: number; requiresAPIKey?: boolean; thinking?: boolean; requestHeaders?: Record<string, string> }>>('chat.customOAIModels', ConfigType.Simple, {});
 	export const AutoFixDiagnostics = defineSetting<boolean>('chat.agent.autoFix', ConfigType.ExperimentBased, true);

extensions/copilot/src/util/common/test/shims/chatTypes.ts

@@ -444,6 +444,31 @@ export enum LanguageModelChatMessageRole {
 	System = 3
 }
 
+export enum LanguageModelChatToolMode {
+	Auto = 1,
+	Required = 2
+}
+
+export class LanguageModelChatMessage implements vscode.LanguageModelChatMessage {
+	role: LanguageModelChatMessageRole;
+	content: Array<any>;
+	name: string | undefined;
+
+	constructor(role: LanguageModelChatMessageRole, content: string | Array<any>, name?: string) {
+		this.role = role;
+		this.content = typeof content === 'string' ? [{ type: 'text', value: content }] : content;
+		this.name = name;
+	}
+
+	static User(content: string | Array<any>, name?: string): LanguageModelChatMessage {
+		return new LanguageModelChatMessage(LanguageModelChatMessageRole.User, content, name);
+	}
+
+	static Assistant(content: string | Array<any>, name?: string): LanguageModelChatMessage {
+		return new LanguageModelChatMessage(LanguageModelChatMessageRole.Assistant, content, name);
+	}
+}
+
 export class ChatToolInvocationPart {
 	toolName: string;
 	toolCallId: string;

extensions/copilot/src/util/common/test/shims/vscodeTypesShim.ts

@@ -3,7 +3,6 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import * as vscodeTypes from '../../../../vscodeTypes';
 import { CancellationTokenSource } from '../../../vs/base/common/cancellation';
 import { Emitter as EventEmitter } from '../../../vs/base/common/event';
 import { URI as Uri } from '../../../vs/base/common/uri';
@@ -18,14 +17,14 @@ import { SnippetString } from '../../../vs/workbench/api/common/extHostTypes/sni
 import { SnippetTextEdit } from '../../../vs/workbench/api/common/extHostTypes/snippetTextEdit';
 import { SymbolInformation, SymbolKind } from '../../../vs/workbench/api/common/extHostTypes/symbolInformation';
 import { EndOfLine, TextEdit } from '../../../vs/workbench/api/common/extHostTypes/textEdit';
-import { AISearchKeyword, ChatErrorLevel, ChatPrepareToolInvocationPart, ChatReferenceBinaryData, ChatReferenceDiagnostic, ChatRequestEditedFileEventKind, ChatRequestEditorData, ChatRequestNotebookData, ChatRequestTurn, ChatResponseAnchorPart, ChatResponseClearToPreviousToolInvocationReason, ChatResponseCodeblockUriPart, ChatResponseCodeCitationPart, ChatResponseCommandButtonPart, ChatResponseConfirmationPart, ChatResponseExtensionsPart, ChatResponseExternalEditPart, ChatResponseFileTreePart, ChatResponseMarkdownPart, ChatResponseMarkdownWithVulnerabilitiesPart, ChatResponseMovePart, ChatResponseNotebookEditPart, ChatResponseProgressPart, ChatResponseProgressPart2, ChatResponsePullRequestPart, ChatResponseReferencePart, ChatResponseReferencePart2, ChatResponseTextEditPart, ChatResponseThinkingProgressPart, ChatResponseTurn, ChatResponseTurn2, ChatResponseWarningPart, ChatSessionStatus, ChatToolInvocationPart, ExcludeSettingOptions, LanguageModelChatMessageRole, LanguageModelDataPart, LanguageModelDataPart2, LanguageModelError, LanguageModelPartAudience, LanguageModelPromptTsxPart, LanguageModelTextPart, LanguageModelTextPart2, LanguageModelThinkingPart, LanguageModelToolCallPart, LanguageModelToolExtensionSource, LanguageModelToolMCPSource, LanguageModelToolResult, LanguageModelToolResult2, LanguageModelToolResultPart, LanguageModelToolResultPart2, TextSearchMatch2 } from './chatTypes';
+import { AISearchKeyword, ChatErrorLevel, ChatPrepareToolInvocationPart, ChatReferenceBinaryData, ChatReferenceDiagnostic, ChatRequestEditedFileEventKind, ChatRequestEditorData, ChatRequestNotebookData, ChatRequestTurn, ChatResponseAnchorPart, ChatResponseClearToPreviousToolInvocationReason, ChatResponseCodeblockUriPart, ChatResponseCodeCitationPart, ChatResponseCommandButtonPart, ChatResponseConfirmationPart, ChatResponseExtensionsPart, ChatResponseExternalEditPart, ChatResponseFileTreePart, ChatResponseMarkdownPart, ChatResponseMarkdownWithVulnerabilitiesPart, ChatResponseMovePart, ChatResponseNotebookEditPart, ChatResponseProgressPart, ChatResponseProgressPart2, ChatResponsePullRequestPart, ChatResponseReferencePart, ChatResponseReferencePart2, ChatResponseTextEditPart, ChatResponseThinkingProgressPart, ChatResponseTurn, ChatResponseTurn2, ChatResponseWarningPart, ChatSessionStatus, ChatToolInvocationPart, ExcludeSettingOptions, LanguageModelChatMessage, LanguageModelChatMessageRole, LanguageModelChatToolMode, LanguageModelDataPart, LanguageModelDataPart2, LanguageModelError, LanguageModelPartAudience, LanguageModelPromptTsxPart, LanguageModelTextPart, LanguageModelTextPart2, LanguageModelThinkingPart, LanguageModelToolCallPart, LanguageModelToolExtensionSource, LanguageModelToolMCPSource, LanguageModelToolResult, LanguageModelToolResult2, LanguageModelToolResultPart, LanguageModelToolResultPart2, TextSearchMatch2 } from './chatTypes';
 import { TextDocumentChangeReason, TextEditorSelectionChangeKind, WorkspaceEdit } from './editing';
 import { ChatLocation, ChatVariableLevel, DiagnosticSeverity, ExtensionMode, FileType, TextEditorCursorStyle, TextEditorLineNumbersStyle, TextEditorRevealType } from './enums';
 import { t } from './l10n';
 import { NewSymbolName, NewSymbolNameTag, NewSymbolNameTriggerKind } from './newSymbolName';
 import { TerminalShellExecutionCommandLineConfidence } from './terminal';
 
-const shim: typeof vscodeTypes = {
+const shim = {
 	Position,
 	Range,
 	Selection,
@@ -108,6 +107,8 @@ const shim: typeof vscodeTypes = {
 	LanguageModelToolResultPart,
 	LanguageModelToolResultPart2,
 	LanguageModelChatMessageRole,
+	LanguageModelChatMessage,
+	LanguageModelChatToolMode,
 	TextEditorSelectionChangeKind,
 	TextDocumentChangeReason,
 	ChatToolInvocationPart,
@@ -118,7 +119,10 @@ const shim: typeof vscodeTypes = {
 	SnippetString,
 	SnippetTextEdit,
 	FileType,
-	ChatSessionStatus
+	ChatSessionStatus,
+	authentication: {
+		getSession: (providerId: string, scopes: readonly string[], options?: { createIfNone?: boolean; silent?: boolean; clearSessionPreference?: boolean; forceNewSession?: boolean }) => Promise.reject(new Error('authentication.getSession not mocked in test'))
+	}
 };
 
 export = shim;

extensions/copilot/src/vscodeTypes.ts

@@ -86,6 +86,8 @@ export import LanguageModelToolCallPart = vscode.LanguageModelToolCallPart;
 export import LanguageModelToolResultPart = vscode.LanguageModelToolResultPart;
 export import LanguageModelToolResultPart2 = vscode.LanguageModelToolResultPart2;
 export import LanguageModelChatMessageRole = vscode.LanguageModelChatMessageRole;
+export import LanguageModelChatMessage = vscode.LanguageModelChatMessage;
+export import LanguageModelChatToolMode = vscode.LanguageModelChatToolMode;
 export import TextEditorSelectionChangeKind = vscode.TextEditorSelectionChangeKind;
 export import TextDocumentChangeReason = vscode.TextDocumentChangeReason;
 export import ChatToolInvocationPart = vscode.ChatToolInvocationPart;