Add ChatAgentResponseStream as a more explict alternative to Progess<ChatAgentProgress>

Author: jrieken

src/vs/workbench/api/common/extHostChatAgents2.ts

@@ -8,22 +8,143 @@ import { DeferredPromise, raceCancellation } from 'vs/base/common/async';
 import { CancellationToken } from 'vs/base/common/cancellation';
 import { toErrorMessage } from 'vs/base/common/errorMessage';
 import { Emitter } from 'vs/base/common/event';
-import { IMarkdownString } from 'vs/base/common/htmlContent';
+import { IMarkdownString, MarkdownString } from 'vs/base/common/htmlContent';
 import { StopWatch } from 'vs/base/common/stopwatch';
 import { URI } from 'vs/base/common/uri';
 import { localize } from 'vs/nls';
 import { IExtensionDescription } from 'vs/platform/extensions/common/extensions';
 import { ILogService } from 'vs/platform/log/common/log';
-import { Progress } from 'vs/platform/progress/common/progress';
 import { ExtHostChatAgentsShape2, IChatAgentCompletionItem, IChatAgentHistoryEntryDto, IMainContext, MainContext, MainThreadChatAgentsShape2 } from 'vs/workbench/api/common/extHost.protocol';
 import { ExtHostChatProvider } from 'vs/workbench/api/common/extHostChatProvider';
 import * as typeConvert from 'vs/workbench/api/common/extHostTypeConverters';
 import * as extHostTypes from 'vs/workbench/api/common/extHostTypes';
 import { IChatAgentCommand, IChatAgentRequest, IChatAgentResult } from 'vs/workbench/contrib/chat/common/chatAgents';
-import { IChatFollowup, IChatReplyFollowup, IChatUserActionEvent, InteractiveSessionVoteDirection } from 'vs/workbench/contrib/chat/common/chatService';
+import { IChatFollowup, IChatProgress, IChatReplyFollowup, IChatUserActionEvent, InteractiveSessionVoteDirection } from 'vs/workbench/contrib/chat/common/chatService';
 import { checkProposedApiEnabled, isProposedApiEnabled } from 'vs/workbench/services/extensions/common/extensions';
+import { Dto } from 'vs/workbench/services/extensions/common/proxyIdentifier';
 import type * as vscode from 'vscode';
 
