feat(deepnote): Enhance notebook resolution management and error handling

tkislan · tkislan · commit 44bd482eb25d · 2026-03-30T15:35:21.000Z
- Introduced `queueNotebookResolution` and `consumePendingNotebookResolution` methods in `DeepnoteNotebookManager` to manage transient notebook resolutions, improving the handling of notebook selections during file operations.
- Updated `DeepnoteNotebookSerializer` to prioritize queued notebook resolutions, ensuring more reliable notebook ID retrieval during deserialization.
- Refactored `DeepnoteExplorerView` to utilize a new `registerNotebookOpenIntent` method, streamlining the process of selecting and opening notebooks.
- Improved error handling in `DeepnoteServerStarter` to log warnings when disposing listeners fails, enhancing diagnostics during server operations.
- Adjusted unit tests to cover new functionality and ensure consistent behavior across notebook management processes.
diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts
@@ -405,7 +405,11 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension
         const existing = this.disposablesByFile.get(fileKey);
         if (existing) {
             for (const d of existing) {
-                d.dispose();
+                try {
+                    d.dispose();
+                } catch (ex) {
+                    logger.warn(`Error disposing listener for ${fileKey}`, ex);
+                }
             }
         }
         const disposables: IDisposable[] = [];
diff --git a/src/notebooks/deepnote/deepnoteExplorerView.ts b/src/notebooks/deepnote/deepnoteExplorerView.ts
@@ -260,7 +260,7 @@ export class DeepnoteExplorerView {
             await this.treeDataProvider.refreshNotebook(treeItem.context.projectId);
 
             // Optionally open the duplicated notebook
-            this.manager.selectNotebookForProject(treeItem.context.projectId, newNotebook.id);
+            this.registerNotebookOpenIntent(treeItem.context.projectId, newNotebook.id);
             const notebookUri = fileUri.with({ query: `notebook=${newNotebook.id}` });
             const document = await workspace.openNotebookDocument(notebookUri);
             await window.showNotebookDocument(document, {
@@ -508,7 +508,7 @@ export class DeepnoteExplorerView {
         await this.treeDataProvider.refreshNotebook(projectData.project.id);
 
         // Open the new notebook
-        this.manager.selectNotebookForProject(projectData.project.id, notebookId);
+        this.registerNotebookOpenIntent(projectData.project.id, notebookId);
         const notebookUri = fileUri.with({ query: `notebook=${notebookId}` });
         const document = await workspace.openNotebookDocument(notebookUri);
         await window.showNotebookDocument(document, {
@@ -517,6 +517,11 @@ export class DeepnoteExplorerView {
         });
     }
 
+    private registerNotebookOpenIntent(projectId: string, notebookId: string): void {
+        this.manager.queueNotebookResolution(projectId, notebookId);
+        this.manager.selectNotebookForProject(projectId, notebookId);
+    }
+
     private refreshExplorer(): void {
         this.treeDataProvider.refresh();
     }
@@ -537,7 +542,7 @@ export class DeepnoteExplorerView {
 
             console.log(`Selecting notebook in manager.`);
 
-            this.manager.selectNotebookForProject(context.projectId, context.notebookId);
+            this.registerNotebookOpenIntent(context.projectId, context.notebookId);
 
             console.log(`Opening notebook document.`, fileUri);
 
@@ -701,7 +706,7 @@ export class DeepnoteExplorerView {
 
             this.treeDataProvider.refresh();
 
-            this.manager.selectNotebookForProject(projectId, notebookId);
+            this.registerNotebookOpenIntent(projectId, notebookId);
 
             const notebookUri = fileUri.with({ query: `notebook=${notebookId}` });
             const document = await workspace.openNotebookDocument(notebookUri);
diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts
@@ -242,17 +242,6 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic
                 doc.notebookType === 'deepnote' && doc.uri.with({ query: '', fragment: '' }).toString() === uriString
         );
 
-        // Clear the global notebook selection so that any VS Code-triggered
-        // re-deserialization (e.g. from workspace.save) falls back to the
-        // active editor rather than a stale global selection.
-        for (const notebook of affectedNotebooks) {
-            const projectId = notebook.metadata?.deepnoteProjectId as string | undefined;
-            if (projectId) {
-                this.notebookManager.clearNotebookSelection(projectId);
-                break;
-            }
-        }
-
         for (const notebook of affectedNotebooks) {
             const nbKey = notebook.uri.toString();
             // main-file-sync always replaces any pending operation
diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts
@@ -135,9 +135,10 @@ suite('DeepnoteFileChangeWatcher', () => {
         mockDisposables = [];
 
         mockedNotebookManager = mock<IDeepnoteNotebookManager>();
+        when(mockedNotebookManager.consumePendingNotebookResolution(anything())).thenReturn(undefined);
         when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(validProject);
         when(mockedNotebookManager.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1');
-        when(mockedNotebookManager.clearNotebookSelection(anything())).thenReturn();
+        when(mockedNotebookManager.queueNotebookResolution(anything(), anything())).thenReturn();
         mockNotebookManager = instance(mockedNotebookManager);
 
         // Set up FileSystemWatcher mock
@@ -452,7 +453,7 @@ project:
 
         onDidChangeFile.fire(uri);
 
-        await waitFor(() => saveCount > 0);
+        await waitFor(() => applyEditCount > 0);
 
         // Only ONE applyEdit call (atomic: replaceCells + metadata in single WorkspaceEdit)
         assert.strictEqual(applyEditCount, 1, 'applyEdit should be called exactly once (atomic edit)');
@@ -1271,9 +1272,10 @@ project:
 
         setup(() => {
             reset(mockedNotebookManager);
+            when(mockedNotebookManager.consumePendingNotebookResolution(anything())).thenReturn(undefined);
             when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject);
             when(mockedNotebookManager.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1');
-            when(mockedNotebookManager.clearNotebookSelection(anything())).thenReturn();
+            when(mockedNotebookManager.queueNotebookResolution(anything(), anything())).thenReturn();
             resetCalls(mockedNotebookManager);
             workspaceSetCaptures = [];
             sinon.stub(watcher, 'applyNotebookEdits' as any).callsFake(async (...args: unknown[]) => {
@@ -1349,7 +1351,7 @@ project:
             assert.notInclude(byUri.get(uriNb2.toString()) ?? '', 'nb1-new');
         });
 
-        test('should clear notebook selection before processing file change', async () => {
+        test('should not clear notebook selection before processing file change', async () => {
             const basePath = Uri.file('/workspace/multi.deepnote');
             const uriNb1 = basePath.with({ query: 'a=1' });
             const uriNb2 = basePath.with({ query: 'b=2' });
@@ -1393,7 +1395,7 @@ project:
 
             await waitFor(() => applyEditCount >= 2);
 
-            verify(mockedNotebookManager.clearNotebookSelection('project-1')).once();
+            verify(mockedNotebookManager.clearNotebookSelection(anything())).never();
         });
 
         test('should not corrupt other notebooks when one notebook triggers a file change', async () => {
diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.ts b/src/notebooks/deepnote/deepnoteNotebookManager.ts
@@ -3,6 +3,13 @@ import { injectable } from 'inversify';
 import { IDeepnoteNotebookManager, ProjectIntegration } from '../types';
 import type { DeepnoteProject } from '../../platform/deepnote/deepnoteTypes';
 
+const pendingNotebookResolutionTtlMs = 2_000;
+
+interface PendingNotebookResolution {
+    notebookId: string;
+    queuedAt: number;
+}
+
 /**
  * Centralized manager for tracking Deepnote notebook selections and project state.
  * Manages per-project state including current selections and project data caching.
@@ -11,17 +18,36 @@ import type { DeepnoteProject } from '../../platform/deepnote/deepnoteTypes';
 export class DeepnoteNotebookManager implements IDeepnoteNotebookManager {
     private readonly currentNotebookId = new Map<string, string>();
     private readonly originalProjects = new Map<string, DeepnoteProject>();
+    private readonly pendingNotebookResolutions = new Map<string, PendingNotebookResolution[]>();
     private readonly projectsWithInitNotebookRun = new Set<string>();
     private readonly selectedNotebookByProject = new Map<string, string>();
 
     /**
-     * Clears the notebook selection for a project so that subsequent
-     * deserializations fall back to the active editor or open documents.
+     * Clears the remembered notebook selection and any pending resolution hints for a project.
      */
     clearNotebookSelection(projectId: string): void {
+        this.pendingNotebookResolutions.delete(projectId);
         this.selectedNotebookByProject.delete(projectId);
     }
 
+    /**
+     * Consumes the next short-lived notebook resolution hint for a project.
+     * These hints are queued immediately before operations that trigger a
+     * deserialize without explicit URI context.
+     */
+    consumePendingNotebookResolution(projectId: string): string | undefined {
+        const pendingResolutions = this.getValidPendingNotebookResolutions(projectId);
+        const nextResolution = pendingResolutions.shift();
+
+        if (pendingResolutions.length > 0) {
+            this.pendingNotebookResolutions.set(projectId, pendingResolutions);
+        } else {
+            this.pendingNotebookResolutions.delete(projectId);
+        }
+
+        return nextResolution?.notebookId;
+    }
+
     /**
      * Gets the currently selected notebook ID for a project.
      * @param projectId Project identifier
@@ -50,9 +76,24 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager {
     }
 
     /**
-     * Associates a notebook ID with a project to remember user's notebook selection.
-     * When a Deepnote project contains multiple notebooks, this mapping persists the user's
-     * choice so we can automatically open the same notebook on subsequent file opens.
+     * Queues a short-lived notebook resolution hint for the next deserialize.
+     *
+     * @param projectId - The project ID that identifies the Deepnote project
+     * @param notebookId - The notebook ID the next deserialize should resolve to
+     */
+    queueNotebookResolution(projectId: string, notebookId: string): void {
+        const pendingResolutions = this.getValidPendingNotebookResolutions(projectId);
+
+        pendingResolutions.push({
+            notebookId,
+            queuedAt: Date.now()
+        });
+
+        this.pendingNotebookResolutions.set(projectId, pendingResolutions);
+    }
+
+    /**
+     * Associates a notebook ID with a project to remember the user's last explicit selection.
      *
      * @param projectId - The project ID that identifies the Deepnote project
      * @param notebookId - The ID of the selected notebook within the project
@@ -134,4 +175,18 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager {
     markInitNotebookAsRun(projectId: string): void {
         this.projectsWithInitNotebookRun.add(projectId);
     }
+
+    private getValidPendingNotebookResolutions(projectId: string): PendingNotebookResolution[] {
+        const cutoffTime = Date.now() - pendingNotebookResolutionTtlMs;
+        const pendingResolutions = (this.pendingNotebookResolutions.get(projectId) ?? []).filter(
+            (resolution) => resolution.queuedAt >= cutoffTime
+        );
+
+        if (pendingResolutions.length === 0) {
+            this.pendingNotebookResolutions.delete(projectId);
+            return [];
+        }
+
+        return pendingResolutions;
+    }
 }
diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts
@@ -93,6 +93,39 @@ suite('DeepnoteNotebookManager', () => {
         });
     });
 
+    suite('consumePendingNotebookResolution', () => {
+        test('should return undefined when no pending resolution exists', () => {
+            const result = manager.consumePendingNotebookResolution('unknown-project');
+
+            assert.strictEqual(result, undefined);
+        });
+
+        test('should consume queued notebook resolutions in order', () => {
+            manager.queueNotebookResolution('project-123', 'notebook-1');
+            manager.queueNotebookResolution('project-123', 'notebook-2');
+
+            assert.strictEqual(manager.consumePendingNotebookResolution('project-123'), 'notebook-1');
+            assert.strictEqual(manager.consumePendingNotebookResolution('project-123'), 'notebook-2');
+            assert.strictEqual(manager.consumePendingNotebookResolution('project-123'), undefined);
+        });
+
+        test('should keep pending resolutions isolated per project', () => {
+            manager.queueNotebookResolution('project-1', 'notebook-1');
+            manager.queueNotebookResolution('project-2', 'notebook-2');
+
+            assert.strictEqual(manager.consumePendingNotebookResolution('project-1'), 'notebook-1');
+            assert.strictEqual(manager.consumePendingNotebookResolution('project-2'), 'notebook-2');
+        });
+    });
+
+    suite('queueNotebookResolution', () => {
+        test('should queue a notebook resolution for later consumption', () => {
+            manager.queueNotebookResolution('project-123', 'notebook-456');
+
+            assert.strictEqual(manager.consumePendingNotebookResolution('project-123'), 'notebook-456');
+        });
+    });
+
     suite('selectNotebookForProject', () => {
         test('should store notebook selection for project', () => {
             manager.selectNotebookForProject('project-123', 'notebook-456');
@@ -148,6 +181,13 @@ suite('DeepnoteNotebookManager', () => {
                 manager.clearNotebookSelection('unknown-project');
             });
         });
+
+        test('should clear pending notebook resolutions for a project', () => {
+            manager.queueNotebookResolution('project-123', 'notebook-456');
+            manager.clearNotebookSelection('project-123');
+
+            assert.strictEqual(manager.consumePendingNotebookResolution('project-123'), undefined);
+        });
     });
 
     suite('storeOriginalProject', () => {
diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts
@@ -1,7 +1,7 @@
 import type { DeepnoteBlock, DeepnoteFile, DeepnoteSnapshot } from '@deepnote/blocks';
 import { deserializeDeepnoteFile, isExecutableBlock, serializeDeepnoteSnapshot } from '@deepnote/blocks';
 import { inject, injectable, optional } from 'inversify';
-import { l10n, window, workspace, type CancellationToken, type NotebookData, type NotebookSerializer } from 'vscode';
+import { l10n, workspace, type CancellationToken, type NotebookData, type NotebookSerializer } from 'vscode';
 
 import { logger } from '../../platform/logging';
 import { IDeepnoteNotebookManager } from '../types';
@@ -62,7 +62,8 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer {
     /**
      * Deserializes a Deepnote YAML file into VS Code notebook format.
      * Parses YAML and converts the selected notebook's blocks to cells.
-     * The notebook to deserialize must be pre-selected and stored in the manager.
+     * Notebook resolution prefers an explicit notebook ID, then transient
+     * resolver state, and finally a deterministic default notebook.
      * @param content Raw file content as bytes
      * @param token Cancellation token (unused)
      * @returns Promise resolving to notebook data
@@ -186,37 +187,29 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer {
     }
 
     /**
-     * Finds the notebook ID to deserialize by checking the manager's stored selection.
-     * The notebook ID should be set via selectNotebookForProject before opening the document.
+     * Finds the notebook ID to deserialize without relying on active-editor state.
+     * Prefers a pending resolution hint, then current/open-document state.
      * @param projectId The project ID to find a notebook for
      * @returns The notebook ID to deserialize, or undefined if none found
      */
     findCurrentNotebookId(projectId: string): string | undefined {
-        // Check the manager's stored selection first - this is set explicitly when the user
-        // picks a notebook from the explorer, and must take priority over the active editor
-        const storedNotebookId = this.notebookManager.getTheSelectedNotebookForAProject(projectId);
+        const pendingNotebookId = this.notebookManager.consumePendingNotebookResolution(projectId);
+        const openNotebookIds = this.findOpenNotebookIds(projectId);
+        const currentNotebookId = this.notebookManager.getCurrentNotebookId(projectId);
 
-        if (storedNotebookId) {
-            return storedNotebookId;
+        if (pendingNotebookId) {
+            return pendingNotebookId;
         }
 
-        // Fallback: prefer the active notebook editor when it matches the project
-        const activeEditorNotebook = window.activeNotebookEditor?.notebook;
-
-        if (
-            activeEditorNotebook?.notebookType === 'deepnote' &&
-            activeEditorNotebook.metadata?.deepnoteProjectId === projectId &&
-            activeEditorNotebook.metadata?.deepnoteNotebookId
-        ) {
-            return activeEditorNotebook.metadata.deepnoteNotebookId;
+        if (currentNotebookId && (openNotebookIds.length === 0 || openNotebookIds.includes(currentNotebookId))) {
+            return currentNotebookId;
         }
 
-        // Last fallback: Check if there's an active notebook document for this project
-        const openNotebook = workspace.notebookDocuments.find(
-            (doc) => doc.notebookType === 'deepnote' && doc.metadata?.deepnoteProjectId === projectId
-        );
+        if (openNotebookIds.length === 1) {
+            return openNotebookIds[0];
+        }
 
-        return openNotebook?.metadata?.deepnoteNotebookId;
+        return undefined;
     }
 
     /**
@@ -264,7 +257,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer {
             logger.debug('SerializeNotebook: Got and cloned original project');
 
             const notebookId =
-                data.metadata?.deepnoteNotebookId || this.notebookManager.getTheSelectedNotebookForAProject(projectId);
+                data.metadata?.deepnoteNotebookId || this.notebookManager.getCurrentNotebookId(projectId);
 
             if (!notebookId) {
                 throw new Error('Cannot determine which notebook to save');
@@ -353,6 +346,11 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer {
 
             // Store the updated project back so subsequent saves start from correct state
             this.notebookManager.storeOriginalProject(projectId, originalProject, notebookId);
+            const openNotebookIdsAtSerialize = this.findOpenNotebookIds(projectId);
+
+            if (openNotebookIdsAtSerialize.length === 0) {
+                this.notebookManager.queueNotebookResolution(projectId, notebookId);
+            }
 
             logger.debug('SerializeNotebook: Serializing to YAML');
 
@@ -555,6 +553,21 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer {
         return false;
     }
 
+    private findOpenNotebookIds(projectId: string): string[] {
+        return [
+            ...new Set(
+                workspace.notebookDocuments
+                    .filter(
+                        (doc) =>
+                            doc.notebookType === 'deepnote' &&
+                            doc.metadata?.deepnoteProjectId === projectId &&
+                            typeof doc.metadata?.deepnoteNotebookId === 'string'
+                    )
+                    .map((doc) => doc.metadata.deepnoteNotebookId as string)
+            )
+        ];
+    }
+
     /**
      * Finds the default notebook to open when no selection is made.
      * @param file
diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts
diff --git a/src/notebooks/types.ts b/src/notebooks/types.ts

Original file line number	Diff line number	Diff line change
`@@ -405,7 +405,11 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension`
`405`	`405`	`const existing = this.disposablesByFile.get(fileKey);`
`406`	`406`	`if (existing) {`
`407`	`407`	`for (const d of existing) {`
`408`		`- d.dispose();`
	`408`	`+ try {`
	`409`	`+ d.dispose();`
	`410`	`+ } catch (ex) {`
	`411`	+ logger.warn(`Error disposing listener for ${fileKey}`, ex);
	`412`	`+ }`
`409`	`413`	`}`
`410`	`414`	`}`
`411`	`415`	`const disposables: IDisposable[] = [];`