Merge pull request microsoft#205164 from microsoft/rebornix/healthy-buzzard

Ensure escape discard/accept changes properly.

Author: rebornix

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

@@ -14,12 +14,12 @@ import { InputFocusedContextKey } from 'vs/platform/contextkey/common/contextkey
 import { ServicesAccessor } from 'vs/platform/instantiation/common/instantiation';
 import { KeybindingWeight } from 'vs/platform/keybinding/common/keybindingsRegistry';
 import { CTX_INLINE_CHAT_FOCUSED, CTX_INLINE_CHAT_HAS_PROVIDER, CTX_INLINE_CHAT_INNER_CURSOR_FIRST, CTX_INLINE_CHAT_INNER_CURSOR_LAST, CTX_INLINE_CHAT_LAST_RESPONSE_TYPE, CTX_INLINE_CHAT_RESPONSE_TYPES, InlineChatResponseFeedbackKind, InlineChatResponseTypes } from 'vs/workbench/contrib/inlineChat/common/inlineChat';
-import { CTX_NOTEBOOK_CELL_CHAT_FOCUSED, CTX_NOTEBOOK_CHAT_HAS_ACTIVE_REQUEST, MENU_CELL_CHAT_INPUT, MENU_CELL_CHAT_WIDGET, MENU_CELL_CHAT_WIDGET_FEEDBACK, MENU_CELL_CHAT_WIDGET_STATUS } from 'vs/workbench/contrib/notebook/browser/controller/chat/notebookChatContext';
+import { CTX_NOTEBOOK_CELL_CHAT_FOCUSED, CTX_NOTEBOOK_CHAT_HAS_ACTIVE_REQUEST, CTX_NOTEBOOK_CHAT_USER_DID_EDIT, MENU_CELL_CHAT_INPUT, MENU_CELL_CHAT_WIDGET, MENU_CELL_CHAT_WIDGET_FEEDBACK, MENU_CELL_CHAT_WIDGET_STATUS } from 'vs/workbench/contrib/notebook/browser/controller/chat/notebookChatContext';
 import { NotebookChatController } from 'vs/workbench/contrib/notebook/browser/controller/chat/notebookChatController';
 import { INotebookActionContext, INotebookCellActionContext, NotebookAction, NotebookCellAction, getEditorFromArgsOrActivePane } from 'vs/workbench/contrib/notebook/browser/controller/coreActions';
 import { CellEditState } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { CellKind, NOTEBOOK_EDITOR_CURSOR_BOUNDARY, NotebookSetting } from 'vs/workbench/contrib/notebook/common/notebookCommon';