+class ChatAgentResponseStream {
+
+	private _stopWatch = StopWatch.create(false);
+	private _isClosed: boolean = false;
+	private _firstProgress: number | undefined;
+	private _apiObject: vscode.ChatAgentExtendedResponseStream | undefined;
+
+	constructor(
+		private readonly _extension: IExtensionDescription,
+		private readonly _request: IChatAgentRequest,
+		private readonly _proxy: MainThreadChatAgentsShape2,
+		@ILogService private readonly _logService: ILogService,
+	) { }
+
+	close() {
+		this._isClosed = true;
+	}
+
+	get timings() {
+		return {
+			firstProgress: this._firstProgress,
+			totalElapsed: this._stopWatch.elapsed()
+		};
+	}
+
+	get apiObject() {
+
+		if (!this._apiObject) {
+
+			const that = this;
+			this._stopWatch.reset();
+
+			function throwIfDone(source: Function | undefined) {
+				if (that._isClosed) {
+					const err = new Error('Response stream has been closed');
+					Error.captureStackTrace(err, source);
+					throw err;
+				}
+			}
+
+			const _report = (progress: Dto<IChatProgress>) => {
+				// Measure the time to the first progress update with real markdown content
+				if (typeof this._firstProgress === 'undefined' && 'content' in progress) {
+					this._firstProgress = this._stopWatch.elapsed();
+				}
+				this._proxy.$handleProgressChunk(this._request.requestId, progress);
+			};
+
+			this._apiObject = {
+				markdown(value) {
+					throwIfDone(this.markdown);
+					_report({
+						kind: 'markdownContent',
+						content: typeConvert.MarkdownString.from(value)
+					});
+					return this;
+				},
+				text(value) {
+					throwIfDone(this.text);
+					this.markdown(new MarkdownString().appendText(value));
+					return this;
+				},
+				files(value) {
+					throwIfDone(this.files);
+					_report({
+						kind: 'treeData',
+						treeData: value
+					});
+					return this;
+				},
+				anchor(value) {
+					throwIfDone(this.anchor);
+					_report({
+						kind: 'inlineReference',
+						inlineReference: !URI.isUri(value) ? typeConvert.Location.from(<vscode.Location>value) : value
+					});
+					return this;
+				},
+				progress(value) {
+					throwIfDone(this.progress);
+					_report({
+						kind: 'progressMessage',
+						content: new MarkdownString(value)
+					});
+					return this;
+				},
+				reference(value) {
+					throwIfDone(this.reference);
+					_report({
+						kind: 'reference',
+						reference: !URI.isUri(value) ? typeConvert.Location.from(<vscode.Location>value) : value
+					});
+					return this;
+				},
+				// annotation(value) {
+				// 	_report(value);
+				// 	return this;
+				// },
+				report(progress) {
+					throwIfDone(this.report);
+					if ('placeholder' in progress && 'resolvedContent' in progress) {
+						// Ignore for now, this is the deleted Task type
+						return;
+					}
+
+					const value = typeConvert.ChatResponseProgress.from(that._extension, progress);
+					if (!value) {
+						that._logService.error('Unknown progress type: ' + JSON.stringify(progress));
+						return;
+					}
+
+					_report(value);
+					return this;
+				}
+			};
+		}
+
+		return this._apiObject;
+	}
+}
+
 export class ExtHostChatAgents2 implements ExtHostChatAgentsShape2 {
 
 	private static _idPool = 0;
@@ -61,44 +182,17 @@ export class ExtHostChatAgents2 implements ExtHostChatAgentsShape2 {
 			throw new Error(`[CHAT](${handle}) CANNOT invoke agent because the agent is not registered`);
 		}
 
-		let done = false;
-		function throwIfDone() {
-			if (done) {
-				throw new Error('Only valid while executing the command');
-			}
-		}
-
 		const commandExecution = new DeferredPromise<void>();
 		token.onCancellationRequested(() => commandExecution.complete());
 		this._extHostChatProvider.allowListExtensionWhile(agent.extension.identifier, commandExecution.p);
 
-		const stopWatch = StopWatch.create(false);
-		let firstProgress: number | undefined;
+		const stream = new ChatAgentResponseStream(agent.extension, request, this._proxy, this._logService);
 		try {
 			const convertedHistory = await this.prepareHistory(agent, request, context);
 			const task = agent.invoke(
 				typeConvert.ChatAgentRequest.to(request),
 				{ history: convertedHistory },
-				new Progress<vscode.ChatAgentExtendedProgress>(progress => {
-					throwIfDone();
-
-					// Measure the time to the first progress update with real markdown content
-					if (typeof firstProgress === 'undefined' && 'content' in progress) {
-						firstProgress = stopWatch.elapsed();
-					}
-
-					const convertedProgress = typeConvert.ChatResponseProgress.from(agent.extension, progress);
-					if (!convertedProgress) {
-						this._logService.error('Unknown progress type: ' + JSON.stringify(progress));
-						return;
-					}
-
-					if ('placeholder' in progress && 'resolvedContent' in progress) {
-						// Ignore for now, this is the deleted Task type
-					} else {
-						this._proxy.$handleProgressChunk(request.requestId, convertedProgress);
-					}
-				}),
+				stream.apiObject,
 				token
 			);
 
@@ -112,8 +206,7 @@ export class ExtHostChatAgents2 implements ExtHostChatAgentsShape2 {
 					}
 					sessionResults.set(request.requestId, result);
 
-					const timings = { firstProgress: firstProgress, totalElapsed: stopWatch.elapsed() };
-					return { errorDetails: result.errorDetails, timings };
+					return { errorDetails: result.errorDetails, timings: stream.timings };
 				} else {
 					this._previousResultMap.delete(request.sessionId);
 				}
@@ -126,7 +219,7 @@ export class ExtHostChatAgents2 implements ExtHostChatAgentsShape2 {
 			return { errorDetails: { message: localize('errorResponse', "Error from provider: {0}", toErrorMessage(e)), responseIsIncomplete: true } };
 
 		} finally {
-			done = true;
+			stream.close();
 			commandExecution.complete();
 		}
 	}
@@ -514,7 +607,7 @@ class ExtHostChatAgent<TResult extends vscode.ChatAgentResult2> {
 		} satisfies vscode.ChatAgent2<TResult>;
 	}
 
-	invoke(request: vscode.ChatAgentRequest, context: vscode.ChatAgentContext, progress: Progress<vscode.ChatAgentExtendedProgress>, token: CancellationToken): vscode.ProviderResult<vscode.ChatAgentResult2> {
-		return this._callback(request, context, progress, token);
+	invoke(request: vscode.ChatAgentRequest, context: vscode.ChatAgentContext, response: vscode.ChatAgentExtendedResponseStream, token: CancellationToken): vscode.ProviderResult<vscode.ChatAgentResult2> {
+		return this._callback(request, context, response, token);
 	}
 }

