make responses of chatRequestAccess proposal easier to consume and digest (microsoft#195479)

* make responses of `chatRequestAccess` proposal easier to consume and digest

* allow same modern JS in API as in the rest

Author: jrieken

src/tsconfig.vscode-dts.json

@@ -13,8 +13,7 @@
 		"forceConsistentCasingInFileNames": true,
 		"types": [],
 		"lib": [
-			"es5",
-			"ES2015.Iterable"
+			"ES2022"
 		],
 	},
 	"include": [

src/vs/base/common/async.ts

@@ -1887,4 +1887,61 @@ export function createCancelableAsyncIterable<T>(callback: (token: CancellationT
 	});
 }
 
+export class DeferredAsyncIterableObject<T> {
+
+	private readonly _deferred = new DeferredPromise<void>();
+	private readonly _asyncIterable: AsyncIterableObject<T>;
+
+	private _errorFn: (error: Error) => void;
+	private _emitFn: (item: T) => void;
+
+	constructor() {
+		this._asyncIterable = new AsyncIterableObject(emitter => {
+
+			if (earlyError) {
+				emitter.reject(earlyError);
+				return;
+			}
+			if (earlyItems) {
+				emitter.emitMany(earlyItems);
+			}
+			this._errorFn = (error: Error) => emitter.reject(error);
+			this._emitFn = (item: T) => emitter.emitOne(item);
+			return this._deferred.p;
+		});
+
+		let earlyError: Error | undefined;
+		let earlyItems: T[] | undefined;
+
+		this._emitFn = (item: T) => {
+			if (!earlyItems) {
+				earlyItems = [];
+			}
+			earlyItems.push(item);
+		};
+		this._errorFn = (error: Error) => {
+			if (!earlyError) {
+				earlyError = error;
+			}
+		};
+	}
+
+	get asyncIterable(): AsyncIterableObject<T> {
+		return this._asyncIterable;
+	}
+
+	complete(): void {
+		this._deferred.complete();
+	}
+
+	error(error: Error): void {
+		this._errorFn(error);
+		this._deferred.complete();
+	}
+
+	emit(item: T): void {
+		this._emitFn(item);
+	}
+}
+
 //#endregion

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

@@ -3,7 +3,7 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { CancellationToken } from 'vs/base/common/cancellation';
+import { CancellationToken, CancellationTokenSource } from 'vs/base/common/cancellation';
 import { IDisposable, toDisposable } from 'vs/base/common/lifecycle';
 import { ILogService } from 'vs/platform/log/common/log';
 import { ExtHostChatProviderShape, IMainContext, MainContext, MainThreadChatProviderShape } from 'vs/workbench/api/common/extHost.protocol';
@@ -12,12 +12,82 @@ import type * as vscode from 'vscode';
 import { Progress } from 'vs/platform/progress/common/progress';
 import { IChatMessage, IChatResponseFragment } from 'vs/workbench/contrib/chat/common/chatProvider';
 import { ExtensionIdentifier, ExtensionIdentifierMap } from 'vs/platform/extensions/common/extensions';
+import { DeferredAsyncIterableObject } from 'vs/base/common/async';
+import { Emitter } from 'vs/base/common/event';
 
 type ProviderData = {
 	readonly extension: ExtensionIdentifier;
 	readonly provider: vscode.ChatResponseProvider;
 };
 
