|
3 | 3 | * Licensed under the MIT License. See License.txt in the project root for license information. |
4 | 4 | *--------------------------------------------------------------------------------------------*/ |
5 | 5 |
|
| 6 | +import { timeout } from 'vs/base/common/async'; |
6 | 7 | import { CancellationToken } from 'vs/base/common/cancellation'; |
7 | 8 | import { Emitter, Event } from 'vs/base/common/event'; |
8 | 9 | import { IMarkdownString } from 'vs/base/common/htmlContent'; |
9 | 10 | import { Iterable } from 'vs/base/common/iterator'; |
10 | 11 | import { IDisposable, toDisposable } from 'vs/base/common/lifecycle'; |
| 12 | +import { IObservable } from 'vs/base/common/observable'; |
| 13 | +import { observableValue } from 'vs/base/common/observableInternal/base'; |
| 14 | +import { equalsIgnoreCase } from 'vs/base/common/strings'; |
11 | 15 | import { ThemeIcon } from 'vs/base/common/themables'; |
12 | 16 | import { URI } from 'vs/base/common/uri'; |
13 | 17 | import { ProviderResult } from 'vs/editor/common/languages'; |
14 | 18 | import { ContextKeyExpr, IContextKey, IContextKeyService } from 'vs/platform/contextkey/common/contextkey'; |
15 | 19 | import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions'; |
16 | 20 | import { createDecorator } from 'vs/platform/instantiation/common/instantiation'; |
| 21 | +import { ILogService } from 'vs/platform/log/common/log'; |
| 22 | +import { IProductService } from 'vs/platform/product/common/productService'; |
| 23 | +import { asJson, IRequestService } from 'vs/platform/request/common/request'; |
| 24 | +import { IStorageService, StorageScope, StorageTarget } from 'vs/platform/storage/common/storage'; |
17 | 25 | import { CONTEXT_CHAT_ENABLED } from 'vs/workbench/contrib/chat/common/chatContextKeys'; |
18 | 26 | import { IChatProgressResponseContent, IChatRequestVariableData } from 'vs/workbench/contrib/chat/common/chatModel'; |
19 | 27 | import { IRawChatCommandContribution, RawChatParticipantLocation } from 'vs/workbench/contrib/chat/common/chatParticipantContribTypes'; |
@@ -351,3 +359,95 @@ export class MergedChatAgent implements IChatAgent { |
351 | 359 | return undefined; |
352 | 360 | } |
353 | 361 | } |
| 362 | + |
| 363 | +export const IChatAgentNameService = createDecorator<IChatAgentNameService>('chatAgentNameService'); |
| 364 | + |
| 365 | +type IChatParticipantRegistry = { [name: string]: string[] }; |
| 366 | + |
| 367 | +interface IChatParticipantRegistryResponse { |
| 368 | + readonly version: number; |
| 369 | + readonly restrictedChatParticipants: IChatParticipantRegistry; |
| 370 | +} |
| 371 | + |
| 372 | +export interface IChatAgentNameService { |
| 373 | + _serviceBrand: undefined; |
| 374 | + getAgentNameRestriction(chatAgentData: IChatAgentData): IObservable<boolean>; |
| 375 | +} |
| 376 | + |
| 377 | +export class ChatAgentNameService implements IChatAgentNameService { |
| 378 | + |
| 379 | + private static readonly StorageKey = 'chat.participantNameRegistry'; |
| 380 | + |
| 381 | + declare _serviceBrand: undefined; |
| 382 | + |
| 383 | + private readonly url!: string; |
| 384 | + private registry = observableValue<IChatParticipantRegistry>(this, Object.create(null)); |
| 385 | + private disposed = false; |
| 386 | + |
| 387 | + constructor( |
| 388 | + @IProductService productService: IProductService, |
| 389 | + @IRequestService private readonly requestService: IRequestService, |
| 390 | + @ILogService private readonly logService: ILogService, |
| 391 | + @IStorageService private readonly storageService: IStorageService |
| 392 | + ) { |
| 393 | + if (!productService.chatParticipantRegistry) { |
| 394 | + return; |
| 395 | + } |
| 396 | + |
| 397 | + this.url = productService.chatParticipantRegistry; |
| 398 | + |
| 399 | + const raw = storageService.get(ChatAgentNameService.StorageKey, StorageScope.APPLICATION); |
| 400 | + |
| 401 | + try { |
| 402 | + this.registry.set(JSON.parse(raw ?? '{}'), undefined); |
| 403 | + } catch (err) { |
| 404 | + storageService.remove(ChatAgentNameService.StorageKey, StorageScope.APPLICATION); |
| 405 | + } |
| 406 | + |
| 407 | + this.refresh(); |
| 408 | + } |
| 409 | + |
| 410 | + private refresh(): void { |
| 411 | + if (this.disposed) { |
| 412 | + return; |
| 413 | + } |
| 414 | + |
| 415 | + this.update() |
| 416 | + .catch(err => this.logService.warn('Failed to fetch chat participant registry', err)) |
| 417 | + .then(() => timeout(5 * 60 * 1000)) // every 5 minutes |
| 418 | + .then(() => this.refresh()); |
| 419 | + } |
| 420 | + |
| 421 | + private async update(): Promise<void> { |
| 422 | + const context = await this.requestService.request({ type: 'GET', url: this.url }, CancellationToken.None); |
| 423 | + |
| 424 | + if (context.res.statusCode !== 200) { |
| 425 | + throw new Error('Could not get extensions report.'); |
| 426 | + } |
| 427 | + |
| 428 | + const result = await asJson<IChatParticipantRegistryResponse>(context); |
| 429 | + |
| 430 | + if (!result || result.version !== 1) { |
| 431 | + throw new Error('Unexpected chat participant registry response.'); |
| 432 | + } |
| 433 | + |
| 434 | + const registry = result.restrictedChatParticipants; |
| 435 | + this.registry.set(registry, undefined); |
| 436 | + this.storageService.store(ChatAgentNameService.StorageKey, JSON.stringify(registry), StorageScope.APPLICATION, StorageTarget.MACHINE); |
| 437 | + } |
| 438 | + |
| 439 | + getAgentNameRestriction(chatAgentData: IChatAgentData): IObservable<boolean> { |
| 440 | + const allowList = this.registry.map<string[] | undefined>(registry => registry[chatAgentData.name.toLowerCase()]); |
| 441 | + return allowList.map(allowList => { |
| 442 | + if (!allowList) { |
| 443 | + return true; |
| 444 | + } |
| 445 | + |
| 446 | + return allowList.some(id => equalsIgnoreCase(id, id.includes('.') ? chatAgentData.extensionId.value : chatAgentData.extensionPublisher)); |
| 447 | + }); |
| 448 | + } |
| 449 | + |
| 450 | + dispose() { |
| 451 | + this.disposed = true; |
| 452 | + } |
| 453 | +} |
0 commit comments