refactor(deepnote): Simplify notebook edit application logic

- Introduced a new `applyNotebookEdits` method in `DeepnoteFileChangeWatcher` to centralize the application of notebook edits, improving code readability and maintainability.
- Updated existing calls to `workspace.applyEdit` to utilize the new method, reducing redundancy in the codebase.
- Adjusted unit tests to reflect changes in the edit application process, ensuring consistent behavior across the application.

Author: tkislan

src/notebooks/deepnote/deepnoteFileChangeWatcher.ts

@@ -122,6 +122,13 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic
         }
     }
 
+    protected async applyNotebookEdits(uri: Uri, edits: NotebookEdit[]): Promise<boolean> {
+        const wsEdit = new WorkspaceEdit();
+        wsEdit.set(uri, edits);
+
+        return workspace.applyEdit(wsEdit);
+    }
+
     private clearAllTimers(): void {
         for (const timer of this.debounceTimers.values()) {
             clearTimeout(timer);
@@ -350,9 +357,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic
         }
 
         // Apply the edit to update in-memory cells immediately (responsive UX).
-        const wsEdit = new WorkspaceEdit();
-        wsEdit.set(notebook.uri, edits);
-        const applied = await workspace.applyEdit(wsEdit);
+        const applied = await this.applyNotebookEdits(notebook.uri, edits);
 
         if (!applied) {
             logger.warn(`[FileChangeWatcher] Failed to apply edit: ${notebook.uri.path}`);
@@ -501,9 +506,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic
                 }
             }
             if (metadataEdits.length > 0) {
-                const wsEdit = new WorkspaceEdit();
-                wsEdit.set(notebook.uri, metadataEdits);
-                await workspace.applyEdit(wsEdit);
+                await this.applyNotebookEdits(notebook.uri, metadataEdits);
             }
 
             logger.info(`[FileChangeWatcher] Updated notebook outputs via execution API: ${notebook.uri.path}`);
@@ -531,9 +534,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic
             );
         }
 
-        const wsEdit = new WorkspaceEdit();
-        wsEdit.set(notebook.uri, replaceEdits);
-        const applied = await workspace.applyEdit(wsEdit);
+        const applied = await this.applyNotebookEdits(notebook.uri, replaceEdits);
 
         if (!applied) {
             logger.warn(`[FileChangeWatcher] Failed to apply snapshot outputs: ${notebook.uri.path}`);
@@ -552,9 +553,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic
                 })
             );
         }
-        const metaEdit = new WorkspaceEdit();
-        metaEdit.set(notebook.uri, metadataEdits);
-        await workspace.applyEdit(metaEdit);
+        await this.applyNotebookEdits(notebook.uri, metadataEdits);
 
         // Save to sync mtime — mark as self-write first
         this.markSelfWrite(notebook.uri);

src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts

@@ -9,8 +9,7 @@ import {
     NotebookCellKind,
     NotebookDocument,
     NotebookEdit,
-    Uri,
-    WorkspaceEdit
+    Uri
 } from 'vscode';
 
 import type { IControllerRegistration } from '../controllers/types';
@@ -94,6 +93,11 @@ const rapidChangeIntervalMs = 100;
 const autoSaveGraceMs = 200;
 const postSnapshotReadGraceMs = 100;
 
+interface NotebookEditCapture {
+    uriKey: string;
+    cellSourceJoined: string;
+}
+
 /**
  * Polls until a condition is met or a timeout is reached.
  */