+class ChatResponseStream {
+
+	readonly apiObj: vscode.ChatResponseStream;
+	readonly stream = new DeferredAsyncIterableObject<string>();
+
+	constructor(option: number, stream?: DeferredAsyncIterableObject<string>) {
+		this.stream = stream ?? new DeferredAsyncIterableObject<string>();
+		const that = this;
+		this.apiObj = {
+			option: option,
+			response: that.stream.asyncIterable
+		};
+	}
+}
+
+class ChatRequest {
+
+	readonly apiObject: vscode.ChatRequest;
+
+	private readonly _onDidStart = new Emitter<vscode.ChatResponseStream>();
+	private readonly _responseStreams = new Map<number, ChatResponseStream>();
+	private readonly _defaultStream = new DeferredAsyncIterableObject<string>();
+	private _isDone: boolean = false;
+
+	constructor(
+		promise: Promise<any>,
+		cts: CancellationTokenSource
+	) {
+		const that = this;
+		this.apiObject = {
+			result: promise,
+			response: that._defaultStream.asyncIterable,
+			onDidStartResponseStream: that._onDidStart.event,
+			cancel() { cts.cancel(); },
+		};
+
+		promise.finally(() => {
+			this._isDone = true;
+			if (this._responseStreams.size > 0) {
+				for (const [, value] of this._responseStreams) {
+					value.stream.complete();
+				}
+			} else {
+				this._defaultStream.complete();
+			}
+		});
+	}
+
+	handleFragment(fragment: IChatResponseFragment): void {
+		if (this._isDone) {
+			return;
+		}
+		let res = this._responseStreams.get(fragment.index);
+		if (!res) {
+			if (this._responseStreams.size === 0) {
+				// the first response claims the default response
+				res = new ChatResponseStream(fragment.index, this._defaultStream);
+			} else {
+				res = new ChatResponseStream(fragment.index);
+			}
+			this._responseStreams.set(fragment.index, res);
+			this._onDidStart.fire(res.apiObj);
+		}
+		res.stream.emit(fragment.part);
+	}
+
+}
+
 export class ExtHostChatProvider implements ExtHostChatProviderShape {
 
 	private static _idPool = 1;
@@ -62,7 +132,7 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape {
 
 	//#region --- making request
 
-	private readonly _pendingRequest = new Map<number, vscode.Progress<vscode.ChatResponseFragment>>();
+	private readonly _pendingRequest = new Map<number, { res: ChatRequest }>();
 
 	private readonly _chatAccessAllowList = new ExtensionIdentifierMap<Promise<unknown>>();
 
@@ -84,24 +154,31 @@ export class ExtHostChatProvider implements ExtHostChatProviderShape {
 			get isRevoked() {
 				return !that._chatAccessAllowList.has(from);
 			},
-			async makeRequest(messages, options, progress, token) {
+			makeRequest(messages, options, token) {
 
 				if (!that._chatAccessAllowList.has(from)) {
 					throw new Error('Access to chat has been revoked');
 				}
 
+				const cts = new CancellationTokenSource(token);
 				const requestId = (Math.random() * 1e6) | 0;
-				that._pendingRequest.set(requestId, progress);
-				try {
-					await that._proxy.$fetchResponse(from, identifier, requestId, messages.map(typeConvert.ChatMessage.from), options, token);
-				} finally {
+				const requestPromise = that._proxy.$fetchResponse(from, identifier, requestId, messages.map(typeConvert.ChatMessage.from), options ?? {}, cts.token);
+				const res = new ChatRequest(requestPromise, cts);
+				that._pendingRequest.set(requestId, { res });
+
+				requestPromise.finally(() => {
 					that._pendingRequest.delete(requestId);
-				}
+				});
+
+				return res.apiObject;
 			},
 		};
 	}
 
 	async $handleResponseFragment(requestId: number, chunk: IChatResponseFragment): Promise<void> {
-		this._pendingRequest.get(requestId)?.report(chunk);
+		const data = this._pendingRequest.get(requestId);//.report(chunk);
+		if (data) {
+			data.res.handleFragment(chunk);
+		}
 	}
 }

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

@@ -5,32 +5,84 @@
 
 declare module 'vscode' {
 
-	export interface ChatResponseFragment {
-		index: number;
-		part: string;
+	export interface ChatResponseStream {
+
+		/**
+		 * The response stream.
+		 */
+		readonly response: AsyncIterable<string>;
+
+		/**
+		 * The variant of multiple responses. This is used to disambiguate between multiple
+		 * response streams when having asked for multiple response options
+		 */
+		readonly option: number;
+	}
+
+	export interface ChatRequest {
+
+		/**
+		 * The overall result of the request which represents failure or success
+		 * but _not_ the actual response or responses
+		 */
+		result: Thenable<any>;
+
+		/**
+		 * The _default response_ stream. This is the stream of the first response option
+		 * receiving data.
+		 *
+		 * Usually there is only one response option and this stream is more convienient to use
+		 * than the {@link onDidStartResponseStream `onDidStartResponseStream`} event.
+		 */
+		response: AsyncIterable<string>;
+
+		/**
+		 * An event that fires whenever a new response option is available. The response
+		 * itself is a stream of the actual response.
+		 *
+		 * *Note* that the first time this event fires, the {@link ChatResponseStream.response response stream}
+		 * is the same as the {@link response `default response stream`}.
+		 *
+		 * *Note* that unless requested there is only one response option, so this event will only fire
+		 * once.
+		 */
+		onDidStartResponseStream: Event<ChatResponseStream>;
+
+		/**
+		 * Cancel this request.
+		 */
+		// TODO@API remove this? We pass a token to makeRequest call already
+		cancel(): void;
 	}
 
 	/**
 	 * Represents access to using a chat provider (LLM). Access is granted and temporary, usually only valid
-	 * for the duration of an user interaction.
+	 * for the duration of an user interaction or specific time frame.
 	 */
 	export interface ChatAccess {
 
 		/**
-		 * Whether the access to chat has been revoked. This happens when the user interaction that allowed for
-		 * chat access is finished.
+		 * Whether the access to chat has been revoked. This happens when the condition that allowed for
+		 * chat access doesn't hold anymore, e.g a user interaction has ended.
 		 */
 		isRevoked: boolean;
 
 		/**
-		 * TODO: return an AsyncIterable instead of asking to pass Progress<...>?
+		 * Make a chat request.
+		 *
+		 * The actual response will be reported back via the `progress` callback. The promise returned by this function
+		 * returns a overall result which represents failure or success of the request.
+		 *
+		 * Chat can be asked for multiple response options. In that case the `progress` callback will be called multiple
+		 * time with different `ChatResponseStream` objects. Each object represents a different response option and the actual
+		 * response will be reported back via their `stream` property.
+		 *
+		 * *Note:* This will throw an error if access has been revoked.
 		 *
 		 * @param messages
 		 * @param options
-		 * @param progress
-		 * @param token
 		 */
-		makeRequest(messages: ChatMessage[], options: { [name: string]: any }, progress: Progress<ChatResponseFragment>, token: CancellationToken): Thenable<any>;
+		makeRequest(messages: ChatMessage[], options: { [name: string]: any }, token: CancellationToken): ChatRequest;
 	}
 
 	export namespace chat {