Fix coding agent cancellation via stop button. (microsoft#258406)

Author: rebornix

src/vs/workbench/api/browser/mainThreadChatSessions.ts

@@ -9,6 +9,8 @@ import { Emitter, Event } from '../../../base/common/event.js';
 import { Disposable, DisposableMap } from '../../../base/common/lifecycle.js';
 import { revive } from '../../../base/common/marshalling.js';
 import { URI, UriComponents } from '../../../base/common/uri.js';
+import { localize } from '../../../nls.js';
+import { IDialogService } from '../../../platform/dialogs/common/dialogs.js';
 import { ILogService } from '../../../platform/log/common/log.js';
 import { ChatViewId } from '../../contrib/chat/browser/chat.js';
 import { ChatViewPane } from '../../contrib/chat/browser/chatViewPane.js';
@@ -42,6 +44,7 @@ export class MainThreadChatSessions extends Disposable implements MainThreadChat
 	constructor(
 		private readonly _extHostContext: IExtHostContext,
 		@IChatSessionsService private readonly _chatSessionsService: IChatSessionsService,
+		@IDialogService private readonly _dialogService: IDialogService,
 		@IEditorService private readonly _editorService: IEditorService,
 		@ILogService private readonly _logService: ILogService,
 		@IViewsService private readonly _viewsService: IViewsService,
@@ -88,12 +91,29 @@ export class MainThreadChatSessions extends Disposable implements MainThreadChat
 			const progressEmitter = new Emitter<IChatProgress[]>;
 			const completionEmitter = new Emitter<void>();
 			let progressEvent: Event<IChatProgress[]> | undefined = undefined;
+			let interruptActiveResponseCallback: (() => Promise<boolean>) | undefined = undefined;
 			if (sessionContent.hasActiveResponseCallback) {
 				const requestId = 'ongoing';
 				// set progress
 				progressEvent = progressEmitter.event;
 				// store the event emitter using a key that combines handle and session id
 				const progressKey = `${providerHandle}_${id}_${requestId}`;
+				interruptActiveResponseCallback = async () => {
+					return this._dialogService.confirm({
+						message: localize('interruptActiveResponse', 'Are you sure you want to interrupt the active session?')
+					}).then(confirmed => {
+						if (confirmed.confirmed) {
+							this._proxy.$interruptChatSessionActiveResponse(providerHandle, id, requestId);
+							return true;
+						} else {
+							progressEmitter.fire([{
+								kind: 'progressMessage',
+								content: { value: '' }
+							}]);
+							return false;
+						}
+					});
+				};
 				this._activeProgressEmitters.set(progressKey, progressEmitter);
 				this._completionEmitters.set(progressKey, completionEmitter);
 			}
@@ -127,6 +147,7 @@ export class MainThreadChatSessions extends Disposable implements MainThreadChat
 				}),
 				progressEvent: progressEvent,
 				requestHandler: requestHandler,
+				interruptActiveResponseCallback: interruptActiveResponseCallback,
 				dispose: () => {
 					progressEmitter.dispose();
 					completionEmitter.dispose();

src/vs/workbench/api/common/extHost.protocol.ts

@@ -3147,6 +3147,7 @@ export interface ExtHostChatSessionsShape {
 	$provideChatSessionItems(providerHandle: number, token: CancellationToken): Promise<Dto<IChatSessionItem>[]>;
 
 	$provideChatSessionContent(providerHandle: number, sessionId: string, token: CancellationToken): Promise<ChatSessionDto>;
+	$interruptChatSessionActiveResponse(providerHandle: number, sessionId: string, requestId: string): Promise<void>;
 	$disposeChatSessionContent(providerHandle: number, sessionId: string): Promise<void>;
 	$invokeChatSessionRequestHandler(providerHandle: number, id: string, request: IChatAgentRequest, history: any[], token: CancellationToken): Promise<IChatAgentResult>;
 }

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

@@ -208,6 +208,12 @@ export class ExtHostChatSessions extends Disposable implements ExtHostChatSessio
 		};
 	}
 
+	async $interruptChatSessionActiveResponse(providerHandle: number, sessionId: string, requestId: string): Promise<void> {
+		const key = `${providerHandle}_${sessionId}`;
+		const entry = this._extHostChatSessions.get(key);
+		entry?.disposeCts.cancel();
+	}
+
 	async $disposeChatSessionContent(providerHandle: number, sessionId: string): Promise<void> {
 		const key = `${providerHandle}_${sessionId}`;
 		const entry = this._extHostChatSessions.get(key);

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

@@ -11,7 +11,7 @@ import { ErrorNoTelemetry } from '../../../../base/common/errors.js';
 import { Emitter, Event } from '../../../../base/common/event.js';
 import { MarkdownString } from '../../../../base/common/htmlContent.js';
 import { Iterable } from '../../../../base/common/iterator.js';
-import { Disposable, DisposableMap, DisposableStore, IDisposable } from '../../../../base/common/lifecycle.js';
+import { Disposable, DisposableMap, DisposableStore, IDisposable, MutableDisposable } from '../../../../base/common/lifecycle.js';
 import { revive } from '../../../../base/common/marshalling.js';
 import { autorun, derived, IObservable, ObservableMap } from '../../../../base/common/observable.js';
 import { StopWatch } from '../../../../base/common/stopwatch.js';
@@ -636,15 +636,34 @@ export class ChatService extends Disposable implements IChatService {
 			}
 		}
 
-		if (content.progressEvent) {
-			disposables.add(content.progressEvent(e => {
-				if (lastRequest) {
-					for (const progress of e) {
-						if (progress.kind === 'progressMessage' && progress.content.value === 'Session completed') {
-							model?.completeResponse(lastRequest);
-						} else {
-							model?.acceptResponseProgress(lastRequest, progress);
+		if (content.progressEvent && lastRequest) {
+			const initialCancellationRequest = this.instantiationService.createInstance(CancellableRequest, new CancellationTokenSource(), undefined);
+			this._pendingRequests.set(model.sessionId, initialCancellationRequest);
+			const cancellationListener = new MutableDisposable();
+
+			const createCancellationListener = (token: CancellationToken) => {
+				return token.onCancellationRequested(() => {
+					content.interruptActiveResponseCallback?.().then(userConfirmedInterruption => {
+						if (!userConfirmedInterruption) {
+							// User cancelled the interruption
+							const newCancellationRequest = this.instantiationService.createInstance(CancellableRequest, new CancellationTokenSource(), undefined);
+							this._pendingRequests.set(model.sessionId, newCancellationRequest);
+							cancellationListener.value = createCancellationListener(newCancellationRequest.cancellationTokenSource.token);
 						}
+					});
+				});
+			};
+
+			cancellationListener.value = createCancellationListener(initialCancellationRequest.cancellationTokenSource.token);
+			disposables.add(cancellationListener);
+
+			disposables.add(content.progressEvent(e => {
+				for (const progress of e) {
+					if (progress.kind === 'progressMessage' && progress.content.value === 'Session completed') {
+						model?.completeResponse(lastRequest);
+						cancellationListener.clear();
+					} else {
+						model?.acceptResponseProgress(lastRequest, progress);
 					}
 				}
 			}));

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

@@ -40,6 +40,7 @@ export interface ChatSession extends IDisposable {
 		| { type: 'response'; parts: IChatProgress[] }>;
 
 	readonly progressEvent?: Event<IChatProgress[]>;
+	readonly interruptActiveResponseCallback?: () => Promise<boolean>;
 
 	requestHandler?: (
 		request: IChatAgentRequest,