From 4ba66bf35e20a55e7261702371db549f2ad396bd Mon Sep 17 00:00:00 2001
From: Rob Lourens <roblourens@gmail.com>
Date: Fri, 16 Feb 2024 16:09:04 +0000
Subject: [PATCH] Chat API: add 'command' to response history (#205377)

---
 .../src/singlefolder-tests/chat.test.ts       | 39 ++++++++++++++-----
 .../api/common/extHostChatAgents2.ts          |  2 +-
 src/vs/workbench/api/common/extHostTypes.ts   |  3 +-
 .../vscode.proposed.chatAgents2.d.ts          |  2 +
 4 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/extensions/vscode-api-tests/src/singlefolder-tests/chat.test.ts b/extensions/vscode-api-tests/src/singlefolder-tests/chat.test.ts
index 6c0406aadbb..bdb033322eb 100644
--- a/extensions/vscode-api-tests/src/singlefolder-tests/chat.test.ts
+++ b/extensions/vscode-api-tests/src/singlefolder-tests/chat.test.ts
@@ -5,8 +5,8 @@
 
 import * as assert from 'assert';
 import 'mocha';
-import { CancellationToken, chat, ChatAgentRequest, ChatAgentResult2, ChatVariableLevel, Disposable, interactive, InteractiveSession, ProviderResult } from 'vscode';
-import { assertNoRpc, closeAllEditors, DeferredPromise, disposeAll } from '../utils';
+import { CancellationToken, ChatAgentContext, ChatAgentRequest, ChatAgentResult2, ChatVariableLevel, Disposable, Event, EventEmitter, InteractiveSession, ProviderResult, chat, interactive } from 'vscode';
+import { DeferredPromise, assertNoRpc, closeAllEditors, disposeAll } from '../utils';
 
 suite('chat', () => {
 
@@ -22,6 +22,15 @@ suite('chat', () => {
 	});
 
 	function getDeferredForRequest(): DeferredPromise<ChatAgentRequest> {
+		const deferred = new DeferredPromise<ChatAgentRequest>();
+		disposables.push(setupAgent()(request => deferred.complete(request.request)));
+
+		return deferred;
+	}
+
+	function setupAgent(): Event<{ request: ChatAgentRequest; context: ChatAgentContext }> {
+		const emitter = new EventEmitter<{ request: ChatAgentRequest; context: ChatAgentContext }>();
+		disposables.push();
 		disposables.push(interactive.registerInteractiveSessionProvider('provider', {
 			prepareSession: (_token: CancellationToken): ProviderResult<InteractiveSession> => {
 				return {
@@ -31,9 +40,8 @@ suite('chat', () => {
 			},
 		}));
 
-		const deferred = new DeferredPromise<ChatAgentRequest>();
-		const agent = chat.createChatAgent('agent', (request, _context, _progress, _token) => {
-			deferred.complete(request);
+		const agent = chat.createChatAgent('agent', (request, context, _progress, _token) => {
+			emitter.fire({ request, context });
 			return null;
 		});
 		agent.isDefault = true;
@@ -43,15 +51,26 @@ suite('chat', () => {
 			}
 		};
 		disposables.push(agent);
-		return deferred;
+		return emitter.event;
 	}
 
 	test('agent and slash command', async () => {
-		const deferred = getDeferredForRequest();
+		const onRequest = setupAgent();
 		interactive.sendInteractiveRequestToProvider('provider', { message: '@agent /hello friend' });
-		const request = await deferred.p;
-		assert.deepStrictEqual(request.command, 'hello');
-		assert.strictEqual(request.prompt, 'friend');
+
+		let i = 0;
+		onRequest(request => {
+			if (i === 0) {
+				assert.deepStrictEqual(request.request.command, 'hello');
+				assert.strictEqual(request.request.prompt, 'friend');
+				i++;
+				interactive.sendInteractiveRequestToProvider('provider', { message: '@agent /hello friend' });
+			} else {
+				assert.strictEqual(request.context.history.length, 1);
+				assert.strictEqual(request.context.history[0].agent.agent, 'agent');
+				assert.strictEqual(request.context.history[0].command, 'hello');
+			}
+		});
 	});
 
 	test('agent and variable', async () => {
diff --git a/src/vs/workbench/api/common/extHostChatAgents2.ts b/src/vs/workbench/api/common/extHostChatAgents2.ts
index 1f40b04c924..399ff60d030 100644
--- a/src/vs/workbench/api/common/extHostChatAgents2.ts
+++ b/src/vs/workbench/api/common/extHostChatAgents2.ts
@@ -239,7 +239,7 @@ export class ExtHostChatAgents2 implements ExtHostChatAgentsShape2 {
 
 			// RESPONSE turn
 			const parts = coalesce(h.response.map(r => typeConvert.ChatResponsePart.from(r, this.commands.converter)));
-			res.push(new extHostTypes.ChatAgentResponseTurn(parts, result, { extensionId: '', agent: h.request.agentId }));
+			res.push(new extHostTypes.ChatAgentResponseTurn(parts, result, { extensionId: '', agent: h.request.agentId }, h.request.command));
 		}
 
 		return res;
diff --git a/src/vs/workbench/api/common/extHostTypes.ts b/src/vs/workbench/api/common/extHostTypes.ts
index 906dc75413e..611314d600f 100644
--- a/src/vs/workbench/api/common/extHostTypes.ts
+++ b/src/vs/workbench/api/common/extHostTypes.ts
@@ -4274,7 +4274,8 @@ export class ChatAgentResponseTurn implements vscode.ChatAgentResponseTurn {
 	constructor(
 		readonly response: ReadonlyArray<ChatResponseTextPart | ChatResponseMarkdownPart | ChatResponseFileTreePart | ChatResponseAnchorPart | ChatResponseCommandButtonPart>,
 		readonly result: vscode.ChatAgentResult2,
-		readonly agent: { extensionId: string; agent: string }
+		readonly agent: { extensionId: string; agent: string },
+		readonly command?: string
 	) { }
 }
 
diff --git a/src/vscode-dts/vscode.proposed.chatAgents2.d.ts b/src/vscode-dts/vscode.proposed.chatAgents2.d.ts
index 63cad62d20a..93938b8ba68 100644
--- a/src/vscode-dts/vscode.proposed.chatAgents2.d.ts
+++ b/src/vscode-dts/vscode.proposed.chatAgents2.d.ts
@@ -54,6 +54,8 @@ declare module 'vscode' {
 		 */
 		readonly agent: { readonly extensionId: string; readonly agent: string };
 
+		readonly command?: string;
+
 		private constructor(response: ReadonlyArray<ChatResponseTextPart | ChatResponseMarkdownPart | ChatResponseFileTreePart | ChatResponseAnchorPart | ChatResponseCommandButtonPart>, result: ChatAgentResult2, agentId: { extensionId: string; agent: string });
 	}