@@ -1263,7 +1267,7 @@ project:
     });
 
     suite('multi-notebook file sync', () => {
-        let workspaceSetCaptures: { uriKey: string; cellSourceJoined: string }[] = [];
+        let workspaceSetCaptures: NotebookEditCapture[] = [];
 
         setup(() => {
             reset(mockedNotebookManager);
@@ -1272,26 +1276,21 @@ project:
             when(mockedNotebookManager.clearNotebookSelection(anything())).thenReturn();
             resetCalls(mockedNotebookManager);
             workspaceSetCaptures = [];
-            when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall((wsEdit: WorkspaceEdit) => {
-                applyEditCount++;
-
-                const ext = wsEdit as WorkspaceEdit & { notebookEntries(): [Uri, NotebookEdit[]][] };
+            sinon.stub(watcher, 'applyNotebookEdits' as any).callsFake(async (...args: unknown[]) => {
+                const uri = args[0] as Uri;
+                const edits = args[1] as NotebookEdit[];
 
-                for (const [uri, edits] of ext.notebookEntries()) {
-                    if (!edits.length) {
-                        continue;
-                    }
+                applyEditCount++;
 
-                    const firstEdit = edits[0];
-                    if (firstEdit?.newCells && firstEdit.newCells.length > 0) {
-                        workspaceSetCaptures.push({
-                            uriKey: uri.toString(),
-                            cellSourceJoined: firstEdit.newCells.map((c) => c.value).join('\n')
-                        });
-                    }
+                const replaceCellsEdit = edits.find((e) => e.newCells?.length > 0);
+                if (replaceCellsEdit) {
+                    workspaceSetCaptures.push({
+                        uriKey: uri.toString(),
+                        cellSourceJoined: replaceCellsEdit.newCells.map((c: any) => c.value).join('\n')
+                    });
                 }
 
-                return Promise.resolve(true);
+                return true;
             });
         });
 

src/test/mocks/vsc/extHostedTypes.ts

@@ -707,7 +707,6 @@ export namespace vscMockExtHostedTypes {
         private _seqPool: number = 0;
 
         private _resourceEdits: { seq: number; from: vscUri.URI; to: vscUri.URI }[] = [];
-        private _notebookEdits = new Map<string, { uri: vscUri.URI; edits: NotebookEdit[] }>();
         private _textEdits = new Map<string, { seq: number; uri: vscUri.URI; edits: TextEdit[] }>();
 
         // createResource(uri: vscode.Uri): void {
@@ -760,20 +759,7 @@ export namespace vscMockExtHostedTypes {
         }
 
         has(uri: vscUri.URI): boolean {
-            return this._textEdits.has(uri.toString()) || this._notebookEdits.has(uri.toString());
-        }
-
-        /**
-         * Notebook edits passed to {@link set} (for tests that inspect edits via workspace.applyEdit).
-         */
-        notebookEntries(): [vscUri.URI, NotebookEdit[]][] {
-            const res: [vscUri.URI, NotebookEdit[]][] = [];
-
-            this._notebookEdits.forEach((value) => {
-                res.push([value.uri, value.edits.slice()]);
-            });
-
-            return res.slice();
+            return this._textEdits.has(uri.toString());
         }
 
         set(uri: vscode.Uri, edits: NotebookEdit[]): void;
@@ -793,54 +779,17 @@ export namespace vscMockExtHostedTypes {
                       | [TextEdit | SnippetTextEdit, vscode.WorkspaceEditEntryMetadata]
                   )[]
         ): void {
-            const uriKey = uri.toString();
-
-            if (!edits) {
-                const data = this._textEdits.get(uriKey);
-                if (data) {
-                    // @ts-ignore
-                    data.edits = undefined;
-                }
-                this._notebookEdits.delete(uriKey);
-
-                return;
-            }
-
-            if (edits.length === 0) {
-                let data = this._textEdits.get(uriKey);
-                if (!data) {
-                    data = { seq: this._seqPool++, uri, edits: [] };
-                    this._textEdits.set(uriKey, data);
-                }
-                data.edits = [];
-
-                return;
-            }
-
-            const first = edits[0];
-            const firstEdit = Array.isArray(first) ? first[0] : first;
-
-            if (firstEdit instanceof NotebookEdit) {
-                const normalized: NotebookEdit[] = edits.map((item) => {
-                    if (Array.isArray(item)) {
-                        return item[0] as NotebookEdit;
-                    }
-
-                    return item as NotebookEdit;
-                });
-
-                this._notebookEdits.set(uriKey, { uri, edits: normalized });
-
-                return;
-            }
-
-            let data = this._textEdits.get(uriKey);
+            let data = this._textEdits.get(uri.toString());
             if (!data) {
                 data = { seq: this._seqPool++, uri, edits: [] };
-                this._textEdits.set(uriKey, data);
+                this._textEdits.set(uri.toString(), data);
+            }
+            if (!edits) {
+                // @ts-ignore
+                data.edits = undefined;
+            } else {
+                //data.edits = edits.slice(0);
             }
-
-            data.edits = edits.slice(0) as TextEdit[];
         }
 
         get(uri: vscUri.URI): TextEdit[] {
@@ -876,7 +825,7 @@ export namespace vscMockExtHostedTypes {
         }
 
         get size(): number {
-            return this._textEdits.size + this._notebookEdits.size + this._resourceEdits.length;
+            return this._textEdits.size + this._resourceEdits.length;
         }
 
         toJSON(): any {