diff --git a/src/vs/workbench/contrib/chat/browser/chatParticipantContributions.ts b/src/vs/workbench/contrib/chat/browser/chatParticipantContributions.ts index c89a6a4d71226a..0b0eb44b222ac1 100644 --- a/src/vs/workbench/contrib/chat/browser/chatParticipantContributions.ts +++ b/src/vs/workbench/contrib/chat/browser/chatParticipantContributions.ts @@ -201,6 +201,12 @@ export class ChatExtensionPointHandler implements IWorkbenchContribution { store.add(this.registerDefaultParticipantView(providerDescriptor)); } + if (providerDescriptor.when && !isProposedApiEnabled(extension.description, 'chatParticipantPrivate')) { + this.logService.error(`Extension '${extension.description.identifier.value}' CANNOT use API proposal: chatParticipantPrivate.`); + continue; + + } + store.add(this._chatAgentService.registerAgent( providerDescriptor.id, { @@ -210,6 +216,7 @@ export class ChatExtensionPointHandler implements IWorkbenchContribution { extensionDisplayName: extension.description.displayName ?? extension.description.name, id: providerDescriptor.id, description: providerDescriptor.description, + when: providerDescriptor.when, metadata: { isSticky: providerDescriptor.isSticky, sampleRequest: providerDescriptor.sampleRequest, diff --git a/src/vs/workbench/contrib/chat/common/chatAgents.ts b/src/vs/workbench/contrib/chat/common/chatAgents.ts index 43cf0f8e680a9c..da0e5eba9d97ab 100644 --- a/src/vs/workbench/contrib/chat/common/chatAgents.ts +++ b/src/vs/workbench/contrib/chat/common/chatAgents.ts @@ -61,6 +61,7 @@ export interface IChatAgentData { name: string; fullName?: string; description?: string; + when?: string; extensionId: ExtensionIdentifier; extensionPublisherId: string; /** This is the extension publisher id, or, in the case of a dynamically registered participant (remote agent), whatever publisher name we have for it */ @@ -193,7 +194,7 @@ export class ChatAgentService implements IChatAgentService { declare _serviceBrand: undefined; - private _agents: IChatAgentEntry[] = []; + private _agents = new Map(); private readonly _onDidChangeAgents = new Emitter(); readonly onDidChangeAgents: Event = this._onDidChangeAgents.event; @@ -221,15 +222,15 @@ export class ChatAgentService implements IChatAgentService { } }; const entry = { data }; - this._agents.push(entry); + this._agents.set(id, entry); return toDisposable(() => { - this._agents = this._agents.filter(a => a !== entry); + this._agents.delete(id); this._onDidChangeAgents.fire(undefined); }); } registerAgentImplementation(id: string, agentImpl: IChatAgentImplementation): IDisposable { - const entry = this._getAgentEntry(id); + const entry = this._agents.get(id); if (!entry) { throw new Error(`Unknown agent: ${JSON.stringify(id)}`); } @@ -258,11 +259,11 @@ export class ChatAgentService implements IChatAgentService { registerDynamicAgent(data: IChatAgentData, agentImpl: IChatAgentImplementation): IDisposable { data.isDynamic = true; const agent = { data, impl: agentImpl }; - this._agents.push(agent); + this._agents.set(data.id, agent); this._onDidChangeAgents.fire(new MergedChatAgent(data, agentImpl)); return toDisposable(() => { - this._agents = this._agents.filter(a => a !== agent); + this._agents.delete(data.id); this._onDidChangeAgents.fire(undefined); }); } @@ -281,7 +282,7 @@ export class ChatAgentService implements IChatAgentService { } updateAgent(id: string, updateMetadata: IChatAgentMetadata): void { - const agent = this._getAgentEntry(id); + const agent = this._agents.get(id); if (!agent?.impl) { throw new Error(`No activated agent with id ${JSON.stringify(id)} registered`); } @@ -302,28 +303,41 @@ export class ChatAgentService implements IChatAgentService { return Iterable.find(this._agents.values(), a => !!a.data.metadata.isSecondary)?.data; } - private _getAgentEntry(id: string): IChatAgentEntry | undefined { - return this._agents.find(a => a.data.id === id); + getAgent(id: string): IChatAgentData | undefined { + if (!this._agentIsEnabled(id)) { + return; + } + + return this._agents.get(id)?.data; } - getAgent(id: string): IChatAgentData | undefined { - return this._getAgentEntry(id)?.data; + private _agentIsEnabled(id: string): boolean { + const entry = this._agents.get(id); + return !entry?.data.when || this.contextKeyService.contextMatchesRules(ContextKeyExpr.deserialize(entry.data.when)); } getAgentByFullyQualifiedId(id: string): IChatAgentData | undefined { - return this._agents.find(a => getFullyQualifiedId(a.data) === id)?.data; + const agent = Iterable.find(this._agents.values(), a => getFullyQualifiedId(a.data) === id)?.data; + if (agent && !this._agentIsEnabled(agent.id)) { + return; + } + + return agent; } /** * Returns all agent datas that exist- static registered and dynamic ones. */ getAgents(): IChatAgentData[] { - return this._agents.map(entry => entry.data); + return Array.from(this._agents.values()) + .map(entry => entry.data) + .filter(a => this._agentIsEnabled(a.id)); } getActivatedAgents(): IChatAgent[] { return Array.from(this._agents.values()) .filter(a => !!a.impl) + .filter(a => this._agentIsEnabled(a.data.id)) .map(a => new MergedChatAgent(a.data, a.impl!)); } @@ -332,7 +346,7 @@ export class ChatAgentService implements IChatAgentService { } async invokeAgent(id: string, request: IChatAgentRequest, progress: (part: IChatProgress) => void, history: IChatAgentHistoryEntry[], token: CancellationToken): Promise { - const data = this._getAgentEntry(id); + const data = this._agents.get(id); if (!data?.impl) { throw new Error(`No activated agent with id "${id}"`); } @@ -341,7 +355,7 @@ export class ChatAgentService implements IChatAgentService { } async getFollowups(id: string, request: IChatAgentRequest, result: IChatAgentResult, history: IChatAgentHistoryEntry[], token: CancellationToken): Promise { - const data = this._getAgentEntry(id); + const data = this._agents.get(id); if (!data?.impl) { throw new Error(`No activated agent with id "${id}"`); } diff --git a/src/vs/workbench/contrib/chat/common/chatParticipantContribTypes.ts b/src/vs/workbench/contrib/chat/common/chatParticipantContribTypes.ts index e6d695c3824537..638dce2633c642 100644 --- a/src/vs/workbench/contrib/chat/common/chatParticipantContribTypes.ts +++ b/src/vs/workbench/contrib/chat/common/chatParticipantContribTypes.ts @@ -18,6 +18,7 @@ export interface IRawChatParticipantContribution { id: string; name: string; fullName: string; + when?: string; description?: string; isDefault?: boolean; isSticky?: boolean; diff --git a/src/vs/workbench/contrib/chat/test/common/chatAgents.test.ts b/src/vs/workbench/contrib/chat/test/common/chatAgents.test.ts new file mode 100644 index 00000000000000..21bd0b3930ffab --- /dev/null +++ b/src/vs/workbench/contrib/chat/test/common/chatAgents.test.ts @@ -0,0 +1,98 @@ +/*--------------------------------------------------------------------------------------------- + * Copyright (c) Microsoft Corporation. All rights reserved. + * Licensed under the MIT License. See License.txt in the project root for license information. + *--------------------------------------------------------------------------------------------*/ + +import { MockObject, mockObject } from 'vs/base/test/common/mock'; +import { ensureNoDisposablesAreLeakedInTestSuite } from 'vs/base/test/common/utils'; +import { IContextKeyService } from 'vs/platform/contextkey/common/contextkey'; +import { ExtensionIdentifier } from 'vs/platform/extensions/common/extensions'; +import { ChatAgentService, IChatAgentData, IChatAgentImplementation, IChatAgentService } from 'vs/workbench/contrib/chat/common/chatAgents'; +import * as assert from 'assert'; + +const testAgentId = 'testAgent'; +const testAgentData: IChatAgentData = { + id: testAgentId, + name: 'Test Agent', + extensionDisplayName: '', + extensionId: new ExtensionIdentifier(''), + extensionPublisherId: '', + locations: [], + metadata: {}, + slashCommands: [], +}; + +suite('ChatAgents', function () { + const store = ensureNoDisposablesAreLeakedInTestSuite(); + + let chatAgentService: IChatAgentService; + let contextKeyService: MockObject; + setup(() => { + contextKeyService = mockObject()(); + chatAgentService = new ChatAgentService(contextKeyService as any); + }); + + test('registerAgent', async () => { + assert.strictEqual(chatAgentService.getAgents().length, 0); + + + const agentRegistration = chatAgentService.registerAgent(testAgentId, testAgentData); + + assert.strictEqual(chatAgentService.getAgents().length, 1); + assert.strictEqual(chatAgentService.getAgents()[0].id, testAgentId); + + assert.throws(() => chatAgentService.registerAgent(testAgentId, testAgentData)); + + agentRegistration.dispose(); + assert.strictEqual(chatAgentService.getAgents().length, 0); + }); + + test('agent when clause', async () => { + assert.strictEqual(chatAgentService.getAgents().length, 0); + + store.add(chatAgentService.registerAgent(testAgentId, { + ...testAgentData, + when: 'myKey' + })); + assert.strictEqual(chatAgentService.getAgents().length, 0); + + contextKeyService.contextMatchesRules.returns(true); + assert.strictEqual(chatAgentService.getAgents().length, 1); + }); + + suite('registerAgentImplementation', function () { + const agentImpl: IChatAgentImplementation = { + invoke: async () => { return {}; }, + provideFollowups: async () => { return []; }, + }; + + test('should register an agent implementation', () => { + store.add(chatAgentService.registerAgent(testAgentId, testAgentData)); + store.add(chatAgentService.registerAgentImplementation(testAgentId, agentImpl)); + + const agents = chatAgentService.getActivatedAgents(); + assert.strictEqual(agents.length, 1); + assert.strictEqual(agents[0].id, testAgentId); + }); + + test('can dispose an agent implementation', () => { + store.add(chatAgentService.registerAgent(testAgentId, testAgentData)); + const implRegistration = chatAgentService.registerAgentImplementation(testAgentId, agentImpl); + implRegistration.dispose(); + + const agents = chatAgentService.getActivatedAgents(); + assert.strictEqual(agents.length, 0); + }); + + test('should throw error if agent does not exist', () => { + assert.throws(() => chatAgentService.registerAgentImplementation('nonexistentAgent', agentImpl)); + }); + + test('should throw error if agent already has an implementation', () => { + store.add(chatAgentService.registerAgent(testAgentId, testAgentData)); + store.add(chatAgentService.registerAgentImplementation(testAgentId, agentImpl)); + + assert.throws(() => chatAgentService.registerAgentImplementation(testAgentId, agentImpl)); + }); + }); +});