diff --git a/src/vs/workbench/api/browser/mainThreadChatProvider.ts b/src/vs/workbench/api/browser/mainThreadChatProvider.ts index dcf7e7ab3e0..1a1d046357a 100644 --- a/src/vs/workbench/api/browser/mainThreadChatProvider.ts +++ b/src/vs/workbench/api/browser/mainThreadChatProvider.ts @@ -4,7 +4,8 @@ *--------------------------------------------------------------------------------------------*/ import { CancellationToken } from 'vs/base/common/cancellation'; -import { DisposableMap, DisposableStore } from 'vs/base/common/lifecycle'; +import { Emitter, Event } from 'vs/base/common/event'; +import { 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'; @@ -12,8 +13,10 @@ import { IProgress, Progress } from 'vs/platform/progress/common/progress'; import { Registry } from 'vs/platform/registry/common/platform'; import { ExtHostChatProviderShape, ExtHostContext, MainContext, MainThreadChatProviderShape } from 'vs/workbench/api/common/extHost.protocol'; import { IChatResponseProviderMetadata, IChatResponseFragment, IChatProviderService, IChatMessage } from 'vs/workbench/contrib/chat/common/chatProvider'; +import { AuthenticationSession, AuthenticationSessionsChangeEvent, IAuthenticationProvider, IAuthenticationProviderCreateSessionOptions, IAuthenticationService, INTERNAL_AUTH_PROVIDER_PREFIX } from 'vs/workbench/services/authentication/common/authentication'; import { Extensions, IExtensionFeaturesManagementService, IExtensionFeaturesRegistry } from 'vs/workbench/services/extensionManagement/common/extensionFeatures'; import { IExtHostContext, extHostNamedCustomer } from 'vs/workbench/services/extensions/common/extHostCustomers'; +import { IExtensionService } from 'vs/workbench/services/extensions/common/extensions'; @extHostNamedCustomer(MainContext.MainThreadChatProvider) export class MainThreadChatProvider implements MainThreadChatProviderShape { @@ -28,6 +31,8 @@ export class MainThreadChatProvider implements MainThreadChatProviderShape { @IChatProviderService private readonly _chatProviderService: IChatProviderService, @IExtensionFeaturesManagementService private readonly _extensionFeaturesManagementService: IExtensionFeaturesManagementService, @ILogService private readonly _logService: ILogService, + @IAuthenticationService private readonly _authenticationService: IAuthenticationService, + @IExtensionService private readonly _extensionService: IExtensionService ) { this._proxy = extHostContext.getProxy(ExtHostContext.ExtHostChatProvider); @@ -54,6 +59,7 @@ export class MainThreadChatProvider implements MainThreadChatProviderShape { } } })); + dipsosables.add(this._registerAuthenticationProvider(identifier)); dipsosables.add(Registry.as(Extensions.ExtensionFeaturesRegistry).registerExtensionFeature({ id: `lm-${identifier}`, label: localize('languageModels', "Language Model ({0})", `${identifier}-${metadata.model}`), @@ -93,4 +99,100 @@ export class MainThreadChatProvider implements MainThreadChatProviderShape { return task; } + + private _registerAuthenticationProvider(identifier: string): IDisposable { + const disposables = new DisposableStore(); + // 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)); + 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 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) { + extensionsToUpdateAccess.push({ + extension: extension.identifier, + enabled: false + }); + } + } + this._proxy.$updateAccesslist(extensionsToUpdateAccess); + } + } + })); + disposables.add(this._authenticationService.onDidChangeExtensionSessionAccess(async (e) => { + const allowedExtensions = this._authenticationService.readAllowedExtensions(authProviderId, authAccountId); + const accessList = []; + for (const allowedExtension of allowedExtensions) { + const extension = await this._extensionService.getExtension(allowedExtension.id); + if (extension) { + accessList.push({ + extension: extension.identifier, + enabled: allowedExtension.allowed ?? true + }); + } + } + this._proxy.$updateAccesslist(accessList); + })); + return disposables; + } +} + +// 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(); + onDidChangeSessions: Event = this._onDidChangeSessions.event; + + private _session: AuthenticationSession | undefined; + + constructor(readonly id: string, private readonly accountName: string) { } + + async getSessions(scopes?: string[] | undefined): Promise { + // If there are no scopes and no session that means no extension has requested a session yet + // and the user is simply opening the Account menu. In that case, we should not return any "sessions". + if (scopes === undefined && !this._session) { + return []; + } + if (this._session) { + return [this._session]; + } + return [await this.createSession(scopes || [], {})]; + } + async createSession(scopes: string[], options: IAuthenticationProviderCreateSessionOptions): Promise { + this._session = this._createFakeSession(scopes); + this._onDidChangeSessions.fire({ added: [this._session], changed: [], removed: [] }); + return this._session; + } + removeSession(sessionId: string): Promise { + if (this._session) { + this._onDidChangeSessions.fire({ added: [], changed: [], removed: [this._session!] }); + this._session = undefined; + } + return Promise.resolve(); + } + + private _createFakeSession(scopes: string[]): AuthenticationSession { + return { + id: 'fake-session', + account: { + id: this.id, + label: this.accountName, + }, + 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 d899ada547e..a820c667796 100644 --- a/src/vs/workbench/api/common/extHost.api.impl.ts +++ b/src/vs/workbench/api/common/extHost.api.impl.ts @@ -207,7 +207,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I const extHostProfileContentHandlers = rpcProtocol.set(ExtHostContext.ExtHostProfileContentHandlers, new ExtHostProfileContentHandlers(rpcProtocol)); rpcProtocol.set(ExtHostContext.ExtHostInteractive, new ExtHostInteractive(rpcProtocol, extHostNotebook, extHostDocumentsAndEditors, extHostCommands, extHostLogService)); const extHostInteractiveEditor = rpcProtocol.set(ExtHostContext.ExtHostInlineChat, new ExtHostInteractiveEditor(rpcProtocol, extHostCommands, extHostDocuments, extHostLogService)); - const extHostChatProvider = rpcProtocol.set(ExtHostContext.ExtHostChatProvider, new ExtHostChatProvider(rpcProtocol, extHostLogService)); + const extHostChatProvider = rpcProtocol.set(ExtHostContext.ExtHostChatProvider, new ExtHostChatProvider(rpcProtocol, extHostLogService, extHostAuthentication)); const extHostChatAgents2 = rpcProtocol.set(ExtHostContext.ExtHostChatAgents2, new ExtHostChatAgents2(rpcProtocol, extHostChatProvider, extHostLogService, extHostCommands)); const extHostChatVariables = rpcProtocol.set(ExtHostContext.ExtHostChatVariables, new ExtHostChatVariables(rpcProtocol)); const extHostChat = rpcProtocol.set(ExtHostContext.ExtHostChat, new ExtHostChat(rpcProtocol)); @@ -1399,7 +1399,7 @@ export function createApiFactoryAndRegisterActors(accessor: ServicesAccessor): I }, requestLanguageModelAccess(id, options) { checkProposedApiEnabled(extension, 'chatRequestAccess'); - return extHostChatProvider.requestLanguageModelAccess(extension.identifier, id, options); + return extHostChatProvider.requestLanguageModelAccess(extension, id, options); }, get languageModels() { checkProposedApiEnabled(extension, 'chatRequestAccess'); diff --git a/src/vs/workbench/api/common/extHostAuthentication.ts b/src/vs/workbench/api/common/extHostAuthentication.ts index ca07cbfef64..84ddbd6fa55 100644 --- a/src/vs/workbench/api/common/extHostAuthentication.ts +++ b/src/vs/workbench/api/common/extHostAuthentication.ts @@ -8,6 +8,7 @@ import { Emitter, Event } from 'vs/base/common/event'; import { IMainContext, MainContext, MainThreadAuthenticationShape, ExtHostAuthenticationShape } from 'vs/workbench/api/common/extHost.protocol'; import { Disposable } from 'vs/workbench/api/common/extHostTypes'; import { IExtensionDescription, ExtensionIdentifier } from 'vs/platform/extensions/common/extensions'; +import { INTERNAL_AUTH_PROVIDER_PREFIX } from 'vs/workbench/services/authentication/common/authentication'; interface ProviderWithMetadata { label: string; @@ -106,7 +107,10 @@ export class ExtHostAuthentication implements ExtHostAuthenticationShape { } $onDidChangeAuthenticationSessions(id: string, label: string) { - this._onDidChangeSessions.fire({ provider: { id, label } }); + // Don't fire events for the internal auth providers + if (!id.startsWith(INTERNAL_AUTH_PROVIDER_PREFIX)) { + this._onDidChangeSessions.fire({ provider: { id, label } }); + } return Promise.resolve(); } } diff --git a/src/vs/workbench/api/common/extHostChatProvider.ts b/src/vs/workbench/api/common/extHostChatProvider.ts index b7a91c72e35..eae0f221ffa 100644 --- a/src/vs/workbench/api/common/extHostChatProvider.ts +++ b/src/vs/workbench/api/common/extHostChatProvider.ts @@ -11,9 +11,12 @@ 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 } from 'vs/platform/extensions/common/extensions'; +import { ExtensionIdentifier, ExtensionIdentifierMap, ExtensionIdentifierSet, IRelaxedExtensionDescription } 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'; +import { localize } from 'vs/nls'; +import { INTERNAL_AUTH_PROVIDER_PREFIX } from 'vs/workbench/services/authentication/common/authentication'; type LanguageModelData = { readonly extension: ExtensionIdentifier; @@ -100,6 +103,7 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape { constructor( mainContext: IMainContext, private readonly _logService: ILogService, + private readonly _extHostAuthentication: ExtHostAuthentication, ) { this._proxy = mainContext.getProxy(MainContext.MainThreadChatProvider); } @@ -186,13 +190,15 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape { this._onDidChangeAccess.fire(updated); } - async requestLanguageModelAccess(from: ExtensionIdentifier, languageModelId: string, options?: vscode.LanguageModelAccessOptions): Promise { + async requestLanguageModelAccess(extension: Readonly, 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) { throw new Error('Extension is NOT allowed to make chat requests'); } - const metadata = await this._proxy.$prepareChatAccess(from, languageModelId, options?.justification); + const justification = options?.justification; + const metadata = await this._proxy.$prepareChatAccess(from, languageModelId, justification); if (!metadata) { if (!this._accesslist.get(from)) { @@ -200,6 +206,7 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape { } throw new Error(`Language model '${languageModelId}' NOT found`); } + await this._checkAuthAccess(extension, languageModelId, justification); const that = this; @@ -244,4 +251,22 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape { data.res.handleFragment(chunk); } } + + // BIG HACK: Using AuthenticationProviders to check access to Language Models + private async _checkAuthAccess(from: Readonly, languageModelId: string, detail?: string): Promise { + // This needs to be done in both MainThread & ExtHost ChatProvider + const providerId = INTERNAL_AUTH_PROVIDER_PREFIX + languageModelId; + 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), + } + }); + } catch (err) { + throw new Error('Access to language model has not been granted'); + } + } + } } diff --git a/src/vs/workbench/services/authentication/browser/authenticationService.ts b/src/vs/workbench/services/authentication/browser/authenticationService.ts index 78963a374a1..78e2a0c4137 100644 --- a/src/vs/workbench/services/authentication/browser/authenticationService.ts +++ b/src/vs/workbench/services/authentication/browser/authenticationService.ts @@ -24,7 +24,7 @@ import { Registry } from 'vs/platform/registry/common/platform'; import { ISecretStorageService } from 'vs/platform/secrets/common/secrets'; import { IStorageService, StorageScope, StorageTarget } from 'vs/platform/storage/common/storage'; import { IActivityService, NumberBadge } from 'vs/workbench/services/activity/common/activity'; -import { IAuthenticationCreateSessionOptions, AuthenticationProviderInformation, AuthenticationSession, AuthenticationSessionsChangeEvent, IAuthenticationProvider, IAuthenticationService } from 'vs/workbench/services/authentication/common/authentication'; +import { IAuthenticationCreateSessionOptions, AuthenticationProviderInformation, AuthenticationSession, AuthenticationSessionsChangeEvent, IAuthenticationProvider, IAuthenticationService, AllowedExtension } from 'vs/workbench/services/authentication/common/authentication'; import { IBrowserWorkbenchEnvironmentService } from 'vs/workbench/services/environment/browser/environmentService'; import { IExtensionFeatureTableRenderer, IRenderedData, ITableData, IRowData, IExtensionFeaturesRegistry, Extensions } from 'vs/workbench/services/extensionManagement/common/extensionFeatures'; import { ActivationKind, IExtensionService } from 'vs/workbench/services/extensions/common/extensions'; @@ -107,15 +107,6 @@ export async function getCurrentAuthenticationSessionInfo( return undefined; } -interface AllowedExtension { - id: string; - name: string; - allowed?: boolean; - lastUsed?: number; - // If true, this comes from the product.json - trusted?: boolean; -} - // OAuth2 spec prohibits space in a scope, so use that to join them. const SCOPESLIST_SEPARATOR = ' '; @@ -245,6 +236,9 @@ export class AuthenticationService extends Disposable implements IAuthentication private _onDidChangeDeclaredProviders: Emitter = this._register(new Emitter()); readonly onDidChangeDeclaredProviders: Event = this._onDidChangeDeclaredProviders.event; + private _onDidChangeExtensionSessionAccess: Emitter<{ providerId: string; accountName: string }> = this._register(new Emitter<{ providerId: string; accountName: string }>()); + readonly onDidChangeExtensionSessionAccess: Event<{ providerId: string; accountName: string }> = this._onDidChangeExtensionSessionAccess.event; + constructor( @IActivityService private readonly activityService: IActivityService, @IExtensionService private readonly extensionService: IExtensionService, @@ -816,7 +810,8 @@ export class AuthenticationService extends Disposable implements IAuthentication } } - private readAllowedExtensions(providerId: string, accountName: string): AllowedExtension[] { + // TODO: pull this stuff out into its own service + readAllowedExtensions(providerId: string, accountName: string): AllowedExtension[] { let trustedExtensions: AllowedExtension[] = []; try { const trustedExtensionSrc = this.storageService.get(`${providerId}-${accountName}`, StorageScope.APPLICATION); @@ -926,6 +921,7 @@ export class AuthenticationService extends Disposable implements IAuthentication .filter((item): item is TrustedExtensionsQuickPickItem => item.type !== 'separator') .map(i => i.extension); this.storageService.store(`${authProvider.id}-${accountName}`, JSON.stringify(updatedAllowedList), StorageScope.APPLICATION, StorageTarget.USER); + this._onDidChangeExtensionSessionAccess.fire({ providerId: authProvider.id, accountName }); quickPick.hide(); })); diff --git a/src/vs/workbench/services/authentication/common/authentication.ts b/src/vs/workbench/services/authentication/common/authentication.ts index 0ff2179bfe2..eaeaffc56ac 100644 --- a/src/vs/workbench/services/authentication/common/authentication.ts +++ b/src/vs/workbench/services/authentication/common/authentication.ts @@ -5,6 +5,11 @@ import { Event } from 'vs/base/common/event'; import { createDecorator } from 'vs/platform/instantiation/common/instantiation'; +/** + * Use this if you don't want the onDidChangeSessions event to fire in the extension host + */ +export const INTERNAL_AUTH_PROVIDER_PREFIX = '__'; + export interface AuthenticationSessionAccount { label: string; id: string; @@ -34,6 +39,20 @@ export interface IAuthenticationCreateSessionOptions { activateImmediate?: boolean; } +export interface AllowedExtension { + id: string; + name: string; + /** + * If true or undefined, the extension is allowed to use the account + * If false, the extension is not allowed to use the account + * TODO: undefined shouldn't be a valid value, but it is for now + */ + allowed?: boolean; + lastUsed?: number; + // If true, this comes from the product.json + trusted?: boolean; +} + export const IAuthenticationService = createDecorator('IAuthenticationService'); export interface IAuthenticationService { @@ -58,6 +77,7 @@ export interface IAuthenticationService { readonly onDidUnregisterAuthenticationProvider: Event; readonly onDidChangeSessions: Event<{ providerId: string; label: string; event: AuthenticationSessionsChangeEvent }>; + readonly onDidChangeExtensionSessionAccess: Event<{ providerId: string; accountName: string }>; // TODO completely remove this property declaredProviders: AuthenticationProviderInformation[]; @@ -70,6 +90,7 @@ export interface IAuthenticationService { removeSession(providerId: string, sessionId: string): Promise; manageTrustedExtensionsForAccount(providerId: string, accountName: string): Promise; + readAllowedExtensions(providerId: string, accountName: string): AllowedExtension[]; removeAccountSessions(providerId: string, accountName: string, sessions: AuthenticationSession[]): Promise; }