diff --git a/extensions/copilot/package.json b/extensions/copilot/package.json index 1cd25151013..541e752f0f2 100644 --- a/extensions/copilot/package.json +++ b/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": {}, diff --git a/extensions/copilot/package.nls.json b/extensions/copilot/package.nls.json index b1efce0cb8d..ee83865de56 100644 --- a/extensions/copilot/package.nls.json +++ b/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.", diff --git a/extensions/copilot/src/extension/byok/node/azureOpenAIEndpoint.ts b/extensions/copilot/src/extension/byok/node/azureOpenAIEndpoint.ts new file mode 100644 index 00000000000..2faf2ee9cbf --- /dev/null +++ b/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 { + 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; + } +} diff --git a/extensions/copilot/src/extension/byok/node/test/azureOpenAIEndpoint.spec.ts b/extensions/copilot/src/extension/byok/node/test/azureOpenAIEndpoint.spec.ts new file mode 100644 index 00000000000..d77e1062e92 --- /dev/null +++ b/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'); + }); + }); +}); diff --git a/extensions/copilot/src/extension/byok/vscode-node/azureProvider.ts b/extensions/copilot/src/extension/byok/vscode-node/azureProvider.ts index 673e8ca845b..de1813e56f5 100644 --- a/extensions/copilot/src/extension/byok/vscode-node/azureProvider.ts +++ b/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 { + // 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, + options: ProvideLanguageModelChatResponseOptions, + progress: Progress, + token: CancellationToken + ): Promise { + 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); + } + } } diff --git a/extensions/copilot/src/extension/byok/vscode-node/customOAIProvider.ts b/extensions/copilot/src/extension/byok/vscode-node/customOAIProvider.ts index ce0448caa64..70c2609737f 100644 --- a/extensions/copilot/src/extension/byok/vscode-node/customOAIProvider.ts +++ b/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; @@ -96,7 +96,7 @@ export class CustomOAIBYOKModelProvider implements BYOKModelProvider { + protected async getAllModels(): Promise { const modelConfig = this.getUserModelConfig(); const models: BYOKKnownModels = {}; @@ -119,7 +119,7 @@ export class CustomOAIBYOKModelProvider implements BYOKModelProvider { + protected async getModelsWithCredentials(silent: boolean): Promise { 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 { 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); diff --git a/extensions/copilot/src/extension/byok/vscode-node/test/azureProvider.spec.ts b/extensions/copilot/src/extension/byok/vscode-node/test/azureProvider.spec.ts new file mode 100644 index 00000000000..a4edbccfdde --- /dev/null +++ b/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; + 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; + 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); + }); + }); +}); diff --git a/extensions/copilot/src/platform/configuration/common/configurationService.ts b/extensions/copilot/src/platform/configuration/common/configurationService.ts index 0f8f4a0f15d..ff7af715fcd 100644 --- a/extensions/copilot/src/platform/configuration/common/configurationService.ts +++ b/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('chat.agent.currentEditorContext.enabled', ConfigType.Simple, true); /** BYOK */ export const OllamaEndpoint = defineSetting('chat.byok.ollamaEndpoint', ConfigType.Simple, 'http://localhost:11434'); + export const AzureAuthType = defineSetting('chat.azureAuthType', ConfigType.Simple, AzureAuthMode.EntraId); export const AzureModels = defineSetting>('chat.azureModels', ConfigType.Simple, {}); export const CustomOAIModels = defineSetting }>>('chat.customOAIModels', ConfigType.Simple, {}); export const AutoFixDiagnostics = defineSetting('chat.agent.autoFix', ConfigType.ExperimentBased, true); diff --git a/extensions/copilot/src/util/common/test/shims/chatTypes.ts b/extensions/copilot/src/util/common/test/shims/chatTypes.ts index ffd46cce883..5cd3f321148 100644 --- a/extensions/copilot/src/util/common/test/shims/chatTypes.ts +++ b/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; + name: string | undefined; + + constructor(role: LanguageModelChatMessageRole, content: string | Array, name?: string) { + this.role = role; + this.content = typeof content === 'string' ? [{ type: 'text', value: content }] : content; + this.name = name; + } + + static User(content: string | Array, name?: string): LanguageModelChatMessage { + return new LanguageModelChatMessage(LanguageModelChatMessageRole.User, content, name); + } + + static Assistant(content: string | Array, name?: string): LanguageModelChatMessage { + return new LanguageModelChatMessage(LanguageModelChatMessageRole.Assistant, content, name); + } +} + export class ChatToolInvocationPart { toolName: string; toolCallId: string; diff --git a/extensions/copilot/src/util/common/test/shims/vscodeTypesShim.ts b/extensions/copilot/src/util/common/test/shims/vscodeTypesShim.ts index c40f615ddf0..0d1a3d50df6 100644 --- a/extensions/copilot/src/util/common/test/shims/vscodeTypesShim.ts +++ b/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; diff --git a/extensions/copilot/src/vscodeTypes.ts b/extensions/copilot/src/vscodeTypes.ts index e0f65a2a81f..4853646c434 100644 --- a/extensions/copilot/src/vscodeTypes.ts +++ b/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;