src/vs/workbench/api/common/extHostTypeConverters.ts

@@ -123,6 +123,14 @@ export namespace Range {
 }
 
 export namespace Location {
+
+	export function from(location: vscode.Location): Dto<languages.Location> {
+		return {
+			uri: location.uri,
+			range: Range.from(location.range)
+		};
+	}
+
 	export function to(location: Dto<languages.Location>): vscode.Location {
 		return new types.Location(URI.revive(location.uri), Range.to(location.range));
 	}

src/vscode-dts/vscode.proposed.chatAgents2.d.ts

@@ -274,16 +274,50 @@ declare module 'vscode' {
 		variables: Record<string, ChatVariableValue[]>;
 	}
 
+	export interface ChatAgentResponseStream {
+
+		// RENDERED
+		markdown(value: string | MarkdownString): ChatAgentResponseStream;
+		text(value: string): ChatAgentResponseStream;
+		files(value: ChatAgentFileTreeData): ChatAgentResponseStream;
+		// TODO@jrieken is this sugar for markdown syntax? should we have more like codeblock, bulletlist etc?
+		anchor(value: Uri | Location, attributes?: { title?: string }): ChatAgentResponseStream;
+
+		// META
+		// TODO@API this influences the rendering, it inserts new lines which is likely a bug
+		progress(value: string): ChatAgentResponseStream;
+
+		// TODO@API support non-file uris, like http://example.com
+		reference(value: Uri | Location): ChatAgentResponseStream;
+
+		// TODO@API define support annotations
+		// annotation(value: string | MarkdownString | ChatXYZAnnotation, references?: string): ChatAgentResponseStream;
+
+		/**
+		 * @deprecated use above methods instread
+		 */
+		report(value: ChatAgentProgress): void;
+	}
+
+	/**
+	 * @deprecated use ChatAgentResponseStream instead
+	 */
 	export type ChatAgentContentProgress =
 		| ChatAgentContent
 		| ChatAgentFileTree
 		| ChatAgentInlineContentReference;
 
+	/**
+	 * @deprecated use ChatAgentResponseStream instead
+	 */
 	export type ChatAgentMetadataProgress =
 		| ChatAgentUsedContext
 		| ChatAgentContentReference
 		| ChatAgentProgressMessage;
 
+	/**
+	 * @deprecated use ChatAgentResponseStream instead
+	 */
 	export type ChatAgentProgress = ChatAgentContentProgress | ChatAgentMetadataProgress;
 
 	/**
@@ -380,7 +414,7 @@ declare module 'vscode' {
 		documents: ChatAgentDocumentContext[];
 	}
 
-	export type ChatAgentHandler = (request: ChatAgentRequest, context: ChatAgentContext, progress: Progress<ChatAgentProgress>, token: CancellationToken) => ProviderResult<ChatAgentResult2>;
+	export type ChatAgentHandler = (request: ChatAgentRequest, context: ChatAgentContext, response: ChatAgentResponseStream, token: CancellationToken) => ProviderResult<ChatAgentResult2>;
 
 	export namespace chat {
 

src/vscode-dts/vscode.proposed.chatAgents2Additions.d.ts

@@ -43,6 +43,8 @@ declare module 'vscode' {
 		| ChatAgentMarkdownContent
 		| ChatAgentDetectedAgent;
 
+	export type ChatAgentExtendedResponseStream = ChatAgentResponseStream & Progress<ChatAgentExtendedProgress>;
+
 	export interface ChatAgent2<TResult extends ChatAgentResult2> {
 		/**
 		 * Provide a set of variables that can only be used with this agent.
@@ -64,7 +66,7 @@ declare module 'vscode' {
 		constructor(label: string | CompletionItemLabel, values: ChatVariableValue[]);
 	}
 
-	export type ChatAgentExtendedHandler = (request: ChatAgentRequest, context: ChatAgentContext, progress: Progress<ChatAgentExtendedProgress>, token: CancellationToken) => ProviderResult<ChatAgentResult2>;
+	export type ChatAgentExtendedHandler = (request: ChatAgentRequest, context: ChatAgentContext, response: ChatAgentExtendedResponseStream, token: CancellationToken) => ProviderResult<ChatAgentResult2>;
 
 	export namespace chat {
 		/**