diff --git a/src/vs/workbench/api/browser/mainThreadChatProvider.ts b/src/vs/workbench/api/browser/mainThreadChatProvider.ts index 1a1d046357a..801b0139627 100644 --- a/src/vs/workbench/api/browser/mainThreadChatProvider.ts +++ b/src/vs/workbench/api/browser/mainThreadChatProvider.ts @@ -5,7 +5,7 @@ import { CancellationToken } from 'vs/base/common/cancellation'; import { Emitter, Event } from 'vs/base/common/event'; -import { DisposableMap, DisposableStore, IDisposable, toDisposable } from 'vs/base/common/lifecycle'; +import { Disposable, DisposableMap, DisposableStore, IDisposable, toDisposable } from 'vs/base/common/lifecycle'; import { localize } from 'vs/nls'; import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions'; import { ILogService } from 'vs/platform/log/common/log'; @@ -59,7 +59,9 @@ export class MainThreadChatProvider implements MainThreadChatProviderShape { } } })); - dipsosables.add(this._registerAuthenticationProvider(identifier)); + if (metadata.auth) { + dipsosables.add(this._registerAuthenticationProvider(metadata.extension, metadata.auth)); + } dipsosables.add(Registry.as(Extensions.ExtensionFeaturesRegistry).registerExtensionFeature({ id: `lm-${identifier}`, label: localize('languageModels', "Language Model ({0})", `${identifier}-${metadata.model}`), @@ -100,48 +102,54 @@ export class MainThreadChatProvider implements MainThreadChatProviderShape { return task; } - private _registerAuthenticationProvider(identifier: string): IDisposable { - const disposables = new DisposableStore(); + private _registerAuthenticationProvider(extension: ExtensionIdentifier, auth: { providerLabel: string; accountLabel?: string | undefined }): IDisposable { // This needs to be done in both MainThread & ExtHost ChatProvider - const authProviderId = INTERNAL_AUTH_PROVIDER_PREFIX + identifier; - // This is what will be displayed in the UI and the account used for managing access via Auth UI - const authAccountId = identifier; - this._authenticationService.registerAuthenticationProvider(authProviderId, new LanguageModelAccessAuthProvider(authProviderId, authAccountId)); + const authProviderId = INTERNAL_AUTH_PROVIDER_PREFIX + extension.value; + + // Only register one auth provider per extension + if (this._authenticationService.getProviderIds().includes(authProviderId)) { + return Disposable.None; + } + + const disposables = new DisposableStore(); + this._authenticationService.registerAuthenticationProvider(authProviderId, new LanguageModelAccessAuthProvider(authProviderId, auth.providerLabel, auth.accountLabel)); disposables.add(toDisposable(() => { this._authenticationService.unregisterAuthenticationProvider(authProviderId); })); disposables.add(this._authenticationService.onDidChangeSessions(async (e) => { if (e.providerId === authProviderId) { if (e.event.removed?.length) { - const allowedExtensions = this._authenticationService.readAllowedExtensions(authProviderId, authAccountId); + const allowedExtensions = this._authenticationService.readAllowedExtensions(authProviderId, authProviderId); const extensionsToUpdateAccess = []; for (const allowed of allowedExtensions) { - const extension = await this._extensionService.getExtension(allowed.id); - this._authenticationService.updateAllowedExtension(authProviderId, authAccountId, allowed.id, allowed.name, false); - if (extension) { + const from = await this._extensionService.getExtension(allowed.id); + this._authenticationService.updateAllowedExtension(authProviderId, authProviderId, allowed.id, allowed.name, false); + if (from) { extensionsToUpdateAccess.push({ - extension: extension.identifier, + from: from.identifier, + to: extension, enabled: false }); } } - this._proxy.$updateAccesslist(extensionsToUpdateAccess); + this._proxy.$updateModelAccesslist(extensionsToUpdateAccess); } } })); disposables.add(this._authenticationService.onDidChangeExtensionSessionAccess(async (e) => { - const allowedExtensions = this._authenticationService.readAllowedExtensions(authProviderId, authAccountId); + const allowedExtensions = this._authenticationService.readAllowedExtensions(authProviderId, authProviderId); const accessList = []; for (const allowedExtension of allowedExtensions) { - const extension = await this._extensionService.getExtension(allowedExtension.id); - if (extension) { + const from = await this._extensionService.getExtension(allowedExtension.id); + if (from) { accessList.push({ - extension: extension.identifier, + from: from.identifier, + to: extension, enabled: allowedExtension.allowed ?? true }); } } - this._proxy.$updateAccesslist(accessList); + this._proxy.$updateModelAccesslist(accessList); })); return disposables; } @@ -150,7 +158,6 @@ export class MainThreadChatProvider implements MainThreadChatProviderShape { // The fake AuthenticationProvider that will be used to gate access to the Language Model. There will be one per provider. class LanguageModelAccessAuthProvider implements IAuthenticationProvider { supportsMultipleAccounts = false; - label = 'Language Model'; // Important for updating the UI private _onDidChangeSessions: Emitter = new Emitter(); @@ -158,7 +165,11 @@ class LanguageModelAccessAuthProvider implements IAuthenticationProvider { private _session: AuthenticationSession | undefined; - constructor(readonly id: string, private readonly accountName: string) { } + constructor( + readonly id: string, + readonly label: string, + private readonly _accountLabel: string = localize('languageModelsAccountId', 'Language Models') + ) { } async getSessions(scopes?: string[] | undefined): Promise { // If there are no scopes and no session that means no extension has requested a session yet @@ -189,7 +200,7 @@ class LanguageModelAccessAuthProvider implements IAuthenticationProvider { id: 'fake-session', account: { id: this.id, - label: this.accountName, + label: this._accountLabel, }, accessToken: 'fake-access-token', scopes, diff --git a/src/vs/workbench/api/common/extHost.api.impl.ts b/src/vs/workbench/api/common/extHost.api.impl.ts index 071e935c051..2b4bfec98a6 100644 --- a/src/vs/workbench/api/common/extHost.api.impl.ts +++ b/src/vs/workbench/api/common/extHost.api.impl.ts @@ -1395,7 +1395,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I const chat: typeof vscode.chat = { registerChatResponseProvider(id: string, provider: vscode.ChatResponseProvider, metadata: vscode.ChatResponseProviderMetadata) { checkProposedApiEnabled(extension, 'chatProvider'); - return extHostChatProvider.registerLanguageModel(extension.identifier, id, provider, metadata); + return extHostChatProvider.registerLanguageModel(extension, id, provider, metadata); }, requestLanguageModelAccess(id, options) { checkProposedApiEnabled(extension, 'chatRequestAccess'); diff --git a/src/vs/workbench/api/common/extHost.protocol.ts b/src/vs/workbench/api/common/extHost.protocol.ts index 964afc507dc..2f931bb83ef 100644 --- a/src/vs/workbench/api/common/extHost.protocol.ts +++ b/src/vs/workbench/api/common/extHost.protocol.ts @@ -1184,6 +1184,7 @@ export interface MainThreadChatProviderShape extends IDisposable { export interface ExtHostChatProviderShape { $updateLanguageModels(data: { added?: string[]; removed?: string[] }): void; $updateAccesslist(data: { extension: ExtensionIdentifier; enabled: boolean }[]): void; + $updateModelAccesslist(data: { from: ExtensionIdentifier; to: ExtensionIdentifier; enabled: boolean }[]): void; $provideLanguageModelResponse(handle: number, requestId: number, from: ExtensionIdentifier, messages: IChatMessage[], options: { [name: string]: any }, token: CancellationToken): Promise; $handleResponseFragment(requestId: number, chunk: IChatResponseFragment): Promise; } diff --git a/src/vs/workbench/api/common/extHostChatProvider.ts b/src/vs/workbench/api/common/extHostChatProvider.ts index eae0f221ffa..108a893fb4d 100644 --- a/src/vs/workbench/api/common/extHostChatProvider.ts +++ b/src/vs/workbench/api/common/extHostChatProvider.ts @@ -11,7 +11,7 @@ import * as typeConvert from 'vs/workbench/api/common/extHostTypeConverters'; import type * as vscode from 'vscode'; import { Progress } from 'vs/platform/progress/common/progress'; import { IChatMessage, IChatResponseFragment } from 'vs/workbench/contrib/chat/common/chatProvider'; -import { ExtensionIdentifier, ExtensionIdentifierMap, ExtensionIdentifierSet, IRelaxedExtensionDescription } from 'vs/platform/extensions/common/extensions'; +import { ExtensionIdentifier, ExtensionIdentifierMap, ExtensionIdentifierSet, IExtensionDescription } from 'vs/platform/extensions/common/extensions'; import { AsyncIterableSource } from 'vs/base/common/async'; import { Emitter, Event } from 'vs/base/common/event'; import { ExtHostAuthentication } from 'vs/workbench/api/common/extHostAuthentication'; @@ -91,12 +91,14 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape { private readonly _proxy: MainThreadChatProviderShape; private readonly _onDidChangeAccess = new Emitter(); + private readonly _onDidChangeModelAccess = new Emitter<{ from: ExtensionIdentifier; to: ExtensionIdentifier }>(); private readonly _onDidChangeProviders = new Emitter(); readonly onDidChangeProviders = this._onDidChangeProviders.event; private readonly _languageModels = new Map(); private readonly _languageModelIds = new Set(); // these are ALL models, not just the one in this EH private readonly _accesslist = new ExtensionIdentifierMap(); + private readonly _modelAccessList = new ExtensionIdentifierMap(); private readonly _pendingRequest = new Map(); @@ -113,11 +115,22 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape { this._onDidChangeProviders.dispose(); } - registerLanguageModel(extension: ExtensionIdentifier, identifier: string, provider: vscode.ChatResponseProvider, metadata: vscode.ChatResponseProviderMetadata): IDisposable { + registerLanguageModel(extension: IExtensionDescription, identifier: string, provider: vscode.ChatResponseProvider, metadata: vscode.ChatResponseProviderMetadata): IDisposable { const handle = ExtHostChatProvider._idPool++; - this._languageModels.set(handle, { extension, provider }); - this._proxy.$registerProvider(handle, identifier, { extension, model: metadata.name ?? '' }); + this._languageModels.set(handle, { extension: extension.identifier, provider }); + let auth; + if (metadata.auth) { + auth = { + providerLabel: extension.displayName || extension.name, + accountLabel: typeof metadata.auth === 'object' ? metadata.auth.label : undefined + }; + } + this._proxy.$registerProvider(handle, identifier, { + extension: extension.identifier, + model: metadata.name ?? '', + auth + }); return toDisposable(() => { this._languageModels.delete(handle); @@ -190,7 +203,26 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape { this._onDidChangeAccess.fire(updated); } - async requestLanguageModelAccess(extension: Readonly, languageModelId: string, options?: vscode.LanguageModelAccessOptions): Promise { + $updateModelAccesslist(data: { from: ExtensionIdentifier; to: ExtensionIdentifier; enabled: boolean }[]): void { + const updated = new Array<{ from: ExtensionIdentifier; to: ExtensionIdentifier }>(); + for (const { from, to, enabled } of data) { + const set = this._modelAccessList.get(from) ?? new ExtensionIdentifierSet(); + const oldValue = set.has(to); + if (oldValue !== enabled) { + if (enabled) { + set.add(to); + } else { + set.delete(to); + } + this._modelAccessList.set(from, set); + const newItem = { from, to }; + updated.push(newItem); + this._onDidChangeModelAccess.fire(newItem); + } + } + } + + async requestLanguageModelAccess(extension: IExtensionDescription, languageModelId: string, options?: vscode.LanguageModelAccessOptions): Promise { const from = extension.identifier; // check if the extension is in the access list and allowed to make chat requests if (this._accesslist.get(from) === false) { @@ -206,7 +238,10 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape { } throw new Error(`Language model '${languageModelId}' NOT found`); } - await this._checkAuthAccess(extension, languageModelId, justification); + + if (metadata.auth) { + await this._checkAuthAccess(extension, { identifier: metadata.extension, displayName: metadata.auth?.providerLabel }, justification); + } const that = this; @@ -215,15 +250,18 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape { return metadata.model; }, get isRevoked() { - return !that._accesslist.get(from) || !that._languageModelIds.has(languageModelId); + return !that._accesslist.get(from) + || (metadata.auth && !that._modelAccessList.get(from)?.has(metadata.extension)) + || !that._languageModelIds.has(languageModelId); }, get onDidChangeAccess() { const onDidChangeAccess = Event.filter(that._onDidChangeAccess.event, set => set.has(from)); const onDidRemoveLM = Event.filter(that._onDidChangeProviders.event, e => e.removed.includes(languageModelId)); - return Event.signal(Event.any(onDidChangeAccess, onDidRemoveLM)); + const onDidChangeModelAccess = Event.filter(that._onDidChangeModelAccess.event, e => ExtensionIdentifier.equals(e.from, from) && ExtensionIdentifier.equals(e.to, metadata.extension)); + return Event.signal(Event.any(onDidChangeAccess, onDidRemoveLM, onDidChangeModelAccess)); }, makeChatRequest(messages, options, token) { - if (!that._accesslist.get(from)) { + if (!that._accesslist.get(from) || (metadata.auth && !that._modelAccessList.get(from)?.has(metadata.extension))) { throw new Error('Access to chat has been revoked'); } if (!that._languageModelIds.has(languageModelId)) { @@ -253,20 +291,22 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape { } // BIG HACK: Using AuthenticationProviders to check access to Language Models - private async _checkAuthAccess(from: Readonly, languageModelId: string, detail?: string): Promise { + private async _checkAuthAccess(from: IExtensionDescription, to: { identifier: ExtensionIdentifier; displayName: string }, detail?: string): Promise { // This needs to be done in both MainThread & ExtHost ChatProvider - const providerId = INTERNAL_AUTH_PROVIDER_PREFIX + languageModelId; + const providerId = INTERNAL_AUTH_PROVIDER_PREFIX + to.identifier.value; const session = await this._extHostAuthentication.getSession(from, providerId, [], { silent: true }); if (!session) { try { await this._extHostAuthentication.getSession(from, providerId, [], { forceNewSession: { - detail: detail ?? localize('chatAccess', "To allow access to the '{0}' language model", languageModelId), + detail: detail ?? localize('chatAccess', "To allow access to the language models provided by {0}", to.displayName), } }); } catch (err) { - throw new Error('Access to language model has not been granted'); + throw new Error('Access to language models has not been granted'); } } + + this.$updateModelAccesslist([{ from: from.identifier, to: to.identifier, enabled: true }]); } } diff --git a/src/vs/workbench/contrib/chat/common/chatProvider.ts b/src/vs/workbench/contrib/chat/common/chatProvider.ts index c4655e35665..c393a73de98 100644 --- a/src/vs/workbench/contrib/chat/common/chatProvider.ts +++ b/src/vs/workbench/contrib/chat/common/chatProvider.ts @@ -30,6 +30,10 @@ export interface IChatResponseProviderMetadata { readonly extension: ExtensionIdentifier; readonly model: string; readonly description?: string; + readonly auth?: { + readonly providerLabel: string; + readonly accountLabel?: string; + }; } export interface IChatResponseProvider { diff --git a/src/vscode-dts/vscode.proposed.chatProvider.d.ts b/src/vscode-dts/vscode.proposed.chatProvider.d.ts index 226416480c0..e35fd5547d1 100644 --- a/src/vscode-dts/vscode.proposed.chatProvider.d.ts +++ b/src/vscode-dts/vscode.proposed.chatProvider.d.ts @@ -26,6 +26,13 @@ declare module 'vscode' { */ // TODO@API rename to model name: string; + + /** + * When present, this gates the use of `requestLanguageModelAccess` behind an authorization flow where + * the user must approve of another extension accessing the models contributed by this extension. + * Additionally, the extension can provide a label that will be shown in the UI. + */ + auth?: true | { label: string }; } export namespace chat {