From 9bd6ee481029ae84b6cd67ab4aa8ec2d86006809 Mon Sep 17 00:00:00 2001
From: Rob Lourens <roblourens@gmail.com>
Date: Mon, 26 Feb 2024 19:02:16 -0300
Subject: [PATCH] Warm up slash command cache (#206290)

Fix #206050
---
 .../workbench/contrib/chat/common/chatAgents.ts   | 15 +++++++++------
 .../contrib/chat/common/chatServiceImpl.ts        | 14 +++++++++++++-
 2 files changed, 22 insertions(+), 7 deletions(-)

diff --git a/src/vs/workbench/contrib/chat/common/chatAgents.ts b/src/vs/workbench/contrib/chat/common/chatAgents.ts
index 4ee4c07f0cf6c..b1ef64dae9d97 100644
--- a/src/vs/workbench/contrib/chat/common/chatAgents.ts
+++ b/src/vs/workbench/contrib/chat/common/chatAgents.ts
@@ -107,7 +107,10 @@ export const IChatAgentService = createDecorator<IChatAgentService>('chatAgentSe
 
 export interface IChatAgentService {
 	_serviceBrand: undefined;
-	readonly onDidChangeAgents: Event<void>;
+	/**
+	 * undefined when an agent was removed
+	 */
+	readonly onDidChangeAgents: Event<IChatAgent | undefined>;
 	registerAgent(agent: IChatAgent): IDisposable;
 	invokeAgent(id: string, request: IChatAgentRequest, progress: (part: IChatProgress) => void, history: IChatAgentHistoryEntry[], token: CancellationToken): Promise<IChatAgentResult>;
 	getFollowups(id: string, request: IChatAgentRequest, result: IChatAgentResult, token: CancellationToken): Promise<IChatFollowup[]>;
@@ -127,8 +130,8 @@ export class ChatAgentService extends Disposable implements IChatAgentService {
 
 	private readonly _agents = new Map<string, { agent: IChatAgent }>();
 
-	private readonly _onDidChangeAgents = this._register(new Emitter<void>());
-	readonly onDidChangeAgents: Event<void> = this._onDidChangeAgents.event;
+	private readonly _onDidChangeAgents = this._register(new Emitter<IChatAgent | undefined>());
+	readonly onDidChangeAgents: Event<IChatAgent | undefined> = this._onDidChangeAgents.event;
 
 	override dispose(): void {
 		super.dispose();
@@ -140,11 +143,11 @@ export class ChatAgentService extends Disposable implements IChatAgentService {
 			throw new Error(`Already registered an agent with id ${agent.id}`);
 		}
 		this._agents.set(agent.id, { agent });
-		this._onDidChangeAgents.fire();
+		this._onDidChangeAgents.fire(agent);
 
 		return toDisposable(() => {
 			if (this._agents.delete(agent.id)) {
-				this._onDidChangeAgents.fire();
+				this._onDidChangeAgents.fire(undefined);
 			}
 		});
 	}
@@ -155,7 +158,7 @@ export class ChatAgentService extends Disposable implements IChatAgentService {
 			throw new Error(`No agent with id ${id} registered`);
 		}
 		data.agent.metadata = { ...data.agent.metadata, ...updateMetadata };
-		this._onDidChangeAgents.fire();
+		this._onDidChangeAgents.fire(data.agent);
 	}
 
 	getDefaultAgent(): IChatAgent | undefined {
diff --git a/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts b/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts
index 89081fe437d25..68cc8da8a53b2 100644
--- a/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts
+++ b/src/vs/workbench/contrib/chat/common/chatServiceImpl.ts
@@ -20,7 +20,7 @@ import { Progress } from 'vs/platform/progress/common/progress';
 import { IStorageService, StorageScope, StorageTarget } from 'vs/platform/storage/common/storage';
 import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
 import { IWorkspaceContextService } from 'vs/platform/workspace/common/workspace';
-import { IChatAgentRequest, IChatAgentResult, IChatAgentService } from 'vs/workbench/contrib/chat/common/chatAgents';
+import { IChatAgent, IChatAgentRequest, IChatAgentResult, IChatAgentService } from 'vs/workbench/contrib/chat/common/chatAgents';
 import { CONTEXT_PROVIDER_EXISTS } from 'vs/workbench/contrib/chat/common/chatContextKeys';
 import { ChatModel, ChatModelInitState, ChatRequestModel, ChatWelcomeMessageModel, IChatModel, IChatRequestVariableData, ISerializableChatData, ISerializableChatsData, getHistoryEntriesFromModel, updateRanges } from 'vs/workbench/contrib/chat/common/chatModel';
 import { ChatRequestAgentPart, ChatRequestAgentSubcommandPart, ChatRequestSlashCommandPart, IParsedChatRequest, getPromptText } from 'vs/workbench/contrib/chat/common/chatParserTypes';
@@ -193,6 +193,12 @@ export class ChatService extends Disposable implements IChatService {
 		}
 
 		this._register(storageService.onWillSaveState(() => this.saveState()));
+
+		this._register(Event.debounce(this.chatAgentService.onDidChangeAgents, () => { }, 500)(() => {
+			for (const model of this._sessionModels.values()) {
+				this.warmSlashCommandCache(model);
+			}
+		}));
 	}
 
 	private saveState(): void {
@@ -358,9 +364,15 @@ export class ChatService extends Disposable implements IChatService {
 		this.initializeSession(model, CancellationToken.None);
 	}
 
+	private warmSlashCommandCache(model: IChatModel, agent?: IChatAgent) {
+		const agents = agent ? [agent] : this.chatAgentService.getAgents();
+		agents.forEach(agent => agent.provideSlashCommands(model, [], CancellationToken.None));
+	}
+
 	private async initializeSession(model: ChatModel, token: CancellationToken): Promise<void> {
 		try {
 			this.trace('initializeSession', `Initialize session ${model.sessionId}`);
+			this.warmSlashCommandCache(model);
 			model.startInitialize();
 			await this.extensionService.activateByEvent(`onInteractiveSession:${model.providerId}`);