Fix microsoft#205987. (microsoft#206012)

Author: rebornix

src/vs/workbench/contrib/notebook/browser/controller/chat/notebookChatController.ts

@@ -177,6 +177,7 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 	private _sessionCtor: CancelablePromise<void> | undefined;
 	private _activeSession?: Session;
 	private _warmupRequestCts?: CancellationTokenSource;
+	private _activeRequestCts?: CancellationTokenSource;
 	private readonly _ctxHasActiveRequest: IContextKey<boolean>;
 	private readonly _ctxCellWidgetFocused: IContextKey<boolean>;
 	private readonly _ctxUserDidEdit: IContextKey<boolean>;
@@ -397,8 +398,9 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 	}
 
 	async acceptInput() {
-		assertType(this._activeSession);
 		assertType(this._widget);
+		await this._sessionCtor;
+		assertType(this._activeSession);
 		this._warmupRequestCts?.dispose(true);
 		this._warmupRequestCts = undefined;
 		this._activeSession.addInput(new SessionPrompt(this._widget.inlineChatWidget.value));
@@ -449,18 +451,19 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 
 		//TODO: update progress in a newly inserted cell below the widget instead of the fake editor
 
-		const requestCts = new CancellationTokenSource();
+		this._activeRequestCts?.cancel();
+		this._activeRequestCts = new CancellationTokenSource();
 		const progressEdits: TextEdit[][] = [];
 
 		const progressiveEditsQueue = new Queue();
 		const progressiveEditsClock = StopWatch.create();
 		const progressiveEditsAvgDuration = new MovingAverage();
-		const progressiveEditsCts = new CancellationTokenSource(requestCts.token);
+		const progressiveEditsCts = new CancellationTokenSource(this._activeRequestCts.token);
 		let progressiveChatResponse: IInlineChatMessageAppender | undefined;
 		const progress = new AsyncProgress<IInlineChatProgressItem>(async data => {
 			// console.log('received chunk', data, request);
 
-			if (requestCts.token.isCancellationRequested) {
+			if (this._activeRequestCts?.token.isCancellationRequested) {
 				return;
 			}
 
@@ -502,7 +505,7 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 			}
 		});
 
-		const task = this._activeSession.provider.provideResponse(this._activeSession.session, request, progress, requestCts.token);
+		const task = this._activeSession.provider.provideResponse(this._activeSession.session, request, progress, this._activeRequestCts.token);
 		let response: ReplyResponse | ErrorResponse | EmptyResponse;
 
 		try {
@@ -512,7 +515,7 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 			this._widget?.inlineChatWidget.updateInfo(!this._activeSession.lastExchange ? localize('thinking', "Thinking\u2026") : '');
 			this._ctxHasActiveRequest.set(true);
 
-			const reply = await raceCancellationError(Promise.resolve(task), requestCts.token);
+			const reply = await raceCancellationError(Promise.resolve(task), this._activeRequestCts.token);
 			if (progressiveEditsQueue.size > 0) {
 				// we must wait for all edits that came in via progress to complete
 				await Event.toPromise(progressiveEditsQueue.onDrained);
@@ -575,7 +578,7 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 		this._widget?.inlineChatWidget.updateInfo('');
 		this._widget?.inlineChatWidget.updateToolbar(true);
 
-		this._activeSession.addExchange(new SessionExchange(this._activeSession.lastInput, response));
+		this._activeSession?.addExchange(new SessionExchange(this._activeSession.lastInput, response));
 		this._ctxLastResponseType.set(response instanceof ReplyResponse ? response.raw.type : undefined);
 	}
 
@@ -762,15 +765,12 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 			this._strategy?.cancel();
 		}
 
-		if (this._activeSession) {
-			this._inlineChatSessionService.releaseSession(this._activeSession);
-		}
-
-		this._activeSession = undefined;
+		this._activeRequestCts?.cancel();
 	}
 
 	discard() {
 		this._strategy?.cancel();
+		this._activeRequestCts?.cancel();
 		this._widget?.discardChange();
 		this.dismiss();
 	}