-import { NOTEBOOK_CELL_LIST_FOCUSED, NOTEBOOK_EDITOR_EDITABLE, NOTEBOOK_EDITOR_FOCUSED } from 'vs/workbench/contrib/notebook/common/notebookContextKeys';
+import { NOTEBOOK_EDITOR_EDITABLE, NOTEBOOK_EDITOR_FOCUSED } from 'vs/workbench/contrib/notebook/common/notebookContextKeys';
 
 
 registerAction2(class extends NotebookAction {
@@ -222,11 +222,18 @@ registerAction2(class extends NotebookAction {
 				shortTitle: localize('apply2', 'Accept'),
 				icon: Codicon.check,
 				tooltip: localize('apply3', 'Accept Changes'),
-				keybinding: {
-					when: ContextKeyExpr.and(CTX_NOTEBOOK_CELL_CHAT_FOCUSED, CTX_INLINE_CHAT_FOCUSED),
-					weight: KeybindingWeight.EditorContrib + 10,
-					primary: KeyMod.CtrlCmd | KeyCode.Enter,
-				},
+				keybinding: [
+					{
+						when: ContextKeyExpr.and(CTX_NOTEBOOK_CELL_CHAT_FOCUSED, CTX_INLINE_CHAT_FOCUSED),
+						weight: KeybindingWeight.EditorContrib + 10,
+						primary: KeyMod.CtrlCmd | KeyCode.Enter,
+					},
+					{
+						when: ContextKeyExpr.and(CTX_NOTEBOOK_CELL_CHAT_FOCUSED, CTX_INLINE_CHAT_FOCUSED, CTX_NOTEBOOK_CHAT_USER_DID_EDIT),
+						weight: KeybindingWeight.EditorCore + 10,
+						primary: KeyCode.Escape
+					}
+				],
 				menu: [
 					{
 						id: MENU_CELL_CHAT_WIDGET_STATUS,
@@ -251,7 +258,7 @@ registerAction2(class extends NotebookAction {
 				title: localize('discard', 'Discard'),
 				icon: Codicon.discard,
 				keybinding: {
-					when: ContextKeyExpr.and(CTX_NOTEBOOK_CELL_CHAT_FOCUSED, CTX_INLINE_CHAT_FOCUSED, NOTEBOOK_CELL_LIST_FOCUSED),
+					when: ContextKeyExpr.and(CTX_NOTEBOOK_CELL_CHAT_FOCUSED, CTX_INLINE_CHAT_FOCUSED, CTX_NOTEBOOK_CHAT_USER_DID_EDIT.negate()),
 					weight: KeybindingWeight.EditorContrib,
 					primary: KeyCode.Escape
 				},

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

@@ -9,6 +9,7 @@ import { RawContextKey } from 'vs/platform/contextkey/common/contextkey';
 
 export const CTX_NOTEBOOK_CELL_CHAT_FOCUSED = new RawContextKey<boolean>('notebookCellChatFocused', false, localize('notebookCellChatFocused', "Whether the cell chat editor is focused"));
 export const CTX_NOTEBOOK_CHAT_HAS_ACTIVE_REQUEST = new RawContextKey<boolean>('notebookChatHasActiveRequest', false, localize('notebookChatHasActiveRequest', "Whether the cell chat editor has an active request"));
+export const CTX_NOTEBOOK_CHAT_USER_DID_EDIT = new RawContextKey<boolean>('notebookChatUserDidEdit', false, localize('notebookChatUserDidEdit', "Whether the user did changes ontop of the notebook cell chat"));
 export const MENU_CELL_CHAT_INPUT = MenuId.for('cellChatInput');
 export const MENU_CELL_CHAT_WIDGET = MenuId.for('cellChatWidget');
 export const MENU_CELL_CHAT_WIDGET_STATUS = MenuId.for('cellChatWidget.status');

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

@@ -13,7 +13,6 @@ import { Schemas } from 'vs/base/common/network';
 import { MovingAverage } from 'vs/base/common/numbers';
 import { StopWatch } from 'vs/base/common/stopwatch';
 import { assertType } from 'vs/base/common/types';
-import { URI } from 'vs/base/common/uri';
 import { generateUuid } from 'vs/base/common/uuid';
 import { IActiveCodeEditor } from 'vs/editor/browser/editorBrowser';
 import { CodeEditorWidget } from 'vs/editor/browser/widget/codeEditorWidget';
@@ -40,11 +39,12 @@ import { IInlineChatMessageAppender, InlineChatWidget } from 'vs/workbench/contr
 import { asProgressiveEdit, performAsyncTextEdit } from 'vs/workbench/contrib/inlineChat/browser/utils';
 import { CTX_INLINE_CHAT_LAST_RESPONSE_TYPE, EditMode, IInlineChatProgressItem, IInlineChatRequest, InlineChatResponseFeedbackKind, InlineChatResponseType } from 'vs/workbench/contrib/inlineChat/common/inlineChat';
 import { insertCell, runDeleteAction } from 'vs/workbench/contrib/notebook/browser/controller/cellOperations';
-import { CTX_NOTEBOOK_CELL_CHAT_FOCUSED, CTX_NOTEBOOK_CHAT_HAS_ACTIVE_REQUEST, MENU_CELL_CHAT_INPUT, MENU_CELL_CHAT_WIDGET, MENU_CELL_CHAT_WIDGET_FEEDBACK, MENU_CELL_CHAT_WIDGET_STATUS } from 'vs/workbench/contrib/notebook/browser/controller/chat/notebookChatContext';
+import { CTX_NOTEBOOK_CELL_CHAT_FOCUSED, CTX_NOTEBOOK_CHAT_HAS_ACTIVE_REQUEST, CTX_NOTEBOOK_CHAT_USER_DID_EDIT, MENU_CELL_CHAT_INPUT, MENU_CELL_CHAT_WIDGET, MENU_CELL_CHAT_WIDGET_FEEDBACK, MENU_CELL_CHAT_WIDGET_STATUS } from 'vs/workbench/contrib/notebook/browser/controller/chat/notebookChatContext';
 import { INotebookEditor, INotebookEditorContribution, INotebookViewZone, ScrollToRevealBehavior } from 'vs/workbench/contrib/notebook/browser/notebookBrowser';
 import { registerNotebookContribution } from 'vs/workbench/contrib/notebook/browser/notebookEditorExtensions';
 import { CellViewModel } from 'vs/workbench/contrib/notebook/browser/viewModel/notebookViewModelImpl';
 import { CellKind } from 'vs/workbench/contrib/notebook/common/notebookCommon';
+import { INotebookExecutionStateService, NotebookExecutionType } from 'vs/workbench/contrib/notebook/common/notebookExecutionStateService';
 
 
 
@@ -100,10 +100,21 @@ class NotebookChatWidget extends Disposable implements INotebookViewZone {
 		this.inlineChatWidget.focus();
 	}
 
-	async getEditingCellEditor() {
+	getEditingCell() {
+		return this._editingCell;
+	}
+
+	async getOrCreateEditingCell(): Promise<{ cell: CellViewModel; editor: IActiveCodeEditor } | undefined> {
 		if (this._editingCell) {
 			await this._notebookEditor.focusNotebookCell(this._editingCell, 'editor');
-			return this._notebookEditor.activeCodeEditor;
+			if (this._notebookEditor.activeCodeEditor?.hasModel()) {
+				return {
+					cell: this._editingCell,
+					editor: this._notebookEditor.activeCodeEditor
+				};
+			} else {
+				return undefined;
+			}
 		}
 
 		if (!this._notebookEditor.hasModel()) {
@@ -117,7 +128,14 @@ class NotebookChatWidget extends Disposable implements INotebookViewZone {
 		}
 
 		await this._notebookEditor.focusNotebookCell(this._editingCell, 'editor', { revealBehavior: ScrollToRevealBehavior.firstLine });
-		return this._notebookEditor.activeCodeEditor;
+		if (this._notebookEditor.activeCodeEditor?.hasModel()) {
+			return {
+				cell: this._editingCell,
+				editor: this._notebookEditor.activeCodeEditor
+			};
+		}
+
+		return undefined;
 	}
 
 	async discardChange() {
@@ -160,6 +178,8 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 	private _activeSession?: Session;
 	private readonly _ctxHasActiveRequest: IContextKey<boolean>;
 	private readonly _ctxCellWidgetFocused: IContextKey<boolean>;
+	private readonly _ctxUserDidEdit: IContextKey<boolean>;
+	private readonly _userEditingDisposables = this._register(new DisposableStore());
 	private readonly _ctxLastResponseType: IContextKey<undefined | InlineChatResponseType>;
 	private _widget: NotebookChatWidget | undefined;
 	private _widgetDisposableStore = this._register(new DisposableStore());
@@ -174,12 +194,14 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 		@IInlineChatSavingService private readonly _inlineChatSavingService: IInlineChatSavingService,
 		@IModelService private readonly _modelService: IModelService,
 		@ILanguageService private readonly _languageService: ILanguageService,
+		@INotebookExecutionStateService private _executionStateService: INotebookExecutionStateService,
 
 	) {
 		super();
 		this._ctxHasActiveRequest = CTX_NOTEBOOK_CHAT_HAS_ACTIVE_REQUEST.bindTo(this._contextKeyService);
 		this._ctxCellWidgetFocused = CTX_NOTEBOOK_CELL_CHAT_FOCUSED.bindTo(this._contextKeyService);
 		this._ctxLastResponseType = CTX_INLINE_CHAT_LAST_RESPONSE_TYPE.bindTo(this._contextKeyService);
+		this._ctxUserDidEdit = CTX_NOTEBOOK_CHAT_USER_DID_EDIT.bindTo(this._contextKeyService);
 	}
 
 	run(index: number, input: string | undefined, autoSend: boolean | undefined): void {
@@ -206,6 +228,10 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 	}
 
 	private _createWidget(index: number, input: string | undefined, autoSend: boolean | undefined) {
+		if (!this._notebookEditor.hasModel()) {
+			return;
+		}
+
 		// Clear the widget if it's already there
 		this._widgetDisposableStore.clear();
 
@@ -230,8 +256,9 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 			{ isSimpleWidget: true }
 		));
 
-		const inputBoxPath = `/notebook-chat-input-${NotebookChatController.counter++}`;
-		const inputUri = URI.from({ scheme: Schemas.untitled, path: inputBoxPath });
+		const inputBoxFragment = `notebook-chat-input-${NotebookChatController.counter++}`;
+		const notebookUri = this._notebookEditor.textModel.uri;
+		const inputUri = notebookUri.with({ scheme: Schemas.untitled, fragment: inputBoxFragment });
 		const result: ITextModel = this._modelService.createModel('', null, inputUri, false);
 		fakeParentEditor.setModel(result);
 
@@ -270,8 +297,9 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 				this._languageService
 			);
 
+			this._ctxCellWidgetFocused.set(true);
+
 			disposableTimeout(() => {
-				this._ctxCellWidgetFocused.set(true);
 				this._focusWidget();
 			}, 0, this._store);
 
@@ -330,6 +358,7 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 		}
 
 		this._notebookEditor.focusContainer(true);
+		this._notebookEditor.setFocus({ start: this._widget.afterModelPosition, end: this._widget.afterModelPosition });
 		this._notebookEditor.setSelections([{
 			start: this._widget.afterModelPosition,
 			end: this._widget.afterModelPosition
@@ -476,6 +505,22 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 						});
 					}
 				}
+
+				this._userEditingDisposables.clear();
+				// monitor user edits
+				const editingCell = this._widget.getEditingCell();
+				if (editingCell) {
+					this._userEditingDisposables.add(editingCell.model.onDidChangeContent(() => this._updateUserEditingState()));
+					this._userEditingDisposables.add(editingCell.model.onDidChangeLanguage(() => this._updateUserEditingState()));
+					this._userEditingDisposables.add(editingCell.model.onDidChangeMetadata(() => this._updateUserEditingState()));
+					this._userEditingDisposables.add(editingCell.model.onDidChangeInternalMetadata(() => this._updateUserEditingState()));
+					this._userEditingDisposables.add(editingCell.model.onDidChangeOutputs(() => this._updateUserEditingState()));
+					this._userEditingDisposables.add(this._executionStateService.onDidChangeExecution(e => {
+						if (e.type === NotebookExecutionType.cell && e.affectsCell(editingCell.uri)) {
+							this._updateUserEditingState();
+						}
+					}));
+				}
 			}
 		} catch (e) {
 			response = new ErrorResponse(e);
@@ -519,12 +564,14 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 		assertType(this._strategy);
 		assertType(this._widget);
 
-		const editor = await this._widget.getEditingCellEditor();
+		const editingCell = await this._widget.getOrCreateEditingCell();
 
-		if (!editor || !editor.hasModel()) {
+		if (!editingCell) {
 			return;
 		}
 
+		const editor = editingCell.editor;
+
 		const moreMinimalEdits = await this._editorWorkerService.computeMoreMinimalEdits(editor.getModel().uri, edits);
 		// this._log('edits from PROVIDER and after making them MORE MINIMAL', this._activeSession.provider.debugName, edits, moreMinimalEdits);
 
@@ -551,6 +598,10 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 		}
 	}
 
+	private _updateUserEditingState() {
+		this._ctxUserDidEdit.set(true);
+	}
+
 	async acceptSession() {
 		assertType(this._activeSession);
 		assertType(this._strategy);
@@ -615,6 +666,7 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 	discard() {
 		this._strategy?.cancel();
 		this._widget?.discardChange();
+		this.dismiss();
 	}
 
 	async feedbackLast(kind: InlineChatResponseFeedbackKind) {
@@ -627,6 +679,7 @@ export class NotebookChatController extends Disposable implements INotebookEdito
 
 	dismiss() {
 		this._ctxCellWidgetFocused.set(false);
+		this._ctxUserDidEdit.set(false);
 		this._sessionCtor?.cancel();
 		this._sessionCtor = undefined;
 		this._widget?.dispose();

src/vs/workbench/contrib/notebook/browser/controller/editActions.ts

@@ -32,6 +32,7 @@ import { IDialogService, IConfirmationResult } from 'vs/platform/dialogs/common/
 import { IConfigurationService } from 'vs/platform/configuration/common/configuration';
 import { ICodeEditor } from 'vs/editor/browser/editorBrowser';
 import { InlineChatController } from 'vs/workbench/contrib/inlineChat/browser/inlineChatController';
+import { CTX_INLINE_CHAT_FOCUSED } from 'vs/workbench/contrib/inlineChat/common/inlineChat';
 
 
 const CLEAR_ALL_CELLS_OUTPUTS_COMMAND_ID = 'notebook.clearAllCellsOutputs';
@@ -85,7 +86,8 @@ registerAction2(class EditCellAction extends NotebookCellAction {
 
 const quitEditCondition = ContextKeyExpr.and(
 	NOTEBOOK_EDITOR_FOCUSED,
-	InputFocusedContext
+	InputFocusedContext,
+	CTX_INLINE_CHAT_FOCUSED.toNegated()
 );
 registerAction2(class QuitEditCellAction extends NotebookCellAction {
 	constructor() {

src/vs/workbench/contrib/notebook/browser/viewModel/notebookViewModelImpl.ts

@@ -361,9 +361,6 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 		this._focused = focused;
 	}
 
-	/**
-	 * Empty selection will be turned to `null`
-	 */
 	validateRange(cellRange: ICellRange | null | undefined): ICellRange | null {
 		if (!cellRange) {
 			return null;
@@ -372,11 +369,7 @@ export class NotebookViewModel extends Disposable implements EditorFoldingStateD
 		const start = clamp(cellRange.start, 0, this.length);
 		const end = clamp(cellRange.end, 0, this.length);
 
-		if (start === end) {
-			return null;
-		}
-
-		if (start < end) {
+		if (start <= end) {
 			return { start, end };
 		} else {
 			return { start: end, end: start };

src/vs/workbench/contrib/notebook/common/notebookRange.ts

@@ -65,7 +65,7 @@ export function reduceCellRanges(ranges: ICellRange[]): ICellRange[] {
 		return [];
 	}
 
-	return sorted.reduce((prev: ICellRange[], curr) => {
+	const reduced = sorted.reduce((prev: ICellRange[], curr) => {
 		const last = prev[prev.length - 1];
 		if (last.end >= curr.start) {
 			last.end = Math.max(last.end, curr.end);
@@ -74,6 +74,13 @@ export function reduceCellRanges(ranges: ICellRange[]): ICellRange[] {
 		}
 		return prev;
 	}, [first] as ICellRange[]);
+
+	if (reduced.length > 1) {
+		// remove the (0, 0) range
+		return reduced.filter(range => !(range.start === range.end && range.start === 0));
+	}
+
+	return reduced;
 }
 
 export function cellRangesEqual(a: ICellRange[], b: ICellRange[]) {

src/vs/workbench/contrib/notebook/test/browser/notebookCommon.test.ts

@@ -399,6 +399,12 @@ suite('CellRange', function () {
 			{ start: 0, end: 4 }
 		]);
 	});
+
+	test('Reduce ranges 2, empty ranges', function () {
+		assert.deepStrictEqual(reduceCellRanges([{ start: 0, end: 0 }, { start: 0, end: 0 }]), [{ start: 0, end: 0 }]);
+		assert.deepStrictEqual(reduceCellRanges([{ start: 0, end: 0 }, { start: 1, end: 2 }]), [{ start: 1, end: 2 }]);
+		assert.deepStrictEqual(reduceCellRanges([{ start: 2, end: 2 }]), [{ start: 2, end: 2 }]);
+	});
 });
 
 suite('NotebookWorkingCopyTypeIdentifier', function () {

src/vs/workbench/contrib/notebook/test/browser/notebookSelection.test.ts

@@ -278,7 +278,7 @@ suite('NotebookCellList focus/selection', () => {
 			(editor, viewModel) => {
 				assert.deepStrictEqual(viewModel.validateRange(null), null);
 				assert.deepStrictEqual(viewModel.validateRange(undefined), null);
-				assert.deepStrictEqual(viewModel.validateRange({ start: 0, end: 0 }), null);
+				assert.deepStrictEqual(viewModel.validateRange({ start: 0, end: 0 }), { start: 0, end: 0 });
 				assert.deepStrictEqual(viewModel.validateRange({ start: 0, end: 2 }), { start: 0, end: 2 });
 				assert.deepStrictEqual(viewModel.validateRange({ start: 0, end: 3 }), { start: 0, end: 2 });
 				assert.deepStrictEqual(viewModel.validateRange({ start: -1, end: 3 }), { start: 0, end: 2 });