wait for diffing to stop before initializing, fix indexing

Author: amunger

src/vs/workbench/contrib/chat/browser/chatEditing/chatEditingModifiedNotebookEntry.ts

@@ -66,8 +66,12 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
 	 */
 	override initialContent: string;
 	/**
-	 * Whether we're in the process of applying edits.
+	 * Whether we're still generating diffs from a response.
 	 */
+	private _isProcessingResponse = observableValue<boolean>('isProcessingResponse', false);
+	get isProcessingResponse(): IObservable<boolean> {
+		return this._isProcessingResponse;
+	}
 	private _isEditFromUs: boolean = false;
 	/**
 	 * Whether all edits are from us, e.g. is possible a user has made edits, then this will be false.
@@ -216,6 +220,7 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
 		}
 		const cellsDiffInfo: CellDiffInfo[] = [];
 		try {
+			this._isProcessingResponse.set(true, undefined);
 			const notebookDiff = await this.notebookEditorWorkerService.computeDiff(this.originalURI, this.modifiedURI);
 			if (id !== this.computeRequestId) {
 				return;
@@ -226,6 +231,8 @@ export class ChatEditingModifiedNotebookEntry extends AbstractChatEditingModifie
 			}
 		} catch (ex) {
 			this.loggingService.error('Notebook Chat', 'Error computing diff:\n' + ex);
+		} finally {
+			this._isProcessingResponse.set(false, undefined);
 		}
 		this.initializeModelsFromDiffImpl(cellsDiffInfo);
 	}

src/vs/workbench/contrib/chat/browser/chatEditing/notebook/chatEditingNotebookEditorIntegration.ts

@@ -11,7 +11,6 @@ import { assertType } from '../../../../../../base/common/types.js';
 import { LineRange } from '../../../../../../editor/common/core/lineRange.js';
 import { Range } from '../../../../../../editor/common/core/range.js';
 import { nullDocumentDiff } from '../../../../../../editor/common/diff/documentDiffProvider.js';
-import { PrefixSumComputer } from '../../../../../../editor/common/model/prefixSumComputer.js';
 import { localize } from '../../../../../../nls.js';
 import { MenuId } from '../../../../../../platform/actions/common/actions.js';
 import { IInstantiationService } from '../../../../../../platform/instantiation/common/instantiation.js';
@@ -25,15 +24,16 @@ import { INotebookEditorService } from '../../../../notebook/browser/services/no
 import { NotebookCellTextModel } from '../../../../notebook/common/model/notebookCellTextModel.js';
 import { NotebookTextModel } from '../../../../notebook/common/model/notebookTextModel.js';
 import { ChatAgentLocation, IChatAgentService } from '../../../common/chatAgents.js';
-import { IModifiedFileEntry, IModifiedFileEntryChangeHunk, IModifiedFileEntryEditorIntegration } from '../../../common/chatEditingService.js';
+import { IModifiedFileEntryChangeHunk, IModifiedFileEntryEditorIntegration } from '../../../common/chatEditingService.js';
 import { ChatEditingCodeEditorIntegration, IDocumentDiff2 } from '../chatEditingCodeEditorIntegration.js';
+import { ChatEditingModifiedNotebookEntry } from '../chatEditingModifiedNotebookEntry.js';
 import { countChanges, ICellDiffInfo, sortCellChanges } from './notebookCellChanges.js';
 
 export class ChatEditingNotebookEditorIntegration extends Disposable implements IModifiedFileEntryEditorIntegration {
 	private integration: ChatEditingNotebookEditorWidgetIntegration;
 	private notebookEditor: INotebookEditor;
 	constructor(
-		_entry: IModifiedFileEntry,
+		_entry: ChatEditingModifiedNotebookEntry,
 		editor: IEditorPane,
 		notebookModel: NotebookTextModel,
 		originalModel: NotebookTextModel,
@@ -88,14 +88,12 @@ class ChatEditingNotebookEditorWidgetIntegration extends Disposable implements I
 	private readonly _currentChange = observableValue<{ change: ICellDiffInfo; index: number } | undefined>(this, undefined);
 	readonly currentChange: IObservable<{ change: ICellDiffInfo; index: number } | undefined> = this._currentChange;
 
-	private diffIndexPrefixSum: PrefixSumComputer = new PrefixSumComputer(new Uint32Array());
-
 	private readonly cellEditorIntegrations = new Map<NotebookCellTextModel, { integration: ChatEditingCodeEditorIntegration; diff: ISettableObservable<IDocumentDiff2> }>();
 
 	private readonly insertDeleteDecorators: IObservable<{ insertedCellDecorator: NotebookInsertedCellDecorator; deletedCellDecorator: NotebookDeletedCellDecorator } | undefined>;
 
 	constructor(
-		private readonly _entry: IModifiedFileEntry,
+		private readonly _entry: ChatEditingModifiedNotebookEntry,
 		private readonly notebookEditor: INotebookEditor,
 		private readonly notebookModel: NotebookTextModel,
 		originalModel: NotebookTextModel,
@@ -130,28 +128,15 @@ class ChatEditingNotebookEditorWidgetIntegration extends Disposable implements I
 			}
 		}));
 
-		// INIT when not streaming anymore, once per request, and when having changes
+		// INIT when not streaming nor diffing the response anymore, once per request, and when having changes
 		let lastModifyingRequestId: string | undefined;
 		this._store.add(autorun(r => {
 
 			if (!_entry.isCurrentlyBeingModifiedBy.read(r)
+				&& !_entry.isProcessingResponse.read(r)
 				&& lastModifyingRequestId !== _entry.lastModifyingRequestId
 				&& cellChanges.read(r).some(c => c.type !== 'unchanged' && !c.diff.read(r).identical)
 			) {
-				lastModifyingRequestId = _entry.lastModifyingRequestId;
-
-				const sortedCellChanges = sortCellChanges(cellChanges.read(r));
-				const values = new Uint32Array(sortedCellChanges.length);
-				for (let i = 0; i < sortedCellChanges.length; i++) {
-					const change = sortedCellChanges[i];
-					values[i] = change.type === 'insert' ? 1
-						: change.type === 'delete' ? 1
-							: change.type === 'modified' ? change.diff.read(r).changes.length
-								: 0;
-				}
-
-				this.diffIndexPrefixSum = new PrefixSumComputer(values);
-
 				this.reveal(true);
 			}
 		}));
@@ -224,17 +209,28 @@ class ChatEditingNotebookEditorWidgetIntegration extends Disposable implements I
 
 		this._register(autorun(r => {
 			const currentChange = this.currentChange.read(r);
-			if (currentChange) {
-				const indexInChange = currentChange.index;
-				const cellIndex = currentChange.change.modifiedCellIndex ?? currentChange.change.originalCellIndex;
-
-				const changesBeforeCell = cellIndex !== undefined && cellIndex > 0 ?
-					this.diffIndexPrefixSum.getPrefixSum(cellIndex - 1) : 0;
-
-				this._currentIndex.set(changesBeforeCell + indexInChange, undefined);
-			} else {
+			if (!currentChange) {
 				this._currentIndex.set(-1, undefined);
+				return;
 			}
+
+			let index = 0;
+			const sortedCellChanges = sortCellChanges(cellChanges.read(r));
+			for (const change of sortedCellChanges) {
+				if (currentChange && currentChange.change === change) {
+					if (change.type === 'modified') {
+						index += currentChange.index;
+					}
+					break;
+				}
+				if (change.type === 'insert' || change.type === 'delete') {
+					index++;
+				} else if (change.type === 'modified') {
+					index += change.diff.read(r).changes.length;
+				}
+			}
+
+			this._currentIndex.set(index, undefined);
 		}));
 
 		this.insertDeleteDecorators = derivedWithStore((r, store) => {