Add "when" clause to chat participants

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,

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<string, IChatAgentEntry>();
 
 	private readonly _onDidChangeAgents = new Emitter<IChatAgent | undefined>();
 	readonly onDidChangeAgents: Event<IChatAgent | undefined> = 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<IChatAgentResult> {
-		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<IChatFollowup[]> {
-		const data = this._getAgentEntry(id);
+		const data = this._agents.get(id);
 		if (!data?.impl) {
 			throw new Error(`No activated agent with id "${id}"`);
 		}

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;

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<IContextKeyService>;
+	setup(() => {
+		contextKeyService = mockObject<IContextKeyService>()();
+		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));
+		});
+	});
+});