Fix issue with "reun without" while request is pending (microsoft#213926)

- send out `ChatRequestRemovalReason`
- make inline chat re-enter `SHOW_REQUEST` when resending an active request
- the `SHOW_REQUEST` state should be prepared to come after the request has started
- more cancellation checking with `asProgressiveEdit`
- add tests

microsoft/vscode-copilot#5997

Author: jrieken

src/vs/workbench/contrib/chat/common/chatModel.ts

@@ -529,10 +529,28 @@ export interface IChatAddResponseEvent {
 	response: IChatResponseModel;
 }
 
+export const enum ChatRequestRemovalReason {
+	/**
+	 * "Normal" remove
+	 */
+	Removal,
+
+	/**
+	 * Removed because the request will be resent
+	 */
+	Resend,
+
+	/**
+	 * Remove because the request is moving to another model
+	 */
+	Adoption
+}
+
 export interface IChatRemoveRequestEvent {
 	kind: 'removeRequest';
 	requestId: string;
 	responseId?: string;
+	reason: ChatRequestRemovalReason;
 }
 
 export interface IChatInitEvent {
@@ -804,7 +822,7 @@ export class ChatModel extends Disposable implements IChatModel {
 		request.response?.adoptTo(this);
 		this._requests.push(request);
 
-		oldOwner._onDidChange.fire({ kind: 'removeRequest', requestId: request.id, responseId: request.response?.id });
+		oldOwner._onDidChange.fire({ kind: 'removeRequest', requestId: request.id, responseId: request.response?.id, reason: ChatRequestRemovalReason.Adoption });
 		this._onDidChange.fire({ kind: 'addRequest', request });
 	}
 
@@ -841,12 +859,12 @@ export class ChatModel extends Disposable implements IChatModel {
 		}
 	}
 
-	removeRequest(id: string): void {
+	removeRequest(id: string, reason: ChatRequestRemovalReason = ChatRequestRemovalReason.Removal): void {
 		const index = this._requests.findIndex(request => request.id === id);
 		const request = this._requests[index];
 
 		if (index !== -1) {
-			this._onDidChange.fire({ kind: 'removeRequest', requestId: request.id, responseId: request.response?.id });
+			this._onDidChange.fire({ kind: 'removeRequest', requestId: request.id, responseId: request.response?.id, reason });
 			this._requests.splice(index, 1);
 			request.response?.dispose();
 		}

src/vs/workbench/contrib/chat/common/chatServiceImpl.ts

@@ -24,7 +24,7 @@ import { IStorageService, StorageScope, StorageTarget } from 'vs/platform/storag
 import { ITelemetryService } from 'vs/platform/telemetry/common/telemetry';
 import { IWorkspaceContextService } from 'vs/platform/workspace/common/workspace';
 import { ChatAgentLocation, IChatAgent, IChatAgentRequest, IChatAgentResult, IChatAgentService } from 'vs/workbench/contrib/chat/common/chatAgents';
-import { ChatModel, ChatRequestModel, ChatWelcomeMessageModel, IChatModel, IChatRequestModel, IChatRequestVariableData, IChatRequestVariableEntry, IChatResponseModel, IExportableChatData, ISerializableChatData, ISerializableChatsData, getHistoryEntriesFromModel, updateRanges } from 'vs/workbench/contrib/chat/common/chatModel';
+import { ChatModel, ChatRequestModel, ChatRequestRemovalReason, ChatWelcomeMessageModel, IChatModel, IChatRequestModel, IChatRequestVariableData, IChatRequestVariableEntry, IChatResponseModel, IExportableChatData, ISerializableChatData, ISerializableChatsData, getHistoryEntriesFromModel, updateRanges } from 'vs/workbench/contrib/chat/common/chatModel';
 import { ChatRequestAgentPart, ChatRequestAgentSubcommandPart, ChatRequestSlashCommandPart, IParsedChatRequest, chatAgentLeader, chatSubcommandLeader, getPromptText } from 'vs/workbench/contrib/chat/common/chatParserTypes';
 import { ChatRequestParser } from 'vs/workbench/contrib/chat/common/chatRequestParser';
 import { ChatCopyKind, IChatCompleteResponse, IChatDetail, IChatFollowup, IChatProgress, IChatSendRequestData, IChatSendRequestOptions, IChatSendRequestResponseState, IChatService, IChatTransferredSessionData, IChatUserActionEvent, ChatAgentVoteDirection } from 'vs/workbench/contrib/chat/common/chatService';
@@ -418,17 +418,18 @@ export class ChatService extends Disposable implements IChatService {
 
 		await model.waitForInitialization();
 
-		if (this._pendingRequests.has(request.session.sessionId)) {
-			this.trace('sendRequest', `Session ${request.session.sessionId} already has a pending request`);
-			return;
+		const cts = this._pendingRequests.get(request.session.sessionId);
+		if (cts) {
+			this.trace('resendRequest', `Session ${request.session.sessionId} already has a pending request, cancelling...`);
+			cts.cancel();
 		}
 
 		const location = options?.location ?? model.initialLocation;
 		const attempt = options?.attempt ?? 0;
 		const enableCommandDetection = !options?.noCommandDetection;
 		const defaultAgent = this.chatAgentService.getDefaultAgent(location)!;
 
-		this.removeRequest(model.sessionId, request.id);
+		model.removeRequest(request.id, ChatRequestRemovalReason.Resend);
 
 		await this._sendRequestAsync(model, model.sessionId, request.message, attempt, enableCommandDetection, defaultAgent, location, options).responseCompletePromise;
 	}

src/vs/workbench/contrib/inlineChat/browser/inlineChatController.ts

@@ -40,7 +40,7 @@ import { StashedSession } from './inlineChatSession';
 import { IModelDeltaDecoration, ITextModel, IValidEditOperation } from 'vs/editor/common/model';
 import { InlineChatContentWidget } from 'vs/workbench/contrib/inlineChat/browser/inlineChatContentWidget';
 import { MessageController } from 'vs/editor/contrib/message/browser/messageController';
-import { ChatModel, IChatRequestModel, IChatTextEditGroup, IChatTextEditGroupState, IResponse } from 'vs/workbench/contrib/chat/common/chatModel';
+import { ChatModel, ChatRequestRemovalReason, IChatRequestModel, IChatTextEditGroup, IChatTextEditGroupState, IResponse } from 'vs/workbench/contrib/chat/common/chatModel';
 import { InlineChatError } from 'vs/workbench/contrib/inlineChat/browser/inlineChatSessionServiceImpl';
 import { ILanguageFeaturesService } from 'vs/editor/common/services/languageFeatures';
 import { ChatInputPart } from 'vs/workbench/contrib/chat/browser/chatInputPart';
@@ -445,10 +445,16 @@ export class InlineChatController implements IEditorContribution {
 				const exchange = this._session!.exchanges.find(candidate => candidate.prompt.request.id === e.requestId);
 				if (exchange && this._editor.hasModel()) {
 					// undo till this point
-					const model = this._editor.getModel();
-					const targetAltVersion = exchange.prompt.modelAltVersionId;
-					while (targetAltVersion < model.getAlternativeVersionId() && model.canUndo()) {
-						await model.undo();
+					this._session!.hunkData.ignoreTextModelNChanges = true;
+					try {
+
+						const model = this._editor.getModel();
+						const targetAltVersion = exchange.prompt.modelAltVersionId;
+						while (targetAltVersion < model.getAlternativeVersionId() && model.canUndo()) {
+							await model.undo();
+						}
+					} finally {
+						this._session!.hunkData.ignoreTextModelNChanges = false;
 					}
 				}
 			}
@@ -641,7 +647,7 @@ export class InlineChatController implements IEditorContribution {
 		const progressiveEditsClock = StopWatch.create();
 		const progressiveEditsQueue = new Queue();
 
-		let next: State.SHOW_RESPONSE | State.CANCEL | State.PAUSE | State.ACCEPT | State.WAIT_FOR_INPUT = State.SHOW_RESPONSE;
+		let next: State.SHOW_RESPONSE | State.SHOW_REQUEST | State.CANCEL | State.PAUSE | State.ACCEPT | State.WAIT_FOR_INPUT = State.SHOW_RESPONSE;
 		store.add(Event.once(this._messages.event)(message => {
 			this._log('state=_makeRequest) message received', message);
 			this._chatService.cancelCurrentRequestForSession(chatModel.sessionId);
@@ -658,7 +664,11 @@ export class InlineChatController implements IEditorContribution {
 			if (e.kind === 'removeRequest' && e.requestId === request.id) {
 				progressiveEditsCts.cancel();
 				responsePromise.complete();
-				next = State.CANCEL;
+				if (e.reason === ChatRequestRemovalReason.Resend) {
+					next = State.SHOW_REQUEST;
+				} else {
+					next = State.CANCEL;
+				}
 			}
 		}));
 
@@ -678,61 +688,59 @@ export class InlineChatController implements IEditorContribution {
 		let localEditGroup: IChatTextEditGroup | undefined;
 
 		// apply edits
-		store.add(response.onDidChange(() => {
-
-			if (response.isCanceled) {
-				progressiveEditsCts.cancel();
-				responsePromise.complete();
-				return;
-			}
-
-			if (response.isComplete) {
-				responsePromise.complete();
-				return;
-			}
+		const handleResponse = () => {
 
 			if (!localEditGroup) {
-				localEditGroup = <IChatTextEditGroup>response.response.value.find(part => part.kind === 'textEditGroup' && isEqual(part.uri, this._session?.textModelN.uri));
+				localEditGroup = <IChatTextEditGroup | undefined>response.response.value.find(part => part.kind === 'textEditGroup' && isEqual(part.uri, this._session?.textModelN.uri));
 			}
 
-			if (!localEditGroup) {
-				return;
-			}
+			if (localEditGroup) {
 
-			localEditGroup.state ??= editState;
+				localEditGroup.state ??= editState;
 
-			const edits = localEditGroup.edits;
-			const newEdits = edits.slice(lastLength);
-			if (newEdits.length === 0) {
-				return; // NO change
-			}
-			lastLength = edits.length;
-			progressiveEditsAvgDuration.update(progressiveEditsClock.elapsed());
-			progressiveEditsClock.reset();
+				const edits = localEditGroup.edits;
+				const newEdits = edits.slice(lastLength);
+				if (newEdits.length > 0) {
+					// NEW changes
+					lastLength = edits.length;
+					progressiveEditsAvgDuration.update(progressiveEditsClock.elapsed());
+					progressiveEditsClock.reset();
 
-			progressiveEditsQueue.queue(async () => {
+					progressiveEditsQueue.queue(async () => {
 
-				const startThen = this._session!.wholeRange.value.getStartPosition();
+						const startThen = this._session!.wholeRange.value.getStartPosition();
 
-				// making changes goes into a queue because otherwise the async-progress time will
-				// influence the time it takes to receive the changes and progressive typing will
-				// become infinitely fast
-				for (const edits of newEdits) {
-					await this._makeChanges(edits, {
-						duration: progressiveEditsAvgDuration.value,
-						token: progressiveEditsCts.token
-					}, isFirstChange);
+						// making changes goes into a queue because otherwise the async-progress time will
+						// influence the time it takes to receive the changes and progressive typing will
+						// become infinitely fast
+						for (const edits of newEdits) {
+							await this._makeChanges(edits, {
+								duration: progressiveEditsAvgDuration.value,
+								token: progressiveEditsCts.token
+							}, isFirstChange);
 
-					isFirstChange = false;
-				}
+							isFirstChange = false;
+						}
 
-				// reshow the widget if the start position changed or shows at the wrong position
-				const startNow = this._session!.wholeRange.value.getStartPosition();
-				if (!startNow.equals(startThen) || !this._ui.value.zone.position?.equals(startNow)) {
-					this._showWidget(false, startNow.delta(-1));
+						// reshow the widget if the start position changed or shows at the wrong position
+						const startNow = this._session!.wholeRange.value.getStartPosition();
+						if (!startNow.equals(startThen) || !this._ui.value.zone.position?.equals(startNow)) {
+							this._showWidget(false, startNow.delta(-1));
+						}
+					});
 				}
-			});
-		}));
+			}
+
+			if (response.isCanceled) {
+				progressiveEditsCts.cancel();
+				responsePromise.complete();
+
+			} else if (response.isComplete) {
+				responsePromise.complete();
+			}
+		};
+		store.add(response.onDidChange(handleResponse));
+		handleResponse();
 
 		// (1) we must wait for the request to finish
 		// (2) we must wait for all edits that came in via progress to complete

src/vs/workbench/contrib/inlineChat/browser/utils.ts

@@ -64,6 +64,9 @@ export function asProgressiveEdit(interval: IntervalTimer, edit: IIdentifiedSing
 	let newText = edit.text ?? '';
 
 	interval.cancelAndSet(() => {
+		if (token.isCancellationRequested) {
+			return;
+		}
 		const r = getNWords(newText, 1);
 		stream.emitOne(r.value);
 		newText = newText.substring(r.value.length);

src/vs/workbench/contrib/inlineChat/test/browser/inlineChatController.test.ts

@@ -5,7 +5,7 @@
 
 import * as assert from 'assert';
 import { equals } from 'vs/base/common/arrays';
-import { timeout } from 'vs/base/common/async';
+import { raceCancellation, timeout } from 'vs/base/common/async';
 import { Emitter, Event } from 'vs/base/common/event';
 import { DisposableStore } from 'vs/base/common/lifecycle';
 import { mock } from 'vs/base/test/common/mock';
@@ -62,6 +62,7 @@ import { TestCommandService } from 'vs/editor/test/browser/editorTestServices';
 import { INotebookEditorService } from 'vs/workbench/contrib/notebook/browser/services/notebookEditorService';
 import { RerunAction } from 'vs/workbench/contrib/inlineChat/browser/inlineChatActions';
 import { CancellationToken } from 'vs/base/common/cancellation';
+import { assertType } from 'vs/base/common/types';
 
 suite('InteractiveChatController', function () {
 
@@ -642,4 +643,84 @@ suite('InteractiveChatController', function () {
 
 		await r;
 	});
+
+	test('Clicking "re-run without /doc" while a request is in progress closes the widget #5997', async function () {
+
+		model.setValue('');
+
+		let count = 0;
+		const commandDetection: (boolean | undefined)[] = [];
+
+		store.add(chatAgentService.registerDynamicAgent({
+			id: 'testEditorAgent2',
+			...agentData
+		}, {
+			async invoke(request, progress, history, token) {
+				commandDetection.push(request.enableCommandDetection);
+				progress({ kind: 'textEdit', uri: model.uri, edits: [{ range: new Range(1, 1, 1, 1), text: request.message + (count++) }] });
+
+				if (count === 1) {
+					// FIRST call waits for cancellation
+					await raceCancellation(new Promise<never>(() => { }), token);
+				} else {
+					await timeout(10);
+				}
+
+				return {};
+			},
+		}));
+		ctrl = instaService.createInstance(TestController, editor);
+
+		// REQUEST 1
+		const p = ctrl.waitFor([...TestController.INIT_SEQUENCE, State.SHOW_REQUEST]);
+		ctrl.run({ message: 'Hello-', autoSend: true });
+		await p;
+
+		// resend pending request without command detection
+		const request = ctrl.chatWidget.viewModel?.model.getRequests().at(-1);
+		assertType(request);
+		const p2 = ctrl.waitFor([State.SHOW_REQUEST, State.SHOW_RESPONSE]);
+		chatService.resendRequest(request, { noCommandDetection: true, attempt: request.attempt + 1, location: ChatAgentLocation.Editor });
+
+		await p2;
+
+		assert.deepStrictEqual(commandDetection, [true, false]);
+		assert.strictEqual(model.getValue(), 'Hello-1');
+	});
+
+	test('Re-run without after request is done', async function () {
+
+		model.setValue('');
+
+		let count = 0;
+		const commandDetection: (boolean | undefined)[] = [];
+
+		store.add(chatAgentService.registerDynamicAgent({
+			id: 'testEditorAgent2',
+			...agentData
+		}, {
+			async invoke(request, progress, history, token) {
+				commandDetection.push(request.enableCommandDetection);
+				progress({ kind: 'textEdit', uri: model.uri, edits: [{ range: new Range(1, 1, 1, 1), text: request.message + (count++) }] });
+				return {};
+			},
+		}));
+		ctrl = instaService.createInstance(TestController, editor);
+
+		// REQUEST 1
+		const p = ctrl.waitFor([...TestController.INIT_SEQUENCE, State.SHOW_REQUEST, State.SHOW_RESPONSE, State.WAIT_FOR_INPUT]);
+		ctrl.run({ message: 'Hello-', autoSend: true });
+		await p;
+
+		// resend pending request without command detection
+		const request = ctrl.chatWidget.viewModel?.model.getRequests().at(-1);
+		assertType(request);
+		const p2 = ctrl.waitFor([State.SHOW_REQUEST, State.SHOW_RESPONSE, State.WAIT_FOR_INPUT]);
+		chatService.resendRequest(request, { noCommandDetection: true, attempt: request.attempt + 1, location: ChatAgentLocation.Editor });
+
+		await p2;
+
+		assert.deepStrictEqual(commandDetection, [true, false]);
+		assert.strictEqual(model.getValue(), 'Hello-1');
+	});
 });