Skip to content

Commit ff93334

Browse files
lramos15lramos15
authored
Fix unable to change API keys in BYOK (#403)
1 parent a7fef93 commit ff93334

File tree

6 files changed

+139
-29
lines changed

6 files changed

+139
-29
lines changed

‎package.json

Lines changed: 16 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1547,31 +1547,37 @@
15471547
},
15481548
{
15491549
"vendor": "azure",
1550-
"displayName": "Azure"
1550+
"displayName": "Azure",
1551+
"managementCommand": "github.copilot.chat.manageBYOK"
15511552
},
15521553
{
15531554
"vendor": "anthropic",
1554-
"displayName": "Anthropic"
1555+
"displayName": "Anthropic",
1556+
"managementCommand": "github.copilot.chat.manageBYOK"
15551557
},
15561558
{
15571559
"vendor": "ollama",
15581560
"displayName": "Ollama"
15591561
},
15601562
{
15611563
"vendor": "openai",
1562-
"displayName": "OpenAI"
1564+
"displayName": "OpenAI",
1565+
"managementCommand": "github.copilot.chat.manageBYOK"
15631566
},
15641567
{
15651568
"vendor": "gemini",
1566-
"displayName": "Gemini"
1569+
"displayName": "Gemini",
1570+
"managementCommand": "github.copilot.chat.manageBYOK"
15671571
},
15681572
{
15691573
"vendor": "groq",
1570-
"displayName": "Groq"
1574+
"displayName": "Groq",
1575+
"managementCommand": "github.copilot.chat.manageBYOK"
15711576
},
15721577
{
15731578
"vendor": "openrouter",
1574-
"displayName": "OpenRouter"
1579+
"displayName": "OpenRouter",
1580+
"managementCommand": "github.copilot.chat.manageBYOK"
15751581
}
15761582
],
15771583
"interactiveSession": [
@@ -2001,6 +2007,10 @@
20012007
"title": "Disable Follow Cell Execution from Chat",
20022008
"shortTitle": "Unfollow",
20032009
"icon": "$(pinned-dirty)"
2010+
},
2011+
{
2012+
"command": "github.copilot.chat.manageBYOK",
2013+
"title": "Manage Bring Your Own Key Vendor"
20042014
}
20052015
],
20062016
"configuration": [

‎src/extension/byok/common/byokProvider.ts

Lines changed: 9 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
* Copyright (c) Microsoft Corporation. All rights reserved.
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
5-
import type { Disposable, LanguageModelChatInformation } from 'vscode';
5+
import type { Disposable, LanguageModelChatInformation, LanguageModelChatProvider2 } from 'vscode';
66
import { CopilotToken } from '../../../platform/authentication/common/copilotToken';
77
import { ICAPIClientService } from '../../../platform/endpoint/common/capiClient';
88
import { IChatModelInformation } from '../../../platform/endpoint/common/endpointProvider';
@@ -62,6 +62,14 @@ export interface BYOKModelRegistry {
6262
registerModel(config: BYOKModelConfig): Promise<Disposable>;
6363
}
6464

65+
export interface BYOKModelProvider<T extends LanguageModelChatInformation> extends LanguageModelChatProvider2<T> {
66+
readonly authType: BYOKAuthType;
67+
/**
68+
* Called when the user is requesting an API key update. The provider should handle all the UI and updating the storage
69+
*/
70+
updateAPIKey(): Promise<void>;
71+
}
72+
6573
// Many model providers don't have robust model lists. This allows us to map id -> information about models, and then if we don't know the model just let the user enter a custom id
6674
export type BYOKKnownModels = Record<string, BYOKModelCapabilities>;
6775

‎src/extension/byok/vscode-node/anthropicProvider.ts

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
*--------------------------------------------------------------------------------------------*/
55

66
import Anthropic from '@anthropic-ai/sdk';
7-
import { CancellationToken, ChatResponseFragment2, LanguageModelChatInformation, LanguageModelChatMessage, LanguageModelChatMessage2, LanguageModelChatProvider2, LanguageModelChatRequestHandleOptions, LanguageModelTextPart, LanguageModelToolCallPart, Progress } from 'vscode';
7+
import { CancellationToken, ChatResponseFragment2, LanguageModelChatInformation, LanguageModelChatMessage, LanguageModelChatMessage2, LanguageModelChatRequestHandleOptions, LanguageModelTextPart, LanguageModelToolCallPart, Progress } from 'vscode';
88
import { ChatFetchResponseType, ChatLocation } from '../../../platform/chat/common/commonTypes';
99
import { ILogService } from '../../../platform/log/common/logService';
1010
import { IResponseDelta, OpenAiFunctionTool } from '../../../platform/networking/common/fetch';
@@ -13,13 +13,14 @@ import { IRequestLogger } from '../../../platform/requestLogger/node/requestLogg
1313
import { RecordedProgress } from '../../../util/common/progressRecorder';
1414
import { toErrorMessage } from '../../../util/vs/base/common/errorMessage';
1515
import { generateUuid } from '../../../util/vs/base/common/uuid';
16-
import { BYOKAuthType, BYOKKnownModels, byokKnownModelsToAPIInfo, BYOKModelCapabilities } from '../common/byokProvider';
16+
import { BYOKAuthType, BYOKKnownModels, byokKnownModelsToAPIInfo, BYOKModelCapabilities, BYOKModelProvider } from '../common/byokProvider';
1717
import { anthropicMessagesToRawMessagesForLogging, apiMessageToAnthropicMessage } from './anthropicMessageConverter';
1818
import { IBYOKStorageService } from './byokStorageService';
1919
import { promptForAPIKey } from './byokUIService';
2020

21-
export class AnthropicLMProvider implements LanguageModelChatProvider2<LanguageModelChatInformation> {
21+
export class AnthropicLMProvider implements BYOKModelProvider<LanguageModelChatInformation> {
2222
public static readonly providerName = 'Anthropic';
23+
public readonly authType: BYOKAuthType = BYOKAuthType.GlobalApiKey;
2324
private _anthropicAPIClient: Anthropic | undefined;
2425
private _apiKey: string | undefined;
2526
constructor(
@@ -49,6 +50,13 @@ export class AnthropicLMProvider implements LanguageModelChatProvider2<LanguageM
4950
}
5051
}
5152

53+
async updateAPIKey(): Promise<void> {
54+
this._apiKey = await promptForAPIKey(AnthropicLMProvider.providerName, await this._byokStorageService.getAPIKey(AnthropicLMProvider.providerName) !== undefined);
55+
if (this._apiKey) {
56+
this._byokStorageService.storeAPIKey(AnthropicLMProvider.providerName, this._apiKey, BYOKAuthType.GlobalApiKey);
57+
}
58+
}
59+
5260
async prepareLanguageModelChat(options: { silent: boolean }, token: CancellationToken): Promise<LanguageModelChatInformation[]> {
5361
if (!this._apiKey) { // If we don't have the API key it might just be in storage, so we try to read it first
5462
this._apiKey = await this._byokStorageService.getAPIKey(AnthropicLMProvider.providerName);
@@ -59,9 +67,8 @@ export class AnthropicLMProvider implements LanguageModelChatProvider2<LanguageM
5967
} else if (options.silent && !this._apiKey) {
6068
return [];
6169
} else { // Not silent, and no api key = good to prompt user for api key
62-
this._apiKey = await promptForAPIKey(AnthropicLMProvider.providerName, false);
70+
await this.updateAPIKey();
6371
if (this._apiKey) {
64-
this._byokStorageService.storeAPIKey(AnthropicLMProvider.providerName, this._apiKey, BYOKAuthType.GlobalApiKey);
6572
return byokKnownModelsToAPIInfo(AnthropicLMProvider.providerName, await this.getAllModels(this._apiKey));
6673
} else {
6774
return [];

‎src/extension/byok/vscode-node/azureProvider.ts

Lines changed: 68 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -3,12 +3,12 @@
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
55

6-
import { CancellationToken, ChatResponseFragment2, LanguageModelChatInformation, LanguageModelChatMessage, LanguageModelChatMessage2, LanguageModelChatProvider2, LanguageModelChatRequestHandleOptions, Progress } from 'vscode';
6+
import { CancellationToken, ChatResponseFragment2, Event, LanguageModelChatInformation, LanguageModelChatMessage, LanguageModelChatMessage2, LanguageModelChatRequestHandleOptions, Progress, QuickPickItem, window } from 'vscode';
77
import { ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService';
88
import { ILogService } from '../../../platform/log/common/logService';
99
import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation';
1010
import { CopilotLanguageModelWrapper } from '../../conversation/vscode-node/languageModelAccess';
11-
import { BYOKAuthType, BYOKKnownModels, resolveModelInfo } from '../common/byokProvider';
11+
import { BYOKAuthType, BYOKKnownModels, BYOKModelProvider, resolveModelInfo } from '../common/byokProvider';
1212
import { OpenAIEndpoint } from '../node/openAIEndpoint';
1313
import { IBYOKStorageService } from './byokStorageService';
1414
import { promptForAPIKey } from './byokUIService';
@@ -43,9 +43,10 @@ interface AzureModelInfo extends LanguageModelChatInformation {
4343
thinking: boolean;
4444
}
4545

46-
export class AzureBYOKModelProvider implements LanguageModelChatProvider2<AzureModelInfo> {
46+
export class AzureBYOKModelProvider implements BYOKModelProvider<AzureModelInfo> {
4747
private readonly _lmWrapper: CopilotLanguageModelWrapper;
4848
static readonly providerName = 'Azure';
49+
public readonly authType: BYOKAuthType = BYOKAuthType.PerModelDeployment;
4950
constructor(
5051
private readonly _byokStorageService: IBYOKStorageService,
5152
@IConfigurationService private readonly _configurationService: IConfigurationService,
@@ -55,6 +56,8 @@ export class AzureBYOKModelProvider implements LanguageModelChatProvider2<AzureM
5556
this._lmWrapper = this._instantiationService.createInstance(CopilotLanguageModelWrapper);
5657
}
5758

59+
onDidChange?: Event<void> | undefined;
60+
5861
private async getAllModels(): Promise<BYOKKnownModels> {
5962
const azureModelConfig = this._configurationService.getConfig(ConfigKey.AzureModels);
6063
const models: BYOKKnownModels = {};
@@ -152,4 +155,66 @@ export class AzureBYOKModelProvider implements LanguageModelChatProvider2<AzureM
152155
return this._lmWrapper.provideTokenCount(openAIChatEndpoint, text);
153156
}
154157

158+
public async updateAPIKey(): Promise<void> {
159+
// Get all available models
160+
const allModels = await this.getAllModels();
161+
162+
if (Object.keys(allModels).length === 0) {
163+
await window.showInformationMessage('No Azure models are configured. Please configure models first.');
164+
return;
165+
}
166+
167+
// Create quick pick items for all models
168+
interface ModelQuickPickItem extends QuickPickItem {
169+
modelId: string;
170+
}
171+
172+
const modelItems: ModelQuickPickItem[] = Object.entries(allModels).map(([modelId, modelInfo]) => ({
173+
label: modelInfo.name || modelId,
174+
description: modelId,
175+
detail: `URL: ${modelInfo.url}`,
176+
modelId: modelId
177+
}));
178+
179+
// Show quick pick to select which model's API key to update
180+
const quickPick = window.createQuickPick<ModelQuickPickItem>();
181+
quickPick.title = 'Update Azure Model API Key';
182+
quickPick.placeholder = 'Select a model to update its API key';
183+
quickPick.items = modelItems;
184+
quickPick.ignoreFocusOut = true;
185+
186+
const selectedModel = await new Promise<ModelQuickPickItem | undefined>((resolve) => {
187+
quickPick.onDidAccept(() => {
188+
const selected = quickPick.selectedItems[0];
189+
quickPick.hide();
190+
resolve(selected);
191+
});
192+
193+
quickPick.onDidHide(() => {
194+
resolve(undefined);
195+
});
196+
197+
quickPick.show();
198+
});
199+
200+
if (!selectedModel) {
201+
return; // User cancelled
202+
}
203+
204+
// Prompt for new API key
205+
const newApiKey = await promptForAPIKey(`Azure - ${selectedModel.modelId}`, true);
206+
207+
if (newApiKey !== undefined) {
208+
if (newApiKey.trim() === '') {
209+
// Empty string means delete the API key
210+
await this._byokStorageService.deleteAPIKey(AzureBYOKModelProvider.providerName, BYOKAuthType.PerModelDeployment, selectedModel.modelId);
211+
await window.showInformationMessage(`API key for ${selectedModel.label} has been deleted.`);
212+
} else {
213+
// Store the new API key
214+
await this._byokStorageService.storeAPIKey(AzureBYOKModelProvider.providerName, newApiKey, BYOKAuthType.PerModelDeployment, selectedModel.modelId);
215+
await window.showInformationMessage(`API key for ${selectedModel.label} has been updated.`);
216+
}
217+
}
218+
}
219+
155220
}

‎src/extension/byok/vscode-node/baseOpenAICompatibleProvider.ts

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -3,18 +3,18 @@
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
55

6-
import { CancellationToken, ChatResponseFragment2, LanguageModelChatInformation, LanguageModelChatMessage, LanguageModelChatMessage2, LanguageModelChatProvider2, LanguageModelChatRequestHandleOptions, Progress } from 'vscode';
6+
import { CancellationToken, ChatResponseFragment2, LanguageModelChatInformation, LanguageModelChatMessage, LanguageModelChatMessage2, LanguageModelChatRequestHandleOptions, Progress } from 'vscode';
77
import { IChatModelInformation } from '../../../platform/endpoint/common/endpointProvider';
88
import { ILogService } from '../../../platform/log/common/logService';
99
import { IFetcherService } from '../../../platform/networking/common/fetcherService';
1010
import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation';
1111
import { CopilotLanguageModelWrapper } from '../../conversation/vscode-node/languageModelAccess';
12-
import { BYOKAuthType, BYOKKnownModels, byokKnownModelsToAPIInfo, BYOKModelCapabilities, resolveModelInfo } from '../common/byokProvider';
12+
import { BYOKAuthType, BYOKKnownModels, byokKnownModelsToAPIInfo, BYOKModelCapabilities, BYOKModelProvider, resolveModelInfo } from '../common/byokProvider';
1313
import { OpenAIEndpoint } from '../node/openAIEndpoint';
1414
import { IBYOKStorageService } from './byokStorageService';
1515
import { promptForAPIKey } from './byokUIService';
1616

17-
export abstract class BaseOpenAICompatibleLMProvider implements LanguageModelChatProvider2<LanguageModelChatInformation> {
17+
export abstract class BaseOpenAICompatibleLMProvider implements BYOKModelProvider<LanguageModelChatInformation> {
1818

1919
private readonly _lmWrapper: CopilotLanguageModelWrapper;
2020
private _apiKey: string | undefined;
@@ -72,9 +72,8 @@ export abstract class BaseOpenAICompatibleLMProvider implements LanguageModelCha
7272
} else if (options.silent && !this._apiKey) {
7373
return [];
7474
} else { // Not silent, and no api key = good to prompt user for api key
75-
this._apiKey = await promptForAPIKey(this._name, false);
75+
await this.updateAPIKey();
7676
if (this._apiKey) {
77-
this._byokStorageService.storeAPIKey(this._name, this._apiKey, BYOKAuthType.GlobalApiKey);
7877
return byokKnownModelsToAPIInfo(this._name, await this.getAllModels());
7978
} else {
8079
return [];
@@ -94,4 +93,14 @@ export abstract class BaseOpenAICompatibleLMProvider implements LanguageModelCha
9493
const openAIChatEndpoint = this._instantiationService.createInstance(OpenAIEndpoint, modelInfo, this._apiKey ?? '', `${this._baseUrl}/chat/completions`);
9594
return this._lmWrapper.provideTokenCount(openAIChatEndpoint, text);
9695
}
96+
97+
async updateAPIKey(): Promise<void> {
98+
if (this.authType === BYOKAuthType.None) {
99+
return;
100+
}
101+
this._apiKey = await promptForAPIKey(this._name, await this._byokStorageService.getAPIKey(this._name) !== undefined);
102+
if (this._apiKey) {
103+
this._byokStorageService.storeAPIKey(this._name, this._apiKey, BYOKAuthType.GlobalApiKey);
104+
}
105+
}
97106
}

‎src/extension/byok/vscode-node/byokContribution.ts

Lines changed: 20 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22
* Copyright (c) Microsoft Corporation. All rights reserved.
33
* Licensed under the MIT License. See License.txt in the project root for license information.
44
*--------------------------------------------------------------------------------------------*/
5-
import { lm } from 'vscode';
5+
import { commands, LanguageModelChatInformation, lm } from 'vscode';
66
import { IAuthenticationService } from '../../../platform/authentication/common/authentication';
77
import { ConfigKey, IConfigurationService } from '../../../platform/configuration/common/configurationService';
88
import { ICAPIClientService } from '../../../platform/endpoint/common/capiClient';
@@ -11,7 +11,7 @@ import { ILogService } from '../../../platform/log/common/logService';
1111
import { IFetcherService } from '../../../platform/networking/common/fetcherService';
1212
import { Disposable } from '../../../util/vs/base/common/lifecycle';
1313
import { IInstantiationService } from '../../../util/vs/platform/instantiation/common/instantiation';
14-
import { BYOKKnownModels, isBYOKEnabled } from '../../byok/common/byokProvider';
14+
import { BYOKKnownModels, BYOKModelProvider, isBYOKEnabled } from '../../byok/common/byokProvider';
1515
import { IExtensionContribution } from '../../common/contributions';
1616
import { AnthropicLMProvider } from './anthropicProvider';
1717
import { AzureBYOKModelProvider } from './azureProvider';
@@ -25,6 +25,7 @@ import { OpenRouterLMProvider } from './openRouterProvider';
2525
export class BYOKContrib extends Disposable implements IExtensionContribution {
2626
public readonly id: string = 'byok-contribution';
2727
private readonly _byokStorageService: IBYOKStorageService;
28+
private readonly _providers: Map<string, BYOKModelProvider<LanguageModelChatInformation>> = new Map();
2829

2930
constructor(
3031
@IFetcherService private readonly _fetcherService: IFetcherService,
@@ -36,6 +37,12 @@ export class BYOKContrib extends Disposable implements IExtensionContribution {
3637
@IInstantiationService private readonly _instantiationService: IInstantiationService,
3738
) {
3839
super();
40+
this._register(commands.registerCommand('github.copilot.chat.manageBYOK', async (vendor: string) => {
41+
const provider = this._providers.get(vendor);
42+
if (provider) {
43+
await provider.updateAPIKey();
44+
}
45+
}));
3946
this._byokStorageService = new BYOKStorageService(extensionContext);
4047
this._authChange(authService, this._instantiationService);
4148

@@ -48,13 +55,17 @@ export class BYOKContrib extends Disposable implements IExtensionContribution {
4855
if (authService.copilotToken && isBYOKEnabled(authService.copilotToken, this._capiClientService)) {
4956
// Update known models list from CDN so all providers have the same list
5057
const knownModels = await this.fetchKnownModelList(this._fetcherService);
51-
this._store.add(lm.registerChatModelProvider(OllamaLMProvider.providerName.toLowerCase(), this._instantiationService.createInstance(OllamaLMProvider, this._configurationService.getConfig(ConfigKey.OllamaEndpoint), this._byokStorageService)));
52-
this._store.add(lm.registerChatModelProvider(AnthropicLMProvider.providerName.toLowerCase(), this._instantiationService.createInstance(AnthropicLMProvider, knownModels[AnthropicLMProvider.providerName], this._byokStorageService)));
53-
this._store.add(lm.registerChatModelProvider(GroqBYOKLMProvider.providerName.toLowerCase(), this._instantiationService.createInstance(GroqBYOKLMProvider, knownModels[GroqBYOKLMProvider.providerName], this._byokStorageService)));
54-
this._store.add(lm.registerChatModelProvider(GeminiBYOKLMProvider.providerName.toLowerCase(), this._instantiationService.createInstance(GeminiBYOKLMProvider, knownModels[GeminiBYOKLMProvider.providerName], this._byokStorageService)));
55-
this._store.add(lm.registerChatModelProvider(OAIBYOKLMProvider.providerName.toLowerCase(), this._instantiationService.createInstance(OAIBYOKLMProvider, knownModels[OAIBYOKLMProvider.providerName], this._byokStorageService)));
56-
this._store.add(lm.registerChatModelProvider(OpenRouterLMProvider.providerName.toLowerCase(), this._instantiationService.createInstance(OpenRouterLMProvider, this._byokStorageService)));
57-
this._store.add(lm.registerChatModelProvider('azure', this._instantiationService.createInstance(AzureBYOKModelProvider, this._byokStorageService)));
58+
this._providers.set(OllamaLMProvider.providerName.toLowerCase(), instantiationService.createInstance(OllamaLMProvider, this._configurationService.getConfig(ConfigKey.OllamaEndpoint), this._byokStorageService));
59+
this._providers.set(AnthropicLMProvider.providerName.toLowerCase(), instantiationService.createInstance(AnthropicLMProvider, knownModels[AnthropicLMProvider.providerName], this._byokStorageService));
60+
this._providers.set(GroqBYOKLMProvider.providerName.toLowerCase(), instantiationService.createInstance(GroqBYOKLMProvider, knownModels[GroqBYOKLMProvider.providerName], this._byokStorageService));
61+
this._providers.set(GeminiBYOKLMProvider.providerName.toLowerCase(), instantiationService.createInstance(GeminiBYOKLMProvider, knownModels[GeminiBYOKLMProvider.providerName], this._byokStorageService));
62+
this._providers.set(OAIBYOKLMProvider.providerName.toLowerCase(), instantiationService.createInstance(OAIBYOKLMProvider, knownModels[OAIBYOKLMProvider.providerName], this._byokStorageService));
63+
this._providers.set(OpenRouterLMProvider.providerName.toLowerCase(), instantiationService.createInstance(OpenRouterLMProvider, this._byokStorageService));
64+
this._providers.set(AzureBYOKModelProvider.providerName.toLowerCase(), instantiationService.createInstance(AzureBYOKModelProvider, this._byokStorageService));
65+
66+
for (const [providerName, provider] of this._providers) {
67+
this._store.add(lm.registerChatModelProvider(providerName, provider));
68+
}
5869
}
5970
}
6071
private async fetchKnownModelList(fetcherService: IFetcherService): Promise<Record<string, BYOKKnownModels>> {

0 commit comments

Comments
 (0)