From 6802960643401d1c278dd4d34da196e16fe6c6a2 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 18 Mar 2026 19:08:34 +0000 Subject: [PATCH 01/47] fix: Fix deepnote notebook deserializer --- src/notebooks/deepnote/deepnoteSerializer.ts | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 17443e93b9..69637db668 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -188,7 +188,15 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { * @returns The notebook ID to deserialize, or undefined if none found */ findCurrentNotebookId(projectId: string): string | undefined { - // Prefer the active notebook editor when it matches the project + // 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); + + if (storedNotebookId) { + return storedNotebookId; + } + + // Fallback: prefer the active notebook editor when it matches the project const activeEditorNotebook = window.activeNotebookEditor?.notebook; if ( @@ -199,14 +207,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return activeEditorNotebook.metadata.deepnoteNotebookId; } - // Check the manager's stored selection - this should be set when opening from explorer - const storedNotebookId = this.notebookManager.getTheSelectedNotebookForAProject(projectId); - - if (storedNotebookId) { - return storedNotebookId; - } - - // Fallback: Check if there's an active notebook document for this project + // 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 ); From 2de7885ff4a08db49042a1582d6c805076e553f8 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 18 Mar 2026 19:12:10 +0000 Subject: [PATCH 02/47] Add tests --- .../deepnote/deepnoteSerializer.unit.test.ts | 79 ++++++++++++++++--- 1 file changed, 68 insertions(+), 11 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index c3332f974d..53334ff5c3 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -275,11 +275,9 @@ project: assert.strictEqual(result2, 'notebook-2'); }); - test('should prioritize active notebook editor over stored selection', () => { - // Store a selection for the project + test('should prioritize stored selection over active editor', () => { manager.selectNotebookForProject('project-123', 'stored-notebook'); - // Mock the active notebook editor to return a different notebook const mockActiveNotebook = { notebookType: 'deepnote', metadata: { @@ -294,14 +292,30 @@ project: const result = serializer.findCurrentNotebookId('project-123'); - // Should return the active editor's notebook, not the stored one + assert.strictEqual(result, 'stored-notebook'); + }); + + test('should return active editor notebook when no stored selection exists', () => { + const mockActiveNotebook = { + notebookType: 'deepnote', + metadata: { + deepnoteProjectId: 'project-123', + deepnoteNotebookId: 'active-editor-notebook' + } + }; + + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ + notebook: mockActiveNotebook + } as any); + + const result = serializer.findCurrentNotebookId('project-123'); + assert.strictEqual(result, 'active-editor-notebook'); }); test('should ignore active editor when project ID does not match', () => { manager.selectNotebookForProject('project-123', 'stored-notebook'); - // Mock active editor with a different project const mockActiveNotebook = { notebookType: 'deepnote', metadata: { @@ -316,14 +330,12 @@ project: const result = serializer.findCurrentNotebookId('project-123'); - // Should fall back to stored selection since active editor is for different project assert.strictEqual(result, 'stored-notebook'); }); test('should ignore active editor when notebook type is not deepnote', () => { manager.selectNotebookForProject('project-123', 'stored-notebook'); - // Mock active editor with non-deepnote notebook type const mockActiveNotebook = { notebookType: 'jupyter-notebook', metadata: { @@ -338,19 +350,16 @@ project: const result = serializer.findCurrentNotebookId('project-123'); - // Should fall back to stored selection since active editor is not a deepnote notebook assert.strictEqual(result, 'stored-notebook'); }); test('should ignore active editor when notebook ID is missing', () => { manager.selectNotebookForProject('project-123', 'stored-notebook'); - // Mock active editor without notebook ID in metadata const mockActiveNotebook = { notebookType: 'deepnote', metadata: { deepnoteProjectId: 'project-123' - // Missing deepnoteNotebookId } }; @@ -360,9 +369,57 @@ project: const result = serializer.findCurrentNotebookId('project-123'); - // Should fall back to stored selection since active editor has no notebook ID assert.strictEqual(result, 'stored-notebook'); }); + + test('switching notebooks: selecting a different notebook while one is open should return the new selection', () => { + manager.selectNotebookForProject('project-123', 'notebook-A'); + + const mockActiveNotebook = { + notebookType: 'deepnote', + metadata: { + deepnoteProjectId: 'project-123', + deepnoteNotebookId: 'notebook-A' + } + }; + + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ + notebook: mockActiveNotebook + } as any); + + assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-A'); + + manager.selectNotebookForProject('project-123', 'notebook-B'); + + assert.strictEqual( + serializer.findCurrentNotebookId('project-123'), + 'notebook-B', + 'Should return the newly selected notebook, not the one currently in the active editor' + ); + }); + + test('switching notebooks: rapidly switching between three notebooks should always return the latest selection', () => { + const mockActiveNotebook = { + notebookType: 'deepnote', + metadata: { + deepnoteProjectId: 'project-123', + deepnoteNotebookId: 'notebook-1' + } + }; + + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ + notebook: mockActiveNotebook + } as any); + + manager.selectNotebookForProject('project-123', 'notebook-1'); + assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-1'); + + manager.selectNotebookForProject('project-123', 'notebook-2'); + assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-2'); + + manager.selectNotebookForProject('project-123', 'notebook-3'); + assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-3'); + }); }); suite('component integration', () => { From 3a418c1f04cdb789aef629a50935613aebdbfbc0 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 19 Mar 2026 14:00:11 +0000 Subject: [PATCH 03/47] Fix file change watcher --- .../deepnote/deepnoteFileChangeWatcher.ts | 52 ++++++++++++++++--- 1 file changed, 44 insertions(+), 8 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index b7da2080f3..4f0ff53847 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -145,7 +145,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic * Consumes a self-write marker. Returns true if the fs event was self-triggered. */ private consumeSelfWrite(uri: Uri): boolean { - const key = uri.toString(); + const key = this.normalizeFileUri(uri); // Check snapshot self-writes first if (this.snapshotSelfWriteUris.has(key)) { @@ -187,6 +187,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic if (liveCells.length !== newCells.length) { return true; } + return liveCells.some( (live, i) => live.kind !== newCells[i].kind || @@ -332,6 +333,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); @@ -341,13 +343,38 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic return; } - // Save to sync mtime — mark as self-write first - this.markSelfWrite(notebook.uri); + // Serialize the notebook and write canonical bytes to disk. This ensures + // the file on disk matches what VS Code's serializer would produce. + // Then save via workspace.save() to clear dirty state and update VS Code's + // internal mtime tracker. Since WE just wrote the file, its mtime is from + // our write (not the external change), avoiding the "content is newer" conflict. + const serializeTokenSource = new CancellationTokenSource(); try { - await workspace.save(notebook.uri); - } catch (error) { - this.consumeSelfWrite(notebook.uri); - throw error; + const serializedBytes = await this.serializer.serializeNotebook(newData, serializeTokenSource.token); + + // Write to disk first — this updates the file mtime to "now" + this.markSelfWrite(fileUri); + try { + await workspace.fs.writeFile(fileUri, serializedBytes); + } catch (writeError) { + this.consumeSelfWrite(fileUri); + logger.warn(`[FileChangeWatcher] Failed to write synced file: ${fileUri.path}`, writeError); + } + + // Now save — VS Code serializes (same bytes), sees the mtime is from our + // recent write (which its internal watcher has picked up), and writes + // successfully without a "content is newer" conflict. + this.markSelfWrite(fileUri); + try { + await workspace.save(notebook.uri); + } catch (saveError) { + this.consumeSelfWrite(fileUri); + logger.warn(`[FileChangeWatcher] Save after sync write failed: ${notebook.uri.path}`, saveError); + } + } catch (serializeError) { + logger.warn(`[FileChangeWatcher] Failed to serialize for sync write: ${fileUri.path}`, serializeError); + } finally { + serializeTokenSource.dispose(); } logger.info(`[FileChangeWatcher] Reloaded notebook from external change: ${notebook.uri.path}`); @@ -523,6 +550,15 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic return (metadata?.__deepnoteBlockId ?? metadata?.id) as string | undefined; } + /** + * Normalizes a URI to the underlying file path by stripping query and fragment. + * Notebook URIs include query params (e.g., ?notebook=id) but the filesystem + * watcher fires with the raw file URI — keys must match for self-write detection. + */ + private normalizeFileUri(uri: Uri): string { + return uri.with({ query: '', fragment: '' }).toString(); + } + private handleFileChange(uri: Uri): void { // Deterministic self-write check — no timers involved if (this.consumeSelfWrite(uri)) { @@ -581,7 +617,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic * Call before workspace.save() to prevent the resulting fs event from triggering a reload. */ private markSelfWrite(uri: Uri): void { - const key = uri.toString(); + const key = this.normalizeFileUri(uri); const count = this.selfWriteCounts.get(key) ?? 0; this.selfWriteCounts.set(key, count + 1); From 77c0347d76d2dfe5b22eff4c2ea757b0661d706d Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 20 Mar 2026 07:46:10 +0000 Subject: [PATCH 04/47] Improve error handling in DeepnoteFileChangeWatcher to prevent stale mtime conflicts during workspace saves --- src/notebooks/deepnote/deepnoteFileChangeWatcher.ts | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index 4f0ff53847..9b7f2a31b1 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -359,11 +359,15 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } catch (writeError) { this.consumeSelfWrite(fileUri); logger.warn(`[FileChangeWatcher] Failed to write synced file: ${fileUri.path}`, writeError); + + // Bail out — without a successful write, workspace.save() would hit + // the stale mtime and show a "content is newer" conflict dialog. + // The notebook stays dirty but cells are already up-to-date from applyEdit. + return; } - // Now save — VS Code serializes (same bytes), sees the mtime is from our - // recent write (which its internal watcher has picked up), and writes - // successfully without a "content is newer" conflict. + // Save to clear dirty state. VS Code serializes (same bytes) and sees the + // mtime from our recent write, so no "content is newer" conflict. this.markSelfWrite(fileUri); try { await workspace.save(notebook.uri); From 6dc77a0ed094954a49e97d38e2aec15d7597066d Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 20 Mar 2026 08:10:50 +0000 Subject: [PATCH 05/47] Add unit tests for DeepnoteFileChangeWatcher to handle scenarios without block IDs and implement a valid project structure for testing --- .../deepnoteFileChangeWatcher.unit.test.ts | 79 +++++++++++++++++-- 1 file changed, 72 insertions(+), 7 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index c9229820e2..ebcb61e987 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -11,6 +11,22 @@ import { IDeepnoteNotebookManager } from '../types'; import { DeepnoteFileChangeWatcher } from './deepnoteFileChangeWatcher'; import { SnapshotService } from './snapshots/snapshotService'; +const validProject = { + version: '1.0.0', + metadata: { createdAt: '2025-01-01T00:00:00Z' }, + project: { + id: 'project-1', + name: 'Test Project', + notebooks: [ + { + id: 'notebook-1', + name: 'Notebook 1', + blocks: [{ id: 'block-1', type: 'code', sortingKey: 'a0', blockGroup: '1', content: 'print("hello")' }] + } + ] + } +} as DeepnoteProject; + const waitForTimeoutMs = 5000; const waitForIntervalMs = 50; const debounceWaitMs = 800; @@ -37,6 +53,7 @@ async function waitFor( suite('DeepnoteFileChangeWatcher', () => { let watcher: DeepnoteFileChangeWatcher; let mockDisposables: IDisposableRegistry; + let mockedNotebookManager: IDeepnoteNotebookManager; let mockNotebookManager: IDeepnoteNotebookManager; let onDidChangeFile: EventEmitter; let onDidCreateFile: EventEmitter; @@ -51,7 +68,11 @@ suite('DeepnoteFileChangeWatcher', () => { saveCount = 0; mockDisposables = []; - mockNotebookManager = instance(mock()); + + mockedNotebookManager = mock(); + when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(validProject); + when(mockedNotebookManager.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1'); + mockNotebookManager = instance(mockedNotebookManager); // Set up FileSystemWatcher mock onDidChangeFile = new EventEmitter(); @@ -126,6 +147,7 @@ suite('DeepnoteFileChangeWatcher', () => { readFileCalls++; return Promise.resolve(new TextEncoder().encode(yamlContent)); }); + when(mockFs.writeFile(anything(), anything())).thenReturn(Promise.resolve()); when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); return mockFs; @@ -325,8 +347,9 @@ project: onDidChangeFile.fire(uri); await waitFor(() => saveCount >= 1); - // The save from the first reload set a self-write marker. - // Simulate the auto-save fs event being consumed (as it would in real VS Code). + // The first reload sets 2 self-write markers (writeFile + save). + // Consume them both with simulated fs events. + onDidChangeFile.fire(uri); onDidChangeFile.fire(uri); // Second real external change: use different YAML content @@ -1053,6 +1076,42 @@ project: }); test('should not apply updates when cells have no block IDs and no fallback', async () => { + const noFallbackDisposables: IDisposableRegistry = []; + const noFallbackOnDidChange = new EventEmitter(); + const noFallbackOnDidCreate = new EventEmitter(); + const fsWatcherNf = mock(); + when(fsWatcherNf.onDidChange).thenReturn(noFallbackOnDidChange.event); + when(fsWatcherNf.onDidCreate).thenReturn(noFallbackOnDidCreate.event); + when(fsWatcherNf.dispose()).thenReturn(); + when(mockedVSCodeNamespaces.workspace.createFileSystemWatcher(anything())).thenReturn( + instance(fsWatcherNf) + ); + + let nfApplyEditCount = 0; + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { + nfApplyEditCount++; + return Promise.resolve(true); + }); + + let nfReadSnapshotCount = 0; + const nfSnapshotService = mock(); + when(nfSnapshotService.isSnapshotsEnabled()).thenReturn(true); + when(nfSnapshotService.readSnapshot(anything())).thenCall(() => { + nfReadSnapshotCount++; + return Promise.resolve(snapshotOutputs); + }); + when(nfSnapshotService.onFileWritten(anything())).thenReturn({ dispose: () => {} } as Disposable); + + const nfManager = mock(); + when(nfManager.getOriginalProject(anything())).thenReturn(undefined); + + const nfWatcher = new DeepnoteFileChangeWatcher( + noFallbackDisposables, + instance(nfManager), + instance(nfSnapshotService) + ); + nfWatcher.activate(); + const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); const notebook = createMockNotebook({ uri: Uri.file('/workspace/test.deepnote'), @@ -1068,16 +1127,22 @@ project: when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - snapshotOnDidChange.fire(snapshotUri); + noFallbackOnDidChange.fire(snapshotUri); - await waitFor(() => readSnapshotCallCount >= 1); + await waitFor(() => nfReadSnapshotCount >= 1); - assert.isAtLeast(readSnapshotCallCount, 1, 'readSnapshot should be called'); + assert.isAtLeast(nfReadSnapshotCount, 1, 'readSnapshot should be called'); assert.strictEqual( - snapshotApplyEditCount, + nfApplyEditCount, 0, 'applyEdit should NOT be called when no block IDs can be resolved' ); + + for (const d of noFallbackDisposables) { + d.dispose(); + } + noFallbackOnDidChange.dispose(); + noFallbackOnDidCreate.dispose(); }); test('should fall back to replaceCells when no kernel is active', async () => { From 8898be926586ef9680227c2a21e05bafa26abf78 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 20 Mar 2026 08:19:44 +0000 Subject: [PATCH 06/47] Format code --- .../deepnote/deepnoteFileChangeWatcher.unit.test.ts | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index ebcb61e987..6bb58b4d3d 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -1132,11 +1132,7 @@ project: await waitFor(() => nfReadSnapshotCount >= 1); assert.isAtLeast(nfReadSnapshotCount, 1, 'readSnapshot should be called'); - assert.strictEqual( - nfApplyEditCount, - 0, - 'applyEdit should NOT be called when no block IDs can be resolved' - ); + assert.strictEqual(nfApplyEditCount, 0, 'applyEdit should NOT be called when no block IDs can be resolved'); for (const d of noFallbackDisposables) { d.dispose(); From c74ffa7b044b390cfcb1b2846700ebfe27be03fe Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 23 Mar 2026 09:54:45 +0000 Subject: [PATCH 07/47] Enhance error handling in DeepnoteFileChangeWatcher to check save operation result and log warnings for undefined saves --- src/notebooks/deepnote/deepnoteFileChangeWatcher.ts | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index 9b7f2a31b1..f01b0950a3 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -370,7 +370,12 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic // mtime from our recent write, so no "content is newer" conflict. this.markSelfWrite(fileUri); try { - await workspace.save(notebook.uri); + const saved = await workspace.save(notebook.uri); + if (!saved) { + this.consumeSelfWrite(fileUri); + logger.warn(`[FileChangeWatcher] Save after sync write returned undefined: ${notebook.uri.path}`); + return; + } } catch (saveError) { this.consumeSelfWrite(fileUri); logger.warn(`[FileChangeWatcher] Save after sync write failed: ${notebook.uri.path}`, saveError); From 8d061ac84cfb1f8d7dbf93181818385e1b664fac Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 23 Mar 2026 09:54:50 +0000 Subject: [PATCH 08/47] Add post-snapshot read grace period in unit tests for DeepnoteFileChangeWatcher to ensure proper handling of snapshot reads --- src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 6bb58b4d3d..72e6713ef4 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -32,6 +32,7 @@ const waitForIntervalMs = 50; const debounceWaitMs = 800; const rapidChangeIntervalMs = 100; const autoSaveGraceMs = 200; +const postSnapshotReadGraceMs = 100; /** * Polls until a condition is met or a timeout is reached. @@ -1130,6 +1131,7 @@ project: noFallbackOnDidChange.fire(snapshotUri); await waitFor(() => nfReadSnapshotCount >= 1); + await new Promise((resolve) => setTimeout(resolve, postSnapshotReadGraceMs)); assert.isAtLeast(nfReadSnapshotCount, 1, 'readSnapshot should be called'); assert.strictEqual(nfApplyEditCount, 0, 'applyEdit should NOT be called when no block IDs can be resolved'); From 5388ead064c41029071cd33c4b45fdbcd3e5a511 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 26 Mar 2026 11:14:15 +0000 Subject: [PATCH 09/47] Remove the second markSelfWrite() on the workspace.save() path. --- src/notebooks/deepnote/deepnoteFileChangeWatcher.ts | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index f01b0950a3..531b102f12 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -368,16 +368,13 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic // Save to clear dirty state. VS Code serializes (same bytes) and sees the // mtime from our recent write, so no "content is newer" conflict. - this.markSelfWrite(fileUri); try { const saved = await workspace.save(notebook.uri); if (!saved) { - this.consumeSelfWrite(fileUri); logger.warn(`[FileChangeWatcher] Save after sync write returned undefined: ${notebook.uri.path}`); return; } } catch (saveError) { - this.consumeSelfWrite(fileUri); logger.warn(`[FileChangeWatcher] Save after sync write failed: ${notebook.uri.path}`, saveError); } } catch (serializeError) { From 3e7dfc97c25e8e5de357a73f64b1c8dc08bc0203 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 27 Mar 2026 07:56:08 +0000 Subject: [PATCH 10/47] feat(deepnote): Add clearNotebookSelection method and update notebook deserialization logic - Introduced `clearNotebookSelection` method in `DeepnoteNotebookManager` to reset notebook selection for a project. - Updated `DeepnoteFileChangeWatcher` to call `clearNotebookSelection` during file change events, ensuring the active editor is prioritized during re-deserialization. - Modified `deserializeNotebook` method in `DeepnoteNotebookSerializer` to accept an optional `notebookId` parameter, preventing race conditions when multiple notebooks from the same project are open. --- .../deepnote/deepnoteFileChangeWatcher.ts | 18 +++++- .../deepnote/deepnoteNotebookManager.ts | 8 +++ .../deepnoteNotebookManager.unit.test.ts | 27 +++++++++ src/notebooks/deepnote/deepnoteSerializer.ts | 14 +++-- .../deepnote/deepnoteSerializer.unit.test.ts | 58 +++++++++++++++++++ src/notebooks/types.ts | 1 + 6 files changed, 120 insertions(+), 6 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index 531b102f12..c31774cd27 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -235,6 +235,17 @@ 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 @@ -276,10 +287,15 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic return; } + // Pass the notebook ID explicitly to avoid mutating the global selection state. + // Multiple notebooks from the same project may be open simultaneously, and + // mutating selectedNotebookByProject would cause race conditions. + const notebookNotebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; + const tokenSource = new CancellationTokenSource(); let newData; try { - newData = await this.serializer.deserializeNotebook(content, tokenSource.token); + newData = await this.serializer.deserializeNotebook(content, tokenSource.token, notebookNotebookId); } catch (error) { logger.warn(`[FileChangeWatcher] Failed to parse changed file: ${fileUri.path}`, error); return; diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.ts b/src/notebooks/deepnote/deepnoteNotebookManager.ts index 069f3a570c..1d80822cc6 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.ts @@ -14,6 +14,14 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { private readonly projectsWithInitNotebookRun = new Set(); private readonly selectedNotebookByProject = new Map(); + /** + * Clears the notebook selection for a project so that subsequent + * deserializations fall back to the active editor or open documents. + */ + clearNotebookSelection(projectId: string): void { + this.selectedNotebookByProject.delete(projectId); + } + /** * Gets the currently selected notebook ID for a project. * @param projectId Project identifier diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts index feb2842848..7c23b7e922 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts @@ -123,6 +123,33 @@ suite('DeepnoteNotebookManager', () => { }); }); + suite('clearNotebookSelection', () => { + test('should clear selection for a project', () => { + manager.selectNotebookForProject('project-123', 'notebook-456'); + manager.clearNotebookSelection('project-123'); + + const selectedNotebook = manager.getTheSelectedNotebookForAProject('project-123'); + + assert.strictEqual(selectedNotebook, undefined); + }); + + test('should not affect other projects', () => { + manager.selectNotebookForProject('project-1', 'notebook-1'); + manager.selectNotebookForProject('project-2', 'notebook-2'); + manager.clearNotebookSelection('project-1'); + + assert.strictEqual(manager.getTheSelectedNotebookForAProject('project-1'), undefined); + assert.strictEqual(manager.getTheSelectedNotebookForAProject('project-2'), 'notebook-2'); + }); + + test('should be idempotent for unknown project', () => { + assert.doesNotThrow(() => { + manager.clearNotebookSelection('unknown-project'); + manager.clearNotebookSelection('unknown-project'); + }); + }); + }); + suite('storeOriginalProject', () => { test('should store both project and current notebook ID', () => { manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 69637db668..532b1eb48e 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -67,7 +67,11 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { * @param token Cancellation token (unused) * @returns Promise resolving to notebook data */ - async deserializeNotebook(content: Uint8Array, token: CancellationToken): Promise { + async deserializeNotebook( + content: Uint8Array, + token: CancellationToken, + notebookId?: string + ): Promise { logger.debug('DeepnoteSerializer: Deserializing Deepnote notebook'); if (token?.isCancellationRequested) { @@ -90,16 +94,16 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } const projectId = deepnoteFile.project.id; - const notebookId = this.findCurrentNotebookId(projectId); + const resolvedNotebookId = notebookId ?? this.findCurrentNotebookId(projectId); - logger.debug(`DeepnoteSerializer: Project ID: ${projectId}, Selected notebook ID: ${notebookId}`); + logger.debug(`DeepnoteSerializer: Project ID: ${projectId}, Selected notebook ID: ${resolvedNotebookId}`); if (deepnoteFile.project.notebooks.length === 0) { throw new Error('Deepnote project contains no notebooks.'); } - const selectedNotebook = notebookId - ? deepnoteFile.project.notebooks.find((nb) => nb.id === notebookId) + const selectedNotebook = resolvedNotebookId + ? deepnoteFile.project.notebooks.find((nb) => nb.id === resolvedNotebookId) : this.findDefaultNotebook(deepnoteFile); if (!selectedNotebook) { diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index 53334ff5c3..0e27d79283 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -147,6 +147,34 @@ project: /no notebooks|notebooks.*must contain at least 1/i ); }); + + test('should deserialize the specified notebook when notebookId is passed', async () => { + const content = projectToYaml(mockProject); + const result = await serializer.deserializeNotebook(content, {} as any, 'notebook-2'); + + assert.strictEqual(result.metadata?.deepnoteNotebookId, 'notebook-2'); + assert.strictEqual(result.metadata?.deepnoteNotebookName, 'Second Notebook'); + assert.strictEqual(result.cells.length, 1); + assert.include(result.cells[0].value, 'Title'); + }); + + test('should ignore stored selection when explicit notebookId is provided', async () => { + manager.selectNotebookForProject('project-123', 'notebook-1'); + const content = projectToYaml(mockProject); + const result = await serializer.deserializeNotebook(content, {} as any, 'notebook-2'); + + assert.strictEqual(result.metadata?.deepnoteNotebookId, 'notebook-2'); + assert.strictEqual(result.metadata?.deepnoteNotebookName, 'Second Notebook'); + }); + + test('should fall back to findCurrentNotebookId when notebookId is undefined', async () => { + manager.selectNotebookForProject('project-123', 'notebook-1'); + const content = projectToYaml(mockProject); + const result = await serializer.deserializeNotebook(content, {} as any); + + assert.strictEqual(result.metadata?.deepnoteNotebookId, 'notebook-1'); + assert.strictEqual(result.metadata?.deepnoteNotebookName, 'First Notebook'); + }); }); suite('serializeNotebook', () => { @@ -420,6 +448,36 @@ project: manager.selectNotebookForProject('project-123', 'notebook-3'); assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-3'); }); + + test('should return undefined when selection is cleared and no active editor', () => { + manager.selectNotebookForProject('project-123', 'notebook-456'); + manager.clearNotebookSelection('project-123'); + + const result = serializer.findCurrentNotebookId('project-123'); + + assert.strictEqual(result, undefined); + }); + + test('should fall back to active editor after selection is cleared', () => { + manager.selectNotebookForProject('project-123', 'stored-wrong'); + manager.clearNotebookSelection('project-123'); + + const mockActiveNotebook = { + notebookType: 'deepnote', + metadata: { + deepnoteProjectId: 'project-123', + deepnoteNotebookId: 'active-editor-notebook' + } + }; + + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ + notebook: mockActiveNotebook + } as any); + + const result = serializer.findCurrentNotebookId('project-123'); + + assert.strictEqual(result, 'active-editor-notebook'); + }); }); suite('component integration', () => { diff --git a/src/notebooks/types.ts b/src/notebooks/types.ts index a12dacf52a..fd2fd83c44 100644 --- a/src/notebooks/types.ts +++ b/src/notebooks/types.ts @@ -37,6 +37,7 @@ export interface ProjectIntegration { export const IDeepnoteNotebookManager = Symbol('IDeepnoteNotebookManager'); export interface IDeepnoteNotebookManager { + clearNotebookSelection(projectId: string): void; getCurrentNotebookId(projectId: string): string | undefined; getOriginalProject(projectId: string): DeepnoteProject | undefined; getTheSelectedNotebookForAProject(projectId: string): string | undefined; From c160ae3af655e8d093541ed7b4df11af1241aeeb Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 27 Mar 2026 07:56:15 +0000 Subject: [PATCH 11/47] Update test --- .../deepnoteFileChangeWatcher.unit.test.ts | 262 +++++++++++++++++- 1 file changed, 257 insertions(+), 5 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 72e6713ef4..98c00e0efb 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -1,7 +1,17 @@ +import { DeepnoteFile, serializeDeepnoteFile } from '@deepnote/blocks'; import { assert } from 'chai'; import * as sinon from 'sinon'; -import { anything, instance, mock, when } from 'ts-mockito'; -import { Disposable, EventEmitter, FileSystemWatcher, NotebookCellKind, NotebookDocument, Uri } from 'vscode'; +import { anything, instance, mock, reset, resetCalls, verify, when } from 'ts-mockito'; +import { + Disposable, + EventEmitter, + FileSystemWatcher, + NotebookCellKind, + NotebookDocument, + NotebookEdit, + Uri, + WorkspaceEdit +} from 'vscode'; import type { IControllerRegistration } from '../controllers/types'; import type { IDisposableRegistry } from '../../platform/common/types'; @@ -11,7 +21,7 @@ import { IDeepnoteNotebookManager } from '../types'; import { DeepnoteFileChangeWatcher } from './deepnoteFileChangeWatcher'; import { SnapshotService } from './snapshots/snapshotService'; -const validProject = { +const validProject: DeepnoteFile = { version: '1.0.0', metadata: { createdAt: '2025-01-01T00:00:00Z' }, project: { @@ -21,11 +31,61 @@ const validProject = { { id: 'notebook-1', name: 'Notebook 1', - blocks: [{ id: 'block-1', type: 'code', sortingKey: 'a0', blockGroup: '1', content: 'print("hello")' }] + blocks: [ + { + id: 'block-1', + type: 'code', + sortingKey: 'a0', + blockGroup: '1', + content: 'print("hello")', + metadata: {} + } + ] + } + ] + } +}; + +const multiNotebookProject: DeepnoteFile = { + version: '1.0.0', + metadata: { createdAt: '2025-01-01T00:00:00Z' }, + project: { + id: 'project-1', + name: 'Multi Notebook Project', + notebooks: [ + { + id: 'notebook-1', + name: 'Notebook 1', + blocks: [ + { + id: 'block-nb1', + type: 'code', + sortingKey: 'a0', + blockGroup: '1', + content: 'print("nb1-new")', + metadata: {} + } + ] + }, + { + id: 'notebook-2', + name: 'Notebook 2', + blocks: [ + { + id: 'block-nb2', + type: 'code', + sortingKey: 'a0', + blockGroup: '1', + content: 'print("nb2-new")', + metadata: {} + } + ] } ] } -} as DeepnoteProject; +}; + +const multiNotebookYaml = serializeDeepnoteFile(multiNotebookProject); const waitForTimeoutMs = 5000; const waitForIntervalMs = 50; @@ -73,6 +133,7 @@ suite('DeepnoteFileChangeWatcher', () => { mockedNotebookManager = mock(); when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(validProject); when(mockedNotebookManager.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1'); + when(mockedNotebookManager.clearNotebookSelection(anything())).thenReturn(); mockNotebookManager = instance(mockedNotebookManager); // Set up FileSystemWatcher mock @@ -1200,4 +1261,195 @@ project: fbOnDidCreate.dispose(); }); }); + + suite('multi-notebook file sync', () => { + let workspaceSetCaptures: { uriKey: string; cellSourceJoined: string }[] = []; + let workspaceEditSetStub: sinon.SinonStub | undefined; + + setup(() => { + reset(mockedNotebookManager); + when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); + when(mockedNotebookManager.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1'); + when(mockedNotebookManager.clearNotebookSelection(anything())).thenReturn(); + resetCalls(mockedNotebookManager); + workspaceSetCaptures = []; + workspaceEditSetStub = sinon.stub(WorkspaceEdit.prototype, 'set').callsFake((uri: Uri, edits: unknown) => { + if (!Array.isArray(edits) || edits.length === 0) { + return; + } + + const firstEdit = edits[0] as NotebookEdit; + if (firstEdit?.newCells && firstEdit.newCells.length > 0) { + workspaceSetCaptures.push({ + uriKey: uri.toString(), + cellSourceJoined: firstEdit.newCells.map((c) => c.value).join('\n') + }); + } + }); + }); + + teardown(() => { + workspaceEditSetStub?.restore(); + workspaceEditSetStub = undefined; + }); + + test('should reload each notebook with its own content when multiple notebooks are open', async () => { + const basePath = Uri.file('/workspace/multi.deepnote'); + const uriNb1 = basePath.with({ query: 'view=1' }); + const uriNb2 = basePath.with({ query: 'view=2' }); + + const notebook1 = createMockNotebook({ + uri: uriNb1, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-1' + }, + cells: [ + { + metadata: { id: 'block-nb1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb1-old")', languageId: 'python' } + } + ] + }); + + const notebook2 = createMockNotebook({ + uri: uriNb2, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-2' + }, + cells: [ + { + metadata: { id: 'block-nb2', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb2-old")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); + setupMockFs(multiNotebookYaml); + + onDidChangeFile.fire(basePath); + + await waitFor(() => applyEditCount >= 2); + + assert.strictEqual(applyEditCount, 2, 'applyEdit should run once per open notebook'); + assert.strictEqual(workspaceSetCaptures.length, 2, 'each notebook should get a replaceCells edit'); + + const byUri = new Map(workspaceSetCaptures.map((c) => [c.uriKey, c.cellSourceJoined])); + + assert.include(byUri.get(uriNb1.toString()) ?? '', 'nb1-new'); + assert.notInclude(byUri.get(uriNb1.toString()) ?? '', 'nb2-new'); + assert.include(byUri.get(uriNb2.toString()) ?? '', 'nb2-new'); + assert.notInclude(byUri.get(uriNb2.toString()) ?? '', 'nb1-new'); + }); + + test('should 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' }); + + const notebook1 = createMockNotebook({ + uri: uriNb1, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-1' + }, + cells: [ + { + metadata: { id: 'block-nb1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb1-old")' } + } + ] + }); + + const notebook2 = createMockNotebook({ + uri: uriNb2, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-2' + }, + cells: [ + { + metadata: { id: 'block-nb2' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb2-old")' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); + setupMockFs(multiNotebookYaml); + + onDidChangeFile.fire(basePath); + + await waitFor(() => applyEditCount >= 2); + + verify(mockedNotebookManager.clearNotebookSelection('project-1')).once(); + }); + + test('should not corrupt other notebooks when one notebook triggers a file change', async () => { + const basePath = Uri.file('/workspace/multi.deepnote'); + const uriNb1 = basePath.with({ query: 'n=1' }); + const uriNb2 = basePath.with({ query: 'n=2' }); + + const notebook1 = createMockNotebook({ + uri: uriNb1, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-1' + }, + cells: [ + { + metadata: { id: 'block-nb1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb1-old")' } + } + ] + }); + + const notebook2 = createMockNotebook({ + uri: uriNb2, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-2' + }, + cells: [ + { + metadata: { id: 'block-nb2' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb2-old")' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); + setupMockFs(multiNotebookYaml); + + onDidChangeFile.fire(basePath); + + await waitFor(() => applyEditCount >= 2); + + const nb1Cells = workspaceSetCaptures.find((c) => c.uriKey === uriNb1.toString())?.cellSourceJoined; + const nb2Cells = workspaceSetCaptures.find((c) => c.uriKey === uriNb2.toString())?.cellSourceJoined; + + assert.isDefined(nb1Cells); + assert.isDefined(nb2Cells); + assert.notStrictEqual(nb1Cells, nb2Cells, 'each open notebook must receive distinct deserialized content'); + + assert.include(nb1Cells!, 'nb1-new'); + assert.include(nb2Cells!, 'nb2-new'); + assert.notInclude(nb1Cells!, 'nb2-new', 'notebook-1 must not receive notebook-2 block content'); + assert.notInclude(nb2Cells!, 'nb1-new', 'notebook-2 must not receive notebook-1 block content'); + }); + }); }); From 0b2504a6e1b53c5281658405c137b9dda70698a7 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 27 Mar 2026 09:52:27 +0000 Subject: [PATCH 12/47] refactor(deepnote): Improve mock child process and server output management - Enhanced `createMockChildProcess` to provide a more comprehensive mock implementation for testing. - Updated `DeepnoteServerStarter` to ensure proper disposal of existing disposables when monitoring server output, improving resource management. - Adjusted error handling in server startup to streamline diagnostics and output tracking. --- .../deepnote/deepnoteServerStarter.node.ts | 8 ++- .../deepnote/deepnoteTestHelpers.node.ts | 59 ++++++++++++++- .../deepnoteFileChangeWatcher.unit.test.ts | 34 ++++----- src/test/mocks/vsc/extHostedTypes.ts | 71 ++++++++++++++++--- 4 files changed, 143 insertions(+), 29 deletions(-) diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 112f7f0e09..6b4219cdd5 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -280,7 +280,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension throw new DeepnoteServerStartupError( interpreter.uri.fsPath, - serverInfo?.jupyterPort ?? 0, + 0, 'unknown', capturedOutput?.stdout || '', capturedOutput?.stderr || '', @@ -402,6 +402,12 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension */ private monitorServerOutput(fileKey: string, serverInfo: DeepnoteServerInfo): void { const proc = serverInfo.process; + const existing = this.disposablesByFile.get(fileKey); + if (existing) { + for (const d of existing) { + d.dispose(); + } + } const disposables: IDisposable[] = []; this.disposablesByFile.set(fileKey, disposables); diff --git a/src/kernels/deepnote/deepnoteTestHelpers.node.ts b/src/kernels/deepnote/deepnoteTestHelpers.node.ts index 524c6e0e47..42d1ae19ec 100644 --- a/src/kernels/deepnote/deepnoteTestHelpers.node.ts +++ b/src/kernels/deepnote/deepnoteTestHelpers.node.ts @@ -5,11 +5,66 @@ import type { ChildProcess } from 'node:child_process'; * Satisfies the ChildProcess interface with minimal stub values. */ export function createMockChildProcess(overrides?: Partial): ChildProcess { - return { + const mockProcess: ChildProcess = { pid: undefined, + stdio: [null, null, null, null, null], + stdin: null, stdout: null, stderr: null, exitCode: null, + killed: false, + connected: false, + signalCode: null, + spawnargs: [], + spawnfile: '', + kill: () => true, + send: () => true, + disconnect: () => true, + unref: () => true, + ref: () => true, + addListener: function () { + return this; + }, + emit: () => true, + on: function () { + return this; + }, + once: function () { + return this; + }, + removeListener: function () { + return this; + }, + removeAllListeners: function () { + return this; + }, + prependListener: function () { + return this; + }, + prependOnceListener: function () { + return this; + }, + [Symbol.dispose]: () => { + return undefined; + }, + off: function () { + return this; + }, + setMaxListeners: function () { + return this; + }, + getMaxListeners: () => 10, + listeners: function () { + return []; + }, + rawListeners: function () { + return []; + }, + eventNames: function () { + return []; + }, + listenerCount: () => 0, ...overrides - } as ChildProcess; + }; + return mockProcess; } diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 98c00e0efb..bb71908937 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -1264,7 +1264,6 @@ project: suite('multi-notebook file sync', () => { let workspaceSetCaptures: { uriKey: string; cellSourceJoined: string }[] = []; - let workspaceEditSetStub: sinon.SinonStub | undefined; setup(() => { reset(mockedNotebookManager); @@ -1273,24 +1272,27 @@ project: when(mockedNotebookManager.clearNotebookSelection(anything())).thenReturn(); resetCalls(mockedNotebookManager); workspaceSetCaptures = []; - workspaceEditSetStub = sinon.stub(WorkspaceEdit.prototype, 'set').callsFake((uri: Uri, edits: unknown) => { - if (!Array.isArray(edits) || edits.length === 0) { - return; - } + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall((wsEdit: WorkspaceEdit) => { + applyEditCount++; + + const ext = wsEdit as WorkspaceEdit & { notebookEntries(): [Uri, NotebookEdit[]][] }; - const firstEdit = edits[0] as NotebookEdit; - if (firstEdit?.newCells && firstEdit.newCells.length > 0) { - workspaceSetCaptures.push({ - uriKey: uri.toString(), - cellSourceJoined: firstEdit.newCells.map((c) => c.value).join('\n') - }); + for (const [uri, edits] of ext.notebookEntries()) { + if (!edits.length) { + continue; + } + + 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') + }); + } } - }); - }); - teardown(() => { - workspaceEditSetStub?.restore(); - workspaceEditSetStub = undefined; + return Promise.resolve(true); + }); }); test('should reload each notebook with its own content when multiple notebooks are open', async () => { diff --git a/src/test/mocks/vsc/extHostedTypes.ts b/src/test/mocks/vsc/extHostedTypes.ts index 4a58f1c077..ed93be5c2c 100644 --- a/src/test/mocks/vsc/extHostedTypes.ts +++ b/src/test/mocks/vsc/extHostedTypes.ts @@ -707,6 +707,7 @@ export namespace vscMockExtHostedTypes { private _seqPool: number = 0; private _resourceEdits: { seq: number; from: vscUri.URI; to: vscUri.URI }[] = []; + private _notebookEdits = new Map(); private _textEdits = new Map(); // createResource(uri: vscode.Uri): void { @@ -759,7 +760,20 @@ export namespace vscMockExtHostedTypes { } has(uri: vscUri.URI): boolean { - return this._textEdits.has(uri.toString()); + 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(); } set(uri: vscode.Uri, edits: NotebookEdit[]): void; @@ -779,17 +793,54 @@ export namespace vscMockExtHostedTypes { | [TextEdit | SnippetTextEdit, vscode.WorkspaceEditEntryMetadata] )[] ): void { - let data = this._textEdits.get(uri.toString()); + 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); if (!data) { data = { seq: this._seqPool++, uri, edits: [] }; - this._textEdits.set(uri.toString(), data); - } - if (!edits) { - // @ts-ignore - data.edits = undefined; - } else { - //data.edits = edits.slice(0); + this._textEdits.set(uriKey, data); } + + data.edits = edits.slice(0) as TextEdit[]; } get(uri: vscUri.URI): TextEdit[] { @@ -825,7 +876,7 @@ export namespace vscMockExtHostedTypes { } get size(): number { - return this._textEdits.size + this._resourceEdits.length; + return this._textEdits.size + this._notebookEdits.size + this._resourceEdits.length; } toJSON(): any { From 9d5c6c055d3a58717f4925691fdd3ea62e92b101 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 27 Mar 2026 17:30:03 +0000 Subject: [PATCH 13/47] 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. --- .../deepnote/deepnoteFileChangeWatcher.ts | 23 +++--- .../deepnoteFileChangeWatcher.unit.test.ts | 37 +++++----- src/test/mocks/vsc/extHostedTypes.ts | 71 +++---------------- 3 files changed, 39 insertions(+), 92 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index c31774cd27..3b0bb7a57c 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -122,6 +122,13 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } } + protected async applyNotebookEdits(uri: Uri, edits: NotebookEdit[]): Promise { + 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); diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index bb71908937..2b1eae7496 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/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; }); }); diff --git a/src/test/mocks/vsc/extHostedTypes.ts b/src/test/mocks/vsc/extHostedTypes.ts index ed93be5c2c..4a58f1c077 100644 --- a/src/test/mocks/vsc/extHostedTypes.ts +++ b/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(); private _textEdits = new Map(); // 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 { From 44bd482eb25df6a2f90bb7414a9fd79336cc1a54 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 30 Mar 2026 15:35:21 +0000 Subject: [PATCH 14/47] feat(deepnote): Enhance notebook resolution management and error handling - 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. --- .../deepnote/deepnoteServerStarter.node.ts | 6 +- .../deepnote/deepnoteExplorerView.ts | 13 +- .../deepnote/deepnoteFileChangeWatcher.ts | 11 - .../deepnoteFileChangeWatcher.unit.test.ts | 12 +- .../deepnote/deepnoteNotebookManager.ts | 65 +++- .../deepnoteNotebookManager.unit.test.ts | 40 +++ src/notebooks/deepnote/deepnoteSerializer.ts | 61 ++-- .../deepnote/deepnoteSerializer.unit.test.ts | 330 +++++++++--------- src/notebooks/types.ts | 7 +- 9 files changed, 324 insertions(+), 221 deletions(-) diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 6b4219cdd5..0ea58022cd 100644 --- 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 index 33599da2f0..6d450401ba 100644 --- 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 index 3b0bb7a57c..caa2d01fb7 100644 --- 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 index 2b1eae7496..5090c26d57 100644 --- 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(); + 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 index 1d80822cc6..11ea318b83 100644 --- 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(); private readonly originalProjects = new Map(); + private readonly pendingNotebookResolutions = new Map(); private readonly projectsWithInitNotebookRun = new Set(); private readonly selectedNotebookByProject = new Map(); /** - * 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 index 7c23b7e922..4888138eef 100644 --- 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 index 532b1eb48e..ad68f90b4f 100644 --- 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 index 0e27d79283..b562c85f1d 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -74,9 +74,8 @@ suite('DeepnoteNotebookSerializer', () => { } suite('deserializeNotebook', () => { - test('should deserialize valid project with selected notebook', async () => { - // Set up the manager to select the first notebook - manager.selectNotebookForProject('project-123', 'notebook-1'); + test('should deserialize valid project with queued notebook resolution', async () => { + manager.queueNotebookResolution('project-123', 'notebook-1'); const yamlContent = ` version: '1.0.0' @@ -168,7 +167,7 @@ project: }); test('should fall back to findCurrentNotebookId when notebookId is undefined', async () => { - manager.selectNotebookForProject('project-123', 'notebook-1'); + manager.queueNotebookResolution('project-123', 'notebook-1'); const content = projectToYaml(mockProject); const result = await serializer.deserializeNotebook(content, {} as any); @@ -233,6 +232,56 @@ project: assert.include(yamlString, 'project-123'); assert.include(yamlString, 'notebook-1'); }); + + test('should use current notebook ID instead of stale selected notebook when metadata notebook ID is missing', async () => { + manager.storeOriginalProject('project-123', mockProject, 'notebook-2'); + manager.selectNotebookForProject('project-123', 'notebook-1'); + + const mockNotebookData = { + cells: [ + { + kind: 1, // NotebookCellKind.Markup + value: '# Updated second notebook', + languageId: 'markdown', + metadata: {} + } + ], + metadata: { + deepnoteProjectId: 'project-123' + } + }; + + const result = await serializer.serializeNotebook(mockNotebookData as any, {} as any); + const serializedProject = deserializeDeepnoteFile(new TextDecoder().decode(result)); + + assert.strictEqual(serializedProject.project.notebooks[0].blocks?.[0].content, 'print("hello")'); + assert.strictEqual(serializedProject.project.notebooks[1].blocks?.[0].content, '# Updated second notebook'); + }); + + test('should queue the serialized notebook for the next resolution hint', async () => { + manager.storeOriginalProject('project-123', mockProject, 'notebook-1'); + + const mockNotebookData = { + cells: [ + { + kind: 1, // NotebookCellKind.Markup + value: '# Updated second notebook', + languageId: 'markdown', + metadata: {} + } + ], + metadata: { + deepnoteProjectId: 'project-123', + deepnoteNotebookId: 'notebook-2' + } + }; + + await serializer.serializeNotebook(mockNotebookData as any, {} as any); + + manager.updateCurrentNotebookId('project-123', 'notebook-1'); + + assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-2'); + }); }); suite('findCurrentNotebookId', () => { @@ -242,22 +291,31 @@ project: when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); }); - test('should return stored notebook ID when available', () => { - manager.selectNotebookForProject('project-123', 'notebook-456'); + test('should return queued notebook resolution when available', () => { + manager.queueNotebookResolution('project-123', 'queued-notebook'); const result = serializer.findCurrentNotebookId('project-123'); - assert.strictEqual(result, 'notebook-456'); + assert.strictEqual(result, 'queued-notebook'); + }); + + test('should consume queued notebook resolution only once', () => { + manager.queueNotebookResolution('project-123', 'queued-notebook'); + + assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'queued-notebook'); + assert.strictEqual(serializer.findCurrentNotebookId('project-123'), undefined); }); - test('should fall back to active notebook document when no stored selection', () => { - // Create a mock notebook document + test('should prioritize queued notebook resolution over current notebook and open documents', () => { + manager.queueNotebookResolution('project-123', 'queued-notebook'); + manager.updateCurrentNotebookId('project-123', 'current-notebook'); + const mockNotebookDoc = { - then: undefined, // Prevent mock from being treated as a Promise-like thenable + then: undefined, notebookType: 'deepnote', metadata: { deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'notebook-from-workspace' + deepnoteNotebookId: 'open-notebook' }, uri: {} as any, version: 1, @@ -270,213 +328,147 @@ project: save: async () => true } as NotebookDocument; - // Configure the mocked workspace.notebookDocuments (same pattern as other tests) when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([mockNotebookDoc]); const result = serializer.findCurrentNotebookId('project-123'); - assert.strictEqual(result, 'notebook-from-workspace'); + assert.strictEqual(result, 'queued-notebook'); }); - test('should return undefined for unknown project', () => { - const result = serializer.findCurrentNotebookId('unknown-project'); - - assert.strictEqual(result, undefined); - }); - - test('should prioritize stored selection over fallback', () => { - manager.selectNotebookForProject('project-123', 'stored-notebook'); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, 'stored-notebook'); - }); - - test('should handle multiple projects independently', () => { - manager.selectNotebookForProject('project-1', 'notebook-1'); - manager.selectNotebookForProject('project-2', 'notebook-2'); - - const result1 = serializer.findCurrentNotebookId('project-1'); - const result2 = serializer.findCurrentNotebookId('project-2'); - - assert.strictEqual(result1, 'notebook-1'); - assert.strictEqual(result2, 'notebook-2'); - }); - - test('should prioritize stored selection over active editor', () => { - manager.selectNotebookForProject('project-123', 'stored-notebook'); - - const mockActiveNotebook = { - notebookType: 'deepnote', - metadata: { - deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'active-editor-notebook' - } - }; - - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ - notebook: mockActiveNotebook - } as any); + test('should return current notebook ID when no pending resolution exists', () => { + manager.updateCurrentNotebookId('project-123', 'current-notebook'); const result = serializer.findCurrentNotebookId('project-123'); - assert.strictEqual(result, 'stored-notebook'); + assert.strictEqual(result, 'current-notebook'); }); - test('should return active editor notebook when no stored selection exists', () => { - const mockActiveNotebook = { + test('should return the only open notebook when current notebook ID is unavailable', () => { + const mockNotebookDoc = { + then: undefined, notebookType: 'deepnote', metadata: { deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'active-editor-notebook' - } - }; + deepnoteNotebookId: 'notebook-from-workspace' + }, + uri: {} as any, + version: 1, + isDirty: false, + isUntitled: false, + isClosed: false, + cellCount: 0, + cellAt: () => ({}) as any, + getCells: () => [], + save: async () => true + } as NotebookDocument; - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ - notebook: mockActiveNotebook - } as any); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([mockNotebookDoc]); const result = serializer.findCurrentNotebookId('project-123'); - assert.strictEqual(result, 'active-editor-notebook'); + assert.strictEqual(result, 'notebook-from-workspace'); }); - test('should ignore active editor when project ID does not match', () => { - manager.selectNotebookForProject('project-123', 'stored-notebook'); + test('should prefer current notebook ID when multiple notebooks are open for a project', () => { + manager.updateCurrentNotebookId('project-123', 'notebook-b'); - const mockActiveNotebook = { + const notebookA = { + then: undefined, notebookType: 'deepnote', - metadata: { - deepnoteProjectId: 'different-project', - deepnoteNotebookId: 'active-editor-notebook' - } - }; - - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ - notebook: mockActiveNotebook - } as any); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, 'stored-notebook'); - }); - - test('should ignore active editor when notebook type is not deepnote', () => { - manager.selectNotebookForProject('project-123', 'stored-notebook'); - - const mockActiveNotebook = { - notebookType: 'jupyter-notebook', metadata: { deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'active-editor-notebook' - } - }; - - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ - notebook: mockActiveNotebook - } as any); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, 'stored-notebook'); - }); - - test('should ignore active editor when notebook ID is missing', () => { - manager.selectNotebookForProject('project-123', 'stored-notebook'); - - const mockActiveNotebook = { + deepnoteNotebookId: 'notebook-a' + }, + uri: {} as any, + version: 1, + isDirty: false, + isUntitled: false, + isClosed: false, + cellCount: 0, + cellAt: () => ({}) as any, + getCells: () => [], + save: async () => true + } as NotebookDocument; + const notebookB = { + then: undefined, notebookType: 'deepnote', metadata: { - deepnoteProjectId: 'project-123' - } - }; + deepnoteProjectId: 'project-123', + deepnoteNotebookId: 'notebook-b' + }, + uri: {} as any, + version: 1, + isDirty: false, + isUntitled: false, + isClosed: false, + cellCount: 0, + cellAt: () => ({}) as any, + getCells: () => [], + save: async () => true + } as NotebookDocument; - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ - notebook: mockActiveNotebook - } as any); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookA, notebookB]); const result = serializer.findCurrentNotebookId('project-123'); - assert.strictEqual(result, 'stored-notebook'); + assert.strictEqual(result, 'notebook-b'); }); - test('switching notebooks: selecting a different notebook while one is open should return the new selection', () => { - manager.selectNotebookForProject('project-123', 'notebook-A'); - - const mockActiveNotebook = { + test('should return undefined when multiple notebooks are open and no stronger signal exists', () => { + const notebookA = { + then: undefined, notebookType: 'deepnote', metadata: { deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'notebook-A' - } - }; - - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ - notebook: mockActiveNotebook - } as any); - - assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-A'); - - manager.selectNotebookForProject('project-123', 'notebook-B'); - - assert.strictEqual( - serializer.findCurrentNotebookId('project-123'), - 'notebook-B', - 'Should return the newly selected notebook, not the one currently in the active editor' - ); - }); - - test('switching notebooks: rapidly switching between three notebooks should always return the latest selection', () => { - const mockActiveNotebook = { + deepnoteNotebookId: 'notebook-a' + }, + uri: {} as any, + version: 1, + isDirty: false, + isUntitled: false, + isClosed: false, + cellCount: 0, + cellAt: () => ({}) as any, + getCells: () => [], + save: async () => true + } as NotebookDocument; + const notebookB = { + then: undefined, notebookType: 'deepnote', metadata: { deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'notebook-1' - } - }; - - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ - notebook: mockActiveNotebook - } as any); + deepnoteNotebookId: 'notebook-b' + }, + uri: {} as any, + version: 1, + isDirty: false, + isUntitled: false, + isClosed: false, + cellCount: 0, + cellAt: () => ({}) as any, + getCells: () => [], + save: async () => true + } as NotebookDocument; - manager.selectNotebookForProject('project-123', 'notebook-1'); - assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-1'); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookA, notebookB]); - manager.selectNotebookForProject('project-123', 'notebook-2'); - assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-2'); + const result = serializer.findCurrentNotebookId('project-123'); - manager.selectNotebookForProject('project-123', 'notebook-3'); - assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-3'); + assert.strictEqual(result, undefined); }); - test('should return undefined when selection is cleared and no active editor', () => { - manager.selectNotebookForProject('project-123', 'notebook-456'); - manager.clearNotebookSelection('project-123'); + test('should ignore stale selected notebook state when no other resolver state exists', () => { + manager.selectNotebookForProject('project-123', 'stored-notebook'); const result = serializer.findCurrentNotebookId('project-123'); assert.strictEqual(result, undefined); }); - test('should fall back to active editor after selection is cleared', () => { - manager.selectNotebookForProject('project-123', 'stored-wrong'); - manager.clearNotebookSelection('project-123'); - - const mockActiveNotebook = { - notebookType: 'deepnote', - metadata: { - deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'active-editor-notebook' - } - }; - - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ - notebook: mockActiveNotebook - } as any); - - const result = serializer.findCurrentNotebookId('project-123'); + test('should return undefined for unknown project', () => { + const result = serializer.findCurrentNotebookId('unknown-project'); - assert.strictEqual(result, 'active-editor-notebook'); + assert.strictEqual(result, undefined); }); }); @@ -503,12 +495,14 @@ project: }); test('should handle manager state operations', () => { + assert.isFunction(manager.consumePendingNotebookResolution, 'has consumePendingNotebookResolution method'); assert.isFunction(manager.getCurrentNotebookId, 'has getCurrentNotebookId method'); assert.isFunction(manager.getOriginalProject, 'has getOriginalProject method'); assert.isFunction( manager.getTheSelectedNotebookForAProject, 'has getTheSelectedNotebookForAProject method' ); + assert.isFunction(manager.queueNotebookResolution, 'has queueNotebookResolution method'); assert.isFunction(manager.selectNotebookForProject, 'has selectNotebookForProject method'); assert.isFunction(manager.storeOriginalProject, 'has storeOriginalProject method'); }); diff --git a/src/notebooks/types.ts b/src/notebooks/types.ts index fd2fd83c44..c460019913 100644 --- a/src/notebooks/types.ts +++ b/src/notebooks/types.ts @@ -38,9 +38,13 @@ export interface ProjectIntegration { export const IDeepnoteNotebookManager = Symbol('IDeepnoteNotebookManager'); export interface IDeepnoteNotebookManager { clearNotebookSelection(projectId: string): void; + consumePendingNotebookResolution(projectId: string): string | undefined; getCurrentNotebookId(projectId: string): string | undefined; getOriginalProject(projectId: string): DeepnoteProject | undefined; getTheSelectedNotebookForAProject(projectId: string): string | undefined; + hasInitNotebookBeenRun(projectId: string): boolean; + markInitNotebookAsRun(projectId: string): void; + queueNotebookResolution(projectId: string, notebookId: string): void; selectNotebookForProject(projectId: string, notebookId: string): void; storeOriginalProject(projectId: string, project: DeepnoteProject, notebookId: string): void; updateCurrentNotebookId(projectId: string, notebookId: string): void; @@ -54,7 +58,4 @@ export interface IDeepnoteNotebookManager { * @returns `true` if the project was found and updated successfully, `false` if the project does not exist */ updateProjectIntegrations(projectId: string, integrations: ProjectIntegration[]): boolean; - - hasInitNotebookBeenRun(projectId: string): boolean; - markInitNotebookAsRun(projectId: string): void; } From 36cc1e61577c1cecadec64cf60555ebe8d5cc770 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 31 Mar 2026 07:46:49 +0000 Subject: [PATCH 15/47] refactor(deepnote): Improve notebook ID retrieval with zod validation --- src/notebooks/deepnote/deepnoteSerializer.ts | 28 +++++++++++--------- 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index ad68f90b4f..2a79a4c7af 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -2,13 +2,14 @@ import type { DeepnoteBlock, DeepnoteFile, DeepnoteSnapshot } from '@deepnote/bl import { deserializeDeepnoteFile, isExecutableBlock, serializeDeepnoteSnapshot } from '@deepnote/blocks'; import { inject, injectable, optional } from 'inversify'; import { l10n, workspace, type CancellationToken, type NotebookData, type NotebookSerializer } from 'vscode'; +import { z } from 'zod'; +import { computeHash } from '../../platform/common/crypto'; +import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; import { logger } from '../../platform/logging'; import { IDeepnoteNotebookManager } from '../types'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; -import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; import { SnapshotService } from './snapshots/snapshotService'; -import { computeHash } from '../../platform/common/crypto'; export type { DeepnoteBlock, DeepnoteFile } from '@deepnote/blocks'; export { DeepnoteNotebook, DeepnoteOutput } from '../../platform/deepnote/deepnoteTypes'; @@ -555,16 +556,19 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { 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) - ) + ...workspace.notebookDocuments.reduce((ids, doc) => { + if (doc.notebookType !== 'deepnote' || doc.metadata.deepnoteProjectId !== projectId) { + return ids; + } + + const parsed = z.string().safeParse(doc.metadata.deepnoteNotebookId); + + if (parsed.success) { + ids.add(parsed.data); + } + + return ids; + }, new Set()) ]; } From 71579cfc8e92ef83acdb7bf8cc3d370e17650c83 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 31 Mar 2026 08:28:07 +0000 Subject: [PATCH 16/47] refactor(deepnote): Improve variable naming and import organization in DeepnoteFileChangeWatcher --- src/notebooks/deepnote/deepnoteFileChangeWatcher.ts | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index caa2d01fb7..60b0d08667 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -1,3 +1,5 @@ +import type { DeepnoteBlock } from '@deepnote/blocks'; +import { inject, injectable, optional } from 'inversify'; import { CancellationTokenSource, NotebookCell, @@ -10,13 +12,11 @@ import { WorkspaceEdit, workspace } from 'vscode'; -import { inject, injectable, optional } from 'inversify'; -import type { DeepnoteBlock } from '@deepnote/blocks'; -import { IControllerRegistration } from '../controllers/types'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { IDisposableRegistry } from '../../platform/common/types'; import { logger } from '../../platform/logging'; +import { IControllerRegistration } from '../controllers/types'; import { IDeepnoteNotebookManager } from '../types'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; import { DeepnoteNotebookSerializer } from './deepnoteSerializer'; @@ -286,12 +286,12 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic // Pass the notebook ID explicitly to avoid mutating the global selection state. // Multiple notebooks from the same project may be open simultaneously, and // mutating selectedNotebookByProject would cause race conditions. - const notebookNotebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; + const targetNotebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; const tokenSource = new CancellationTokenSource(); let newData; try { - newData = await this.serializer.deserializeNotebook(content, tokenSource.token, notebookNotebookId); + newData = await this.serializer.deserializeNotebook(content, tokenSource.token, targetNotebookId); } catch (error) { logger.warn(`[FileChangeWatcher] Failed to parse changed file: ${fileUri.path}`, error); return; @@ -378,13 +378,16 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic // Save to clear dirty state. VS Code serializes (same bytes) and sees the // mtime from our recent write, so no "content is newer" conflict. + this.markSelfWrite(fileUri); try { const saved = await workspace.save(notebook.uri); if (!saved) { + this.consumeSelfWrite(fileUri); logger.warn(`[FileChangeWatcher] Save after sync write returned undefined: ${notebook.uri.path}`); return; } } catch (saveError) { + this.consumeSelfWrite(fileUri); logger.warn(`[FileChangeWatcher] Save after sync write failed: ${notebook.uri.path}`, saveError); } } catch (serializeError) { From 23fdccac6e6ea3a44f8f97104d1cc44b7f0280de Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 31 Mar 2026 11:55:29 +0000 Subject: [PATCH 17/47] Fix notebook deserialization race conditions --- .../deepnote/deepnoteFileChangeWatcher.ts | 9 ++++ .../deepnoteFileChangeWatcher.unit.test.ts | 7 +++ .../deepnote/deepnoteNotebookManager.ts | 14 +++++- src/notebooks/deepnote/deepnoteSerializer.ts | 43 +++++++++++++++++-- src/notebooks/types.ts | 1 + 5 files changed, 70 insertions(+), 4 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index 60b0d08667..e413200a12 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -99,6 +99,15 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic this.disposables.push(watcher.onDidCreate((uri) => this.handleFileChange(uri))); this.disposables.push({ dispose: () => this.clearAllTimers() }); + this.disposables.push( + workspace.onDidSaveNotebookDocument((notebook) => { + if (notebook.notebookType === 'deepnote') { + const fileUri = notebook.uri.with({ query: '', fragment: '' }); + this.markSelfWrite(fileUri); + } + }) + ); + if (this.snapshotService) { this.disposables.push( this.snapshotService.onFileWritten((uri) => { diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 5090c26d57..5d2d375d48 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -122,6 +122,7 @@ suite('DeepnoteFileChangeWatcher', () => { let mockNotebookManager: IDeepnoteNotebookManager; let onDidChangeFile: EventEmitter; let onDidCreateFile: EventEmitter; + let onDidSaveNotebook: EventEmitter; let readFileCalls: number; let applyEditCount: number; let saveCount: number; @@ -139,6 +140,7 @@ suite('DeepnoteFileChangeWatcher', () => { when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(validProject); when(mockedNotebookManager.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1'); when(mockedNotebookManager.queueNotebookResolution(anything(), anything())).thenReturn(); + when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); mockNotebookManager = instance(mockedNotebookManager); // Set up FileSystemWatcher mock @@ -151,6 +153,9 @@ suite('DeepnoteFileChangeWatcher', () => { when(mockedVSCodeNamespaces.workspace.createFileSystemWatcher(anything())).thenReturn(instance(fsWatcher)); + onDidSaveNotebook = new EventEmitter(); + when(mockedVSCodeNamespaces.workspace.onDidSaveNotebookDocument).thenReturn(onDidSaveNotebook.event); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { applyEditCount++; return Promise.resolve(true); @@ -171,6 +176,7 @@ suite('DeepnoteFileChangeWatcher', () => { } onDidChangeFile.dispose(); onDidCreateFile.dispose(); + onDidSaveNotebook.dispose(); }); function createMockNotebook(opts: { @@ -1276,6 +1282,7 @@ project: when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); when(mockedNotebookManager.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1'); when(mockedNotebookManager.queueNotebookResolution(anything(), anything())).thenReturn(); + when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); resetCalls(mockedNotebookManager); workspaceSetCaptures = []; sinon.stub(watcher, 'applyNotebookEdits' as any).callsFake(async (...args: unknown[]) => { diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.ts b/src/notebooks/deepnote/deepnoteNotebookManager.ts index 11ea318b83..68360c7df4 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.ts @@ -3,7 +3,7 @@ import { injectable } from 'inversify'; import { IDeepnoteNotebookManager, ProjectIntegration } from '../types'; import type { DeepnoteProject } from '../../platform/deepnote/deepnoteTypes'; -const pendingNotebookResolutionTtlMs = 2_000; +const pendingNotebookResolutionTtlMs = 60_000; interface PendingNotebookResolution { notebookId: string; @@ -120,6 +120,18 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { this.currentNotebookId.set(projectId, notebookId); } + /** + * Updates the stored project data without changing the current notebook selection. + * Used during serialization where we need to cache the updated project state + * but must not alter notebook routing for other open notebooks. + * @param projectId Project identifier + * @param project Updated project data to store + */ + updateOriginalProject(projectId: string, project: DeepnoteProject): void { + const clonedProject = structuredClone(project); + this.originalProjects.set(projectId, clonedProject); + } + /** * Updates the current notebook ID for a project. * Used when switching notebooks within the same project. diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 2a79a4c7af..a9a0c8491d 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -51,10 +51,21 @@ function cloneWithoutCircularRefs(obj: T, seen = new WeakSet()): T { * Serializer for converting between Deepnote YAML files and VS Code notebook format. * Handles reading/writing .deepnote files and manages project state persistence. */ +/** Window (ms) during which a post-serialize deserialize excludes the serialized notebook. */ +const recentSerializeTtlMs = 5_000; + @injectable() export class DeepnoteNotebookSerializer implements NotebookSerializer { private converter = new DeepnoteDataConverter(); + /** + * Tracks the most recently serialized notebook per project. + * When VS Code calls $dataToNotebook after a save, it's re-reading the + * file for a SIBLING tab (the saved notebook already has its content). + * This tracker lets findCurrentNotebookId pick a sibling instead. + */ + private readonly recentSerializations = new Map(); + constructor( @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, @inject(SnapshotService) @optional() private readonly snapshotService?: SnapshotService @@ -202,7 +213,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return pendingNotebookId; } - if (currentNotebookId && (openNotebookIds.length === 0 || openNotebookIds.includes(currentNotebookId))) { + if (currentNotebookId && openNotebookIds.length === 0) { return currentNotebookId; } @@ -210,6 +221,27 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return openNotebookIds[0]; } + // Multiple notebooks open — VS Code may be re-reading the file for a + // sibling tab after a save. Pick a sibling (any open notebook that is + // NOT the one just serialized). If no recent serialization, fall back + // to currentNotebookId for backward compatibility. + if (openNotebookIds.length > 1) { + const recent = this.recentSerializations.get(projectId); + + if (recent && Date.now() - recent.timestamp < recentSerializeTtlMs) { + this.recentSerializations.delete(projectId); + const sibling = openNotebookIds.find((id) => id !== recent.notebookId); + + if (sibling) { + return sibling; + } + } + + if (currentNotebookId && openNotebookIds.includes(currentNotebookId)) { + return currentNotebookId; + } + } + return undefined; } @@ -345,8 +377,13 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { originalProject.metadata.modifiedAt = new Date().toISOString(); } - // Store the updated project back so subsequent saves start from correct state - this.notebookManager.storeOriginalProject(projectId, originalProject, notebookId); + // Store the updated project back so subsequent saves start from correct state. + // Use updateOriginalProject (not storeOriginalProject) to avoid overwriting + // currentNotebookId — when multiple notebooks share the same file, changing + // currentNotebookId here would cause VS Code's follow-up deserialize calls + // for other open notebooks to resolve to the wrong notebook. + this.notebookManager.updateOriginalProject(projectId, originalProject); + this.recentSerializations.set(projectId, { notebookId, timestamp: Date.now() }); const openNotebookIdsAtSerialize = this.findOpenNotebookIds(projectId); if (openNotebookIdsAtSerialize.length === 0) { diff --git a/src/notebooks/types.ts b/src/notebooks/types.ts index c460019913..527a56ddee 100644 --- a/src/notebooks/types.ts +++ b/src/notebooks/types.ts @@ -48,6 +48,7 @@ export interface IDeepnoteNotebookManager { selectNotebookForProject(projectId: string, notebookId: string): void; storeOriginalProject(projectId: string, project: DeepnoteProject, notebookId: string): void; updateCurrentNotebookId(projectId: string, notebookId: string): void; + updateOriginalProject(projectId: string, project: DeepnoteProject): void; /** * Updates the integrations list in the project data. From 824f471a753ee5fe44dbd88e133ba82189ac5e24 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 31 Mar 2026 16:38:14 +0000 Subject: [PATCH 18/47] feat(deepnote): Enhance testing for DeepnoteFileChangeWatcher and DeepnoteNotebookManager - Added new unit tests for self-write detection in `DeepnoteFileChangeWatcher`, ensuring that saving a deepnote notebook suppresses subsequent file change events. - Implemented tests in `DeepnoteNotebookManager` to verify project data updates without altering the current notebook ID, ensuring data integrity during updates. - Introduced a utility function for creating notebook documents in tests, improving test setup consistency. - Expanded multi-notebook save scenarios in `DeepnoteNotebookSerializer` to validate notebook ID resolution during serialization and deserialization processes. --- .../deepnoteFileChangeWatcher.unit.test.ts | 298 ++++++++++++++---- .../deepnoteNotebookManager.unit.test.ts | 70 ++++ .../deepnote/deepnoteSerializer.unit.test.ts | 233 ++++++++++++++ 3 files changed, 547 insertions(+), 54 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 5d2d375d48..f3b03a4935 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -1,5 +1,8 @@ import { DeepnoteFile, serializeDeepnoteFile } from '@deepnote/blocks'; import { assert } from 'chai'; +import { mkdtempSync, rmSync } from 'fs'; +import { tmpdir } from 'os'; +import { join } from 'path'; import * as sinon from 'sinon'; import { anything, instance, mock, reset, resetCalls, verify, when } from 'ts-mockito'; import { @@ -116,6 +119,20 @@ async function waitFor( } suite('DeepnoteFileChangeWatcher', () => { + let testFixturesDir: string; + + suiteSetup(() => { + testFixturesDir = mkdtempSync(join(tmpdir(), 'deepnote-fcw-')); + }); + + suiteTeardown(() => { + rmSync(testFixturesDir, { recursive: true, force: true }); + }); + + function testFileUri(...pathSegments: string[]): Uri { + return Uri.file(join(testFixturesDir, ...pathSegments)); + } + let watcher: DeepnoteFileChangeWatcher; let mockDisposables: IDisposableRegistry; let mockedNotebookManager: IDeepnoteNotebookManager; @@ -160,9 +177,9 @@ suite('DeepnoteFileChangeWatcher', () => { applyEditCount++; return Promise.resolve(true); }); - when(mockedVSCodeNamespaces.workspace.save(anything())).thenCall(() => { + when(mockedVSCodeNamespaces.workspace.save(anything())).thenCall((uri: Uri) => { saveCount++; - return Promise.resolve(Uri.file('/workspace/test.deepnote')); + return Promise.resolve(uri); }); watcher = new DeepnoteFileChangeWatcher(mockDisposables, mockNotebookManager); @@ -244,8 +261,78 @@ project: content: print("hello") `; + suite('save-triggered self-write detection', () => { + test('saving a deepnote notebook should suppress the next FS change event', async () => { + const uri = testFileUri('self-write.deepnote'); + const notebook = createMockNotebook({ + notebookType: 'deepnote', + uri + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(validYaml); + + onDidSaveNotebook.fire(notebook); + onDidChangeFile.fire(uri); + + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); + + assert.strictEqual(applyEditCount, 0, 'FS change after deepnote save should be treated as self-write'); + }); + + test('saving a non-deepnote notebook should not suppress FS change events', async function () { + this.timeout(8000); + const fileUri = testFileUri('jupyter-save.deepnote'); + const jupyterNotebook = createMockNotebook({ + cellCount: 0, + notebookType: 'jupyter-notebook', + uri: fileUri + }); + const deepnoteNotebook = createMockNotebook({ + cellCount: 0, + uri: fileUri.with({ query: 'view=1' }) + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([jupyterNotebook, deepnoteNotebook]); + setupMockFs(validYaml); + + onDidSaveNotebook.fire(jupyterNotebook); + onDidChangeFile.fire(fileUri); + + await waitFor(() => applyEditCount >= 1); + + assert.isAtLeast(applyEditCount, 1); + }); + + test('self-write is consumed exactly once', async function () { + this.timeout(8000); + const uri = testFileUri('self-write-once.deepnote'); + const notebook = createMockNotebook({ + notebookType: 'deepnote', + uri, + cellCount: 0 + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(validYaml); + + onDidSaveNotebook.fire(notebook); + onDidChangeFile.fire(uri); + + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); + + assert.strictEqual(applyEditCount, 0, 'first FS event after save should be skipped'); + + onDidChangeFile.fire(uri); + + await waitFor(() => applyEditCount >= 1); + + assert.isAtLeast(applyEditCount, 1); + }); + }); + test('should skip reload when content matches notebook cells', async () => { - const uri = Uri.file('/workspace/test.deepnote'); + const uri = testFileUri('test.deepnote'); // Create a notebook whose cell content already matches validYaml const notebook = createMockNotebook({ uri, @@ -273,7 +360,7 @@ project: }); test('should reload on external change', async () => { - const uri = Uri.file('/workspace/test.deepnote'); + const uri = testFileUri('test.deepnote'); const notebook = createMockNotebook({ uri, cellCount: 0 }); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); @@ -290,7 +377,7 @@ project: }); test('should skip snapshot files when no SnapshotService', async () => { - const snapshotUri = Uri.file('/workspace/snapshots/project_abc_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'project_abc_latest.snapshot.deepnote'); setupMockFs(validYaml); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); @@ -305,7 +392,7 @@ project: }); test('should reload dirty notebooks', async () => { - const uri = Uri.file('/workspace/test.deepnote'); + const uri = testFileUri('test.deepnote'); const notebook = createMockNotebook({ uri, isDirty: true }); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); @@ -323,7 +410,7 @@ project: }); test('should debounce rapid changes', async () => { - const uri = Uri.file('/workspace/test.deepnote'); + const uri = testFileUri('test.deepnote'); const notebook = createMockNotebook({ uri, cellCount: 0 }); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); @@ -344,7 +431,7 @@ project: }); test('should handle parse errors gracefully', async () => { - const uri = Uri.file('/workspace/test.deepnote'); + const uri = testFileUri('test.deepnote'); const notebook = createMockNotebook({ uri, cellCount: 0 }); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); @@ -360,7 +447,7 @@ project: }); test('should preserve live cell outputs during reload', async () => { - const uri = Uri.file('/workspace/test.deepnote'); + const uri = testFileUri('test.deepnote'); const fakeOutput = { items: [{ mime: 'text/plain', data: new Uint8Array([72]) }] }; const notebook = createMockNotebook({ uri, @@ -384,7 +471,7 @@ project: }); test('should reload dirty notebooks and preserve outputs', async () => { - const uri = Uri.file('/workspace/test.deepnote'); + const uri = testFileUri('test.deepnote'); const fakeOutput = { items: [{ mime: 'text/plain', data: new Uint8Array([72]) }] }; const notebook = createMockNotebook({ uri, @@ -410,7 +497,7 @@ project: test('should not suppress real changes after auto-save', async function () { this.timeout(5000); - const uri = Uri.file('/workspace/test.deepnote'); + const uri = testFileUri('test.deepnote'); // First change: notebook has no cells, YAML has one cell -> different -> reload const notebook = createMockNotebook({ uri, cellCount: 0, cells: [] }); @@ -451,7 +538,7 @@ project: }); test('should use atomic edit (single applyEdit for replaceCells + metadata)', async () => { - const uri = Uri.file('/workspace/test.deepnote'); + const uri = testFileUri('test.deepnote'); const notebook = createMockNotebook({ uri, cellCount: 0 }); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); @@ -466,7 +553,7 @@ project: }); test('should skip auto-save-triggered changes via content comparison', async () => { - const uri = Uri.file('/workspace/test.deepnote'); + const uri = testFileUri('test.deepnote'); // Notebook already has the same source as validYaml but with outputs const fakeOutput = { items: [{ mime: 'text/plain', data: new Uint8Array([72]) }] }; const notebook = createMockNotebook({ @@ -554,9 +641,9 @@ project: snapshotApplyEditCount++; return Promise.resolve(true); }); - when(mockedVSCodeNamespaces.workspace.save(anything())).thenCall(() => { + when(mockedVSCodeNamespaces.workspace.save(anything())).thenCall((uri: Uri) => { snapshotSaveCount++; - return Promise.resolve(Uri.file('/workspace/test.deepnote')); + return Promise.resolve(uri); }); snapshotWatcher = new DeepnoteFileChangeWatcher( @@ -576,9 +663,9 @@ project: }); test('should update outputs when snapshot file changes', async () => { - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: { id: 'block-1', type: 'code' }, @@ -605,7 +692,7 @@ project: const noSnapshotWatcher = new DeepnoteFileChangeWatcher(noSnapshotDisposables, mockNotebookManager); noSnapshotWatcher.activate(); - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); @@ -622,9 +709,9 @@ project: }); test('should skip self-triggered snapshot writes via onFileWritten', async () => { - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [{ metadata: { id: 'block-1', type: 'code' }, outputs: [] }] }); @@ -645,7 +732,7 @@ project: test('should skip when snapshots are disabled', async () => { when(mockSnapshotService.isSnapshotsEnabled()).thenReturn(false); - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); snapshotOnDidChange.fire(snapshotUri); @@ -655,12 +742,10 @@ project: }); test('should debounce rapid snapshot changes for same project', async () => { - const snapshotUri1 = Uri.file( - '/workspace/snapshots/my-project_project-1_2025-01-15T10-31-48.snapshot.deepnote' - ); - const snapshotUri2 = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri1 = testFileUri('snapshots', 'my-project_project-1_2025-01-15T10-31-48.snapshot.deepnote'); + const snapshotUri2 = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: { id: 'block-1', type: 'code' }, @@ -683,9 +768,9 @@ project: }); test('should handle onDidCreate for new snapshot files', async () => { - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: { id: 'block-1', type: 'code' }, @@ -707,9 +792,9 @@ project: }); test('should skip update when snapshot outputs match live state', async () => { - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: { id: 'block-1', type: 'code' }, @@ -734,7 +819,7 @@ project: data: new TextEncoder().encode('Hello World') }; const notebookWithOutputs = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: { id: 'block-1', type: 'code' }, @@ -753,10 +838,10 @@ project: }); test('should update outputs when content changed but count is the same', async () => { - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); const existingOutput = { items: [{ mime: 'text/plain', data: new Uint8Array([72]) }] }; const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: { id: 'block-1', type: 'code' }, @@ -778,8 +863,8 @@ project: }); test('should skip main-file reload after snapshot update via self-write tracking', async () => { - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); - const notebookUri = Uri.file('/workspace/test.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); + const notebookUri = testFileUri('test.deepnote'); const notebook = createMockNotebook({ uri: notebookUri, cells: [ @@ -822,9 +907,9 @@ project: }); test('should use two-phase edit for snapshot updates (replaceCells + metadata restore)', async () => { - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: { id: 'block-1', type: 'code' }, @@ -850,9 +935,9 @@ project: }); test('should call workspace.save after snapshot fallback output update', async () => { - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: { id: 'block-1', type: 'code' }, @@ -874,10 +959,10 @@ project: }); test('should preserve outputs for cells not covered by snapshot', async () => { - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); const existingOutput = { items: [{ mime: 'text/plain', data: new Uint8Array([72]) }] }; const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: { id: 'block-1', type: 'code' }, @@ -949,10 +1034,10 @@ project: ); fallbackWatcher.activate(); - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); // Cell has NO id in metadata — simulates VS Code losing metadata after replaceCells const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), metadata: { deepnoteProjectId: 'project-1', deepnoteNotebookId: 'notebook-1' }, cells: [ { @@ -996,7 +1081,7 @@ project: }); test('should only update cells whose outputs changed (per-cell updates)', async () => { - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); // Two cells: block-1 has no outputs (will get updated), block-2 already has matching outputs const outputItem = { @@ -1004,7 +1089,7 @@ project: data: new TextEncoder().encode('Existing output') }; const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: { id: 'block-1', type: 'code' }, @@ -1116,9 +1201,9 @@ project: ); execWatcher.activate(); - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: { id: 'block-1', type: 'code' }, @@ -1185,9 +1270,9 @@ project: ); nfWatcher.activate(); - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: {}, @@ -1244,9 +1329,9 @@ project: ); fbWatcher.activate(); - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); const notebook = createMockNotebook({ - uri: Uri.file('/workspace/test.deepnote'), + uri: testFileUri('test.deepnote'), cells: [ { metadata: { id: 'block-1', type: 'code' }, @@ -1304,7 +1389,7 @@ project: }); test('should reload each notebook with its own content when multiple notebooks are open', async () => { - const basePath = Uri.file('/workspace/multi.deepnote'); + const basePath = testFileUri('multi.deepnote'); const uriNb1 = basePath.with({ query: 'view=1' }); const uriNb2 = basePath.with({ query: 'view=2' }); @@ -1359,7 +1444,7 @@ project: }); test('should not clear notebook selection before processing file change', async () => { - const basePath = Uri.file('/workspace/multi.deepnote'); + const basePath = testFileUri('multi.deepnote'); const uriNb1 = basePath.with({ query: 'a=1' }); const uriNb2 = basePath.with({ query: 'b=2' }); @@ -1406,7 +1491,7 @@ project: }); test('should not corrupt other notebooks when one notebook triggers a file change', async () => { - const basePath = Uri.file('/workspace/multi.deepnote'); + const basePath = testFileUri('multi.deepnote'); const uriNb1 = basePath.with({ query: 'n=1' }); const uriNb2 = basePath.with({ query: 'n=2' }); @@ -1461,5 +1546,110 @@ project: assert.notInclude(nb1Cells!, 'nb2-new', 'notebook-1 must not receive notebook-2 block content'); assert.notInclude(nb2Cells!, 'nb1-new', 'notebook-2 must not receive notebook-1 block content'); }); + + test('external edit to disk should update each open notebook and not be suppressed as self-write', async () => { + const basePath = testFileUri('multi-external.deepnote'); + const uriNb1 = basePath.with({ query: 'view=1' }); + const uriNb2 = basePath.with({ query: 'view=2' }); + + const notebook1 = createMockNotebook({ + uri: uriNb1, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-1' + }, + cells: [ + { + metadata: { id: 'block-nb1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb1-old")', languageId: 'python' } + } + ] + }); + + const notebook2 = createMockNotebook({ + uri: uriNb2, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-2' + }, + cells: [ + { + metadata: { id: 'block-nb2', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb2-old")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); + setupMockFs(multiNotebookYaml); + + onDidChangeFile.fire(basePath); + + await waitFor(() => applyEditCount >= 2); + + const byUri = new Map(workspaceSetCaptures.map((c) => [c.uriKey, c.cellSourceJoined])); + + assert.include(byUri.get(uriNb1.toString()) ?? '', 'nb1-new'); + assert.include(byUri.get(uriNb2.toString()) ?? '', 'nb2-new'); + }); + + test('external edit after a user save should still be processed', async function () { + this.timeout(8000); + const basePath = testFileUri('multi-save-then-external.deepnote'); + const uriNb1 = basePath.with({ query: 'view=1' }); + const uriNb2 = basePath.with({ query: 'view=2' }); + + const notebook1 = createMockNotebook({ + uri: uriNb1, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-1' + }, + cells: [ + { + metadata: { id: 'block-nb1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb1-old")' } + } + ] + }); + + const notebook2 = createMockNotebook({ + uri: uriNb2, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-2' + }, + cells: [ + { + metadata: { id: 'block-nb2' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb2-old")' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); + setupMockFs(multiNotebookYaml); + + onDidSaveNotebook.fire(notebook1); + onDidChangeFile.fire(basePath); + + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); + + assert.strictEqual(applyEditCount, 0, 'first FS event after save should be suppressed as self-write'); + + onDidChangeFile.fire(basePath); + + await waitFor(() => applyEditCount >= 1); + + assert.isAtLeast(applyEditCount, 1); + }); }); }); diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts index 4888138eef..b4276c605e 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts @@ -221,6 +221,76 @@ suite('DeepnoteNotebookManager', () => { }); }); + suite('updateOriginalProject', () => { + test('should update project data without changing currentNotebookId', () => { + const updatedProject: DeepnoteProject = { + ...mockProject, + project: { + ...mockProject.project, + name: 'Updated Name Only' + } + }; + + manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); + manager.updateOriginalProject('project-123', updatedProject); + + const storedProject = manager.getOriginalProject('project-123'); + const currentNotebookId = manager.getCurrentNotebookId('project-123'); + + assert.deepStrictEqual(storedProject, updatedProject); + assert.strictEqual(currentNotebookId, 'notebook-456'); + }); + + test('should deep-clone project data so mutations to input do not affect stored state', () => { + const updatedProject: DeepnoteProject = { + ...mockProject, + project: { + ...mockProject.project, + name: 'Before Mutation' + } + }; + + manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); + manager.updateOriginalProject('project-123', updatedProject); + + updatedProject.project.name = 'After Mutation'; + + const storedProject = manager.getOriginalProject('project-123'); + + assert.strictEqual(storedProject?.project.name, 'Before Mutation'); + }); + + test('should overwrite existing project data while preserving currentNotebookId', () => { + const firstUpdate: DeepnoteProject = { + ...mockProject, + project: { ...mockProject.project, name: 'First Update' } + }; + const secondUpdate: DeepnoteProject = { + ...mockProject, + project: { ...mockProject.project, name: 'Second Update' } + }; + + manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); + manager.updateOriginalProject('project-123', firstUpdate); + manager.updateOriginalProject('project-123', secondUpdate); + + assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-456'); + assert.strictEqual(manager.getOriginalProject('project-123')?.project.name, 'Second Update'); + }); + + test('should store project when no currentNotebookId has been set', () => { + const projectOnly: DeepnoteProject = { + ...mockProject, + project: { ...mockProject.project, name: 'No Notebook Id Yet' } + }; + + manager.updateOriginalProject('project-123', projectOnly); + + assert.strictEqual(manager.getCurrentNotebookId('project-123'), undefined); + assert.deepStrictEqual(manager.getOriginalProject('project-123'), projectOnly); + }); + }); + suite('updateCurrentNotebookId', () => { test('should update notebook ID for existing project', () => { manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index b562c85f1d..baa5e49591 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -73,6 +73,26 @@ suite('DeepnoteNotebookSerializer', () => { return new TextEncoder().encode(yamlString); } + function createNotebookDoc(notebookId: string, projectId: string = 'project-123'): NotebookDocument { + return { + cellAt: () => ({}) as any, + cellCount: 0, + getCells: () => [], + isClosed: false, + isDirty: false, + isUntitled: false, + metadata: { + deepnoteNotebookId: notebookId, + deepnoteProjectId: projectId + }, + notebookType: 'deepnote', + save: async () => true, + then: undefined, + uri: {} as any, + version: 1 + } as NotebookDocument; + } + suite('deserializeNotebook', () => { test('should deserialize valid project with queued notebook resolution', async () => { manager.queueNotebookResolution('project-123', 'notebook-1'); @@ -282,6 +302,125 @@ project: assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-2'); }); + + suite('multi-notebook save scenarios', () => { + teardown(() => { + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(undefined); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); + }); + + test('serializing notebook-1 should not change currentNotebookId', async () => { + manager.queueNotebookResolution('project-123', 'notebook-2'); + await serializer.deserializeNotebook(projectToYaml(mockProject), {} as any); + + assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-2'); + + const mockNotebookData = { + cells: [ + { + kind: 2, + languageId: 'python', + metadata: {}, + value: 'print("edited nb1")' + } + ], + metadata: { + deepnoteNotebookId: 'notebook-1', + deepnoteProjectId: 'project-123' + } + }; + + await serializer.serializeNotebook(mockNotebookData as any, {} as any); + + assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-2'); + }); + + test('serializing notebook-1 then deserializing without explicit ID should resolve to sibling notebook-2', async () => { + const yamlBytes = projectToYaml(mockProject); + + manager.queueNotebookResolution('project-123', 'notebook-1'); + await serializer.deserializeNotebook(yamlBytes, {} as any); + manager.queueNotebookResolution('project-123', 'notebook-2'); + await serializer.deserializeNotebook(yamlBytes, {} as any); + + const notebookDoc1 = createNotebookDoc('notebook-1'); + const notebookDoc2 = createNotebookDoc('notebook-2'); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookDoc1, notebookDoc2]); + + const mockNotebookData = { + cells: [ + { + kind: 2, + languageId: 'python', + metadata: {}, + value: 'print("saved from nb1")' + } + ], + metadata: { + deepnoteNotebookId: 'notebook-1', + deepnoteProjectId: 'project-123' + } + }; + + const saved = await serializer.serializeNotebook(mockNotebookData as any, {} as any); + const afterSave = await serializer.deserializeNotebook(saved, {} as any); + + assert.strictEqual(afterSave.metadata?.deepnoteNotebookId, 'notebook-2'); + }); + + test('sequential saves: contextless deserialize returns the sibling after each serialize', async () => { + const yamlBytes = projectToYaml(mockProject); + + manager.queueNotebookResolution('project-123', 'notebook-1'); + await serializer.deserializeNotebook(yamlBytes, {} as any); + manager.queueNotebookResolution('project-123', 'notebook-2'); + await serializer.deserializeNotebook(yamlBytes, {} as any); + + const notebookDoc1 = createNotebookDoc('notebook-1'); + const notebookDoc2 = createNotebookDoc('notebook-2'); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookDoc1, notebookDoc2]); + + const dataNb1 = { + cells: [ + { + kind: 2, + languageId: 'python', + metadata: {}, + value: 'print("round-a")' + } + ], + metadata: { + deepnoteNotebookId: 'notebook-1', + deepnoteProjectId: 'project-123' + } + }; + let bytes = await serializer.serializeNotebook(dataNb1 as any, {} as any); + let meta = await serializer.deserializeNotebook(bytes, {} as any); + + assert.strictEqual(meta.metadata?.deepnoteNotebookId, 'notebook-2'); + + const dataNb2 = { + cells: [ + { + kind: 1, + languageId: 'markdown', + metadata: {}, + value: '# round-b' + } + ], + metadata: { + deepnoteNotebookId: 'notebook-2', + deepnoteProjectId: 'project-123' + } + }; + bytes = await serializer.serializeNotebook(dataNb2 as any, {} as any); + meta = await serializer.deserializeNotebook(bytes, {} as any); + + assert.strictEqual(meta.metadata?.deepnoteNotebookId, 'notebook-1'); + }); + }); }); suite('findCurrentNotebookId', () => { @@ -465,6 +604,100 @@ project: assert.strictEqual(result, undefined); }); + test('after serializing notebook-1 with both notebooks open should return sibling notebook-2', async () => { + manager.storeOriginalProject('project-123', mockProject, 'notebook-1'); + + const notebookDoc1 = createNotebookDoc('notebook-1'); + const notebookDoc2 = createNotebookDoc('notebook-2'); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookDoc1, notebookDoc2]); + + const mockNotebookData = { + cells: [ + { + kind: 2, + languageId: 'python', + metadata: {}, + value: 'print("findOpen")' + } + ], + metadata: { + deepnoteNotebookId: 'notebook-1', + deepnoteProjectId: 'project-123' + } + }; + + await serializer.serializeNotebook(mockNotebookData as any, {} as any); + + const result = serializer.findCurrentNotebookId('project-123'); + + assert.strictEqual(result, 'notebook-2'); + }); + + test('recent serialization hint is one-shot', async () => { + manager.storeOriginalProject('project-123', mockProject, 'notebook-1'); + + const notebookDoc1 = createNotebookDoc('notebook-1'); + const notebookDoc2 = createNotebookDoc('notebook-2'); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookDoc1, notebookDoc2]); + + const mockNotebookData = { + cells: [ + { + kind: 2, + languageId: 'python', + metadata: {}, + value: 'print("one-shot")' + } + ], + metadata: { + deepnoteNotebookId: 'notebook-1', + deepnoteProjectId: 'project-123' + } + }; + + await serializer.serializeNotebook(mockNotebookData as any, {} as any); + + assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-2'); + assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-1'); + }); + + test('recent serialization hint expires after TTL', async () => { + manager.storeOriginalProject('project-123', mockProject, 'notebook-1'); + + const notebookDoc1 = createNotebookDoc('notebook-1'); + const notebookDoc2 = createNotebookDoc('notebook-2'); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookDoc1, notebookDoc2]); + + const mockNotebookData = { + cells: [ + { + kind: 2, + languageId: 'python', + metadata: {}, + value: 'print("ttl")' + } + ], + metadata: { + deepnoteNotebookId: 'notebook-1', + deepnoteProjectId: 'project-123' + } + }; + + await serializer.serializeNotebook(mockNotebookData as any, {} as any); + + (serializer as any).recentSerializations.set('project-123', { + notebookId: 'notebook-1', + timestamp: Date.now() - 6000 + }); + + const result = serializer.findCurrentNotebookId('project-123'); + + assert.strictEqual(result, 'notebook-1'); + }); + test('should return undefined for unknown project', () => { const result = serializer.findCurrentNotebookId('unknown-project'); From c6d21e6b45735f9cebf0545be99141c905eb7aa1 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 1 Apr 2026 07:52:31 +0000 Subject: [PATCH 19/47] feat(deepnote): Add snapshot interaction tests for DeepnoteFileChangeWatcher - Introduced new unit tests to validate snapshot changes and deserialization interactions in `DeepnoteFileChangeWatcher`. - Enhanced test setup to capture interactions with notebook edits, ensuring accurate application of changes across multi-notebook projects. - Improved organization of imports and added missing type definitions for better clarity and maintainability. --- .../deepnoteFileChangeWatcher.unit.test.ts | 560 +++++++++++++++++- 1 file changed, 558 insertions(+), 2 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index f3b03a4935..35591319b5 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -2,7 +2,6 @@ import { DeepnoteFile, serializeDeepnoteFile } from '@deepnote/blocks'; import { assert } from 'chai'; import { mkdtempSync, rmSync } from 'fs'; import { tmpdir } from 'os'; -import { join } from 'path'; import * as sinon from 'sinon'; import { anything, instance, mock, reset, resetCalls, verify, when } from 'ts-mockito'; import { @@ -14,11 +13,12 @@ import { NotebookEdit, Uri } from 'vscode'; +import { join } from '../../platform/vscode-path/path'; -import type { IControllerRegistration } from '../controllers/types'; import type { IDisposableRegistry } from '../../platform/common/types'; import type { DeepnoteOutput, DeepnoteProject } from '../../platform/deepnote/deepnoteTypes'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; +import type { IControllerRegistration } from '../controllers/types'; import { IDeepnoteNotebookManager } from '../types'; import { DeepnoteFileChangeWatcher } from './deepnoteFileChangeWatcher'; import { SnapshotService } from './snapshots/snapshotService'; @@ -101,6 +101,12 @@ interface NotebookEditCapture { cellSourceJoined: string; } +interface SnapshotInteractionCapture { + cellSourcesJoined: string; + outputPlainJoined: string; + uriKey: string; +} + /** * Polls until a condition is met or a timeout is reached. */ @@ -1356,6 +1362,556 @@ project: fbOnDidChange.dispose(); fbOnDidCreate.dispose(); }); + + suite('snapshot and deserialization interaction', () => { + let interactionCaptures: SnapshotInteractionCapture[]; + let snapshotApplyEditStub: sinon.SinonStub; + + setup(function () { + this.timeout(12_000); + interactionCaptures = []; + + reset(mockedNotebookManager); + when(mockedNotebookManager.consumePendingNotebookResolution(anything())).thenReturn(undefined); + when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(validProject); + when(mockedNotebookManager.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1'); + when(mockedNotebookManager.queueNotebookResolution(anything(), anything())).thenReturn(); + when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); + resetCalls(mockedNotebookManager); + + snapshotApplyEditStub = sinon + .stub(snapshotWatcher, 'applyNotebookEdits' as keyof DeepnoteFileChangeWatcher) + .callsFake(async function (this: DeepnoteFileChangeWatcher, ...args: unknown[]) { + const uri = args[0] as Uri; + const edits = args[1] as NotebookEdit[]; + + const replaceCellsEdit = edits.find((e) => (e as { newCells?: unknown[] }).newCells?.length); + if (replaceCellsEdit) { + const newCells = ( + replaceCellsEdit as { + newCells: Array<{ + outputs?: Array<{ items: Array<{ data?: Uint8Array }> }>; + value: string; + }>; + } + ).newCells; + const outputPlainJoined = newCells + .map((c) => { + const data = c.outputs?.[0]?.items?.[0]?.data; + + return data ? new TextDecoder().decode(data) : ''; + }) + .filter(Boolean) + .join(';'); + + interactionCaptures.push({ + uriKey: uri.toString(), + cellSourcesJoined: newCells.map((c) => c.value).join('\n'), + outputPlainJoined + }); + } + + return DeepnoteFileChangeWatcher.prototype.applyNotebookEdits.apply(this, [ + uri, + edits + ] as never); + }); + }); + + teardown(() => { + snapshotApplyEditStub.restore(); + }); + + test('snapshot change with multi-notebook project applies only matching block outputs per notebook', async () => { + when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); + + const multiOutputs = new Map([ + [ + 'block-nb1', + [ + { + output_type: 'execute_result', + data: { 'text/plain': 'OutputForNb1Only' }, + execution_count: 1 + } as DeepnoteOutput + ] + ], + [ + 'block-nb2', + [ + { + output_type: 'execute_result', + data: { 'text/plain': 'OutputForNb2Only' }, + execution_count: 1 + } as DeepnoteOutput + ] + ] + ]); + when(mockSnapshotService.readSnapshot(anything())).thenCall(() => { + readSnapshotCallCount++; + + return Promise.resolve(multiOutputs); + }); + + const basePath = testFileUri('multi-snap.deepnote'); + const uriNb1 = basePath.with({ query: 'view=1' }); + const uriNb2 = basePath.with({ query: 'view=2' }); + + const notebook1 = createMockNotebook({ + uri: uriNb1, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-1' + }, + cells: [ + { + metadata: { id: 'block-nb1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb1-old")', languageId: 'python' } + } + ] + }); + + const notebook2 = createMockNotebook({ + uri: uriNb2, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-2' + }, + cells: [ + { + metadata: { id: 'block-nb2', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb2-old")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); + + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); + snapshotOnDidChange.fire(snapshotUri); + + await waitFor(() => snapshotApplyEditCount >= 4); + + const byUri = new Map(interactionCaptures.map((c) => [c.uriKey, c])); + + assert.include(byUri.get(uriNb1.toString())?.outputPlainJoined ?? '', 'OutputForNb1Only'); + assert.notInclude(byUri.get(uriNb1.toString())?.outputPlainJoined ?? '', 'OutputForNb2Only'); + + assert.include(byUri.get(uriNb2.toString())?.outputPlainJoined ?? '', 'OutputForNb2Only'); + assert.notInclude(byUri.get(uriNb2.toString())?.outputPlainJoined ?? '', 'OutputForNb1Only'); + }); + + test('main file change after snapshot update deserializes updated cell source', async function () { + this.timeout(10_000); + const notebookUri = testFileUri('after-snap.deepnote'); + const notebook = createMockNotebook({ + uri: notebookUri, + cells: [ + { + metadata: { id: 'block-1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); + when(mockSnapshotService.readSnapshot(anything())).thenCall(() => { + readSnapshotCallCount++; + + return Promise.resolve(snapshotOutputs); + }); + + snapshotOnDidChange.fire(snapshotUri); + await waitFor(() => snapshotApplyEditCount >= 2); + + const changedYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00Z' +project: + id: project-1 + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + type: code + sortingKey: a0 + blockGroup: '1' + content: print("after-snapshot-main-sync") +`; + + const mockFs = mock(); + when(mockFs.readFile(anything())).thenCall(() => { + return Promise.resolve(new TextEncoder().encode(changedYaml)); + }); + when(mockFs.writeFile(anything(), anything())).thenReturn(Promise.resolve()); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + // Snapshot save marks self-write; first FS event consumes it without reloading. + snapshotOnDidChange.fire(notebookUri); + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); + + snapshotOnDidChange.fire(notebookUri); + + await waitFor(() => + interactionCaptures.some((c) => c.cellSourcesJoined.includes('after-snapshot-main-sync')) + ); + + const mainSyncCapture = interactionCaptures.find((c) => + c.cellSourcesJoined.includes('after-snapshot-main-sync') + ); + + assert.isDefined(mainSyncCapture); + assert.include( + mainSyncCapture!.cellSourcesJoined, + 'after-snapshot-main-sync', + 'main-file sync should deserialize new source after snapshot outputs were applied' + ); + }); + + test('snapshot save self-write is consumed once then external main-file change applies', async function () { + this.timeout(10_000); + const baseUri = testFileUri('snap-selfwrite.deepnote'); + const notebook = createMockNotebook({ + uri: baseUri, + cells: [ + { + metadata: { id: 'block-1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); + snapshotOnDidChange.fire(snapshotUri); + await waitFor(() => snapshotSaveCount >= 1); + + const editsBefore = snapshotApplyEditCount; + + snapshotOnDidChange.fire(baseUri); + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); + + assert.strictEqual( + snapshotApplyEditCount, + editsBefore, + 'first main-file FS event after snapshot save should be consumed as self-write' + ); + + const externalYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00Z' +project: + id: project-1 + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + type: code + sortingKey: a0 + blockGroup: '1' + content: print("external-after-self-write") +`; + + const mockFs = mock(); + when(mockFs.readFile(anything())).thenCall(() => { + return Promise.resolve(new TextEncoder().encode(externalYaml)); + }); + when(mockFs.writeFile(anything(), anything())).thenReturn(Promise.resolve()); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + snapshotOnDidChange.fire(baseUri); + + await waitFor(() => + interactionCaptures.some((c) => c.cellSourcesJoined.includes('external-after-self-write')) + ); + + assert.isTrue( + interactionCaptures.some((c) => c.cellSourcesJoined.includes('external-after-self-write')), + 'second main-file change should deserialize external content' + ); + }); + + test('main-file sync runs after in-flight snapshot when both are triggered close together', async function () { + this.timeout(12_000); + let releaseSnapshot!: () => void; + const snapshotGate = new Promise((resolve) => { + releaseSnapshot = resolve; + }); + + when(mockSnapshotService.readSnapshot(anything())).thenCall(() => { + readSnapshotCallCount++; + + return snapshotGate.then(() => snapshotOutputs); + }); + + const baseUri = testFileUri('coalesce.deepnote'); + const notebook = createMockNotebook({ + uri: baseUri, + cells: [ + { + metadata: { id: 'block-1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); + snapshotOnDidChange.fire(snapshotUri); + + await waitFor(() => readSnapshotCallCount >= 1); + + const coalescedYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00Z' +project: + id: project-1 + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + type: code + sortingKey: a0 + blockGroup: '1' + content: print("main-wins-after-snapshot") +`; + + const mockFs = mock(); + when(mockFs.readFile(anything())).thenCall(() => { + return Promise.resolve(new TextEncoder().encode(coalescedYaml)); + }); + when(mockFs.writeFile(anything(), anything())).thenReturn(Promise.resolve()); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + snapshotOnDidChange.fire(baseUri); + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); + + releaseSnapshot(); + + await waitFor(() => + interactionCaptures.some((c) => c.cellSourcesJoined.includes('main-wins-after-snapshot')) + ); + + assert.isTrue( + interactionCaptures.some((c) => c.cellSourcesJoined.includes('main-wins-after-snapshot')), + 'main-file sync should deserialize disk YAML after snapshot operation completes' + ); + }); + + test('multi-notebook: snapshot outputs then external YAML update keeps per-notebook sources', async function () { + this.timeout(12_000); + when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); + + const multiOutputs = new Map([ + [ + 'block-nb1', + [ + { + output_type: 'execute_result', + data: { 'text/plain': 'SnapNb1' }, + execution_count: 1 + } as DeepnoteOutput + ] + ], + [ + 'block-nb2', + [ + { + output_type: 'execute_result', + data: { 'text/plain': 'SnapNb2' }, + execution_count: 1 + } as DeepnoteOutput + ] + ] + ]); + when(mockSnapshotService.readSnapshot(anything())).thenCall(() => { + readSnapshotCallCount++; + + return Promise.resolve(multiOutputs); + }); + + const basePath = testFileUri('multi-snap-then-yaml.deepnote'); + const uriNb1 = basePath.with({ query: 'view=1' }); + const uriNb2 = basePath.with({ query: 'view=2' }); + + const notebook1 = createMockNotebook({ + uri: uriNb1, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-1' + }, + cells: [ + { + metadata: { id: 'block-nb1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb1-old")', languageId: 'python' } + } + ] + }); + + const notebook2 = createMockNotebook({ + uri: uriNb2, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-2' + }, + cells: [ + { + metadata: { id: 'block-nb2', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb2-old")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); + + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); + snapshotOnDidChange.fire(snapshotUri); + await waitFor(() => snapshotApplyEditCount >= 4); + + const round2Project = structuredClone(multiNotebookProject); + round2Project.project.notebooks[0].blocks![0].content = 'print("nb1-round2")'; + round2Project.project.notebooks[1].blocks![0].content = 'print("nb2-round2")'; + const yamlRound2 = serializeDeepnoteFile(round2Project); + + const mockFs = mock(); + when(mockFs.readFile(anything())).thenCall(() => { + return Promise.resolve(new TextEncoder().encode(yamlRound2)); + }); + when(mockFs.writeFile(anything(), anything())).thenReturn(Promise.resolve()); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + // Two snapshot saves increment self-write count to 2 for the shared base file URI. + snapshotOnDidChange.fire(basePath); + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); + snapshotOnDidChange.fire(basePath); + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); + + snapshotOnDidChange.fire(basePath); + + await waitFor(() => + interactionCaptures.some( + (c) => c.uriKey === uriNb1.toString() && c.cellSourcesJoined.includes('nb1-round2') + ) + ); + + assert.isTrue( + interactionCaptures.some( + (c) => c.uriKey === uriNb1.toString() && c.outputPlainJoined.includes('SnapNb1') + ), + 'snapshot phase should apply SnapNb1 output to notebook-1' + ); + assert.isTrue( + interactionCaptures.some( + (c) => c.uriKey === uriNb2.toString() && c.outputPlainJoined.includes('SnapNb2') + ), + 'snapshot phase should apply SnapNb2 output to notebook-2' + ); + + const nb1Main = interactionCaptures.find( + (c) => c.uriKey === uriNb1.toString() && c.cellSourcesJoined.includes('nb1-round2') + ); + const nb2Main = interactionCaptures.find( + (c) => c.uriKey === uriNb2.toString() && c.cellSourcesJoined.includes('nb2-round2') + ); + + assert.isDefined(nb1Main); + assert.include(nb1Main!.cellSourcesJoined, 'nb1-round2'); + assert.notInclude(nb1Main!.cellSourcesJoined, 'nb2-round2'); + + assert.isDefined(nb2Main); + assert.include(nb2Main!.cellSourcesJoined, 'nb2-round2'); + assert.notInclude(nb2Main!.cellSourcesJoined, 'nb1-round2'); + }); + + test('snapshot outputs for sibling notebook blocks do not leak into a single open notebook', async () => { + when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); + + const multiOutputs = new Map([ + [ + 'block-nb1', + [ + { + output_type: 'execute_result', + data: { 'text/plain': 'OnlyNb1' }, + execution_count: 1 + } as DeepnoteOutput + ] + ], + [ + 'block-nb2', + [ + { + output_type: 'execute_result', + data: { 'text/plain': 'LeakIfApplied' }, + execution_count: 1 + } as DeepnoteOutput + ] + ] + ]); + when(mockSnapshotService.readSnapshot(anything())).thenCall(() => { + readSnapshotCallCount++; + + return Promise.resolve(multiOutputs); + }); + + const uriNb1 = testFileUri('only-nb1.deepnote'); + const notebook1Only = createMockNotebook({ + uri: uriNb1, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: 'notebook-1' + }, + cells: [ + { + metadata: { id: 'block-nb1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb1-only")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1Only]); + + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); + snapshotOnDidChange.fire(snapshotUri); + + await waitFor(() => snapshotApplyEditCount >= 2); + + const cap = interactionCaptures.find((c) => c.uriKey === uriNb1.toString()); + + assert.isDefined(cap); + assert.include(cap!.outputPlainJoined, 'OnlyNb1'); + assert.notInclude(cap!.outputPlainJoined, 'LeakIfApplied'); + }); + }); }); suite('multi-notebook file sync', () => { From 5d9d6e2f6d3d0fcd15237f81367c7984da745e6b Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 1 Apr 2026 08:35:15 +0000 Subject: [PATCH 20/47] Fix cspell --- src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 35591319b5..2284e75304 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -1581,7 +1581,7 @@ project: test('snapshot save self-write is consumed once then external main-file change applies', async function () { this.timeout(10_000); - const baseUri = testFileUri('snap-selfwrite.deepnote'); + const baseUri = testFileUri('snap-self-write.deepnote'); const notebook = createMockNotebook({ uri: baseUri, cells: [ From 496261623b261481bbbae5228f3c82956a564893 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 1 Apr 2026 08:44:18 +0000 Subject: [PATCH 21/47] Fix tests --- .../deepnote/deepnoteFileChangeWatcher.ts | 2 +- .../deepnoteFileChangeWatcher.unit.test.ts | 70 +++++++++---------- .../deepnote/deepnoteSerializer.unit.test.ts | 2 + 3 files changed, 37 insertions(+), 37 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index e413200a12..940ae8a3ec 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -131,7 +131,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } } - protected async applyNotebookEdits(uri: Uri, edits: NotebookEdit[]): Promise { + public async applyNotebookEdits(uri: Uri, edits: NotebookEdit[]): Promise { const wsEdit = new WorkspaceEdit(); wsEdit.set(uri, edits); diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 2284e75304..47fe1668d8 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -1379,43 +1379,41 @@ project: when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); resetCalls(mockedNotebookManager); - snapshotApplyEditStub = sinon - .stub(snapshotWatcher, 'applyNotebookEdits' as keyof DeepnoteFileChangeWatcher) - .callsFake(async function (this: DeepnoteFileChangeWatcher, ...args: unknown[]) { - const uri = args[0] as Uri; - const edits = args[1] as NotebookEdit[]; - - const replaceCellsEdit = edits.find((e) => (e as { newCells?: unknown[] }).newCells?.length); - if (replaceCellsEdit) { - const newCells = ( - replaceCellsEdit as { - newCells: Array<{ - outputs?: Array<{ items: Array<{ data?: Uint8Array }> }>; - value: string; - }>; - } - ).newCells; - const outputPlainJoined = newCells - .map((c) => { - const data = c.outputs?.[0]?.items?.[0]?.data; - - return data ? new TextDecoder().decode(data) : ''; - }) - .filter(Boolean) - .join(';'); - - interactionCaptures.push({ - uriKey: uri.toString(), - cellSourcesJoined: newCells.map((c) => c.value).join('\n'), - outputPlainJoined - }); - } + snapshotApplyEditStub = sinon.stub(snapshotWatcher, 'applyNotebookEdits').callsFake(async function ( + this: DeepnoteFileChangeWatcher, + ...args: unknown[] + ) { + const uri = args[0] as Uri; + const edits = args[1] as NotebookEdit[]; + + const replaceCellsEdit = edits.find((e) => (e as { newCells?: unknown[] }).newCells?.length); + if (replaceCellsEdit) { + const newCells = ( + replaceCellsEdit as { + newCells: Array<{ + outputs?: Array<{ items: Array<{ data?: Uint8Array }> }>; + value: string; + }>; + } + ).newCells; + const outputPlainJoined = newCells + .map((c) => { + const data = c.outputs?.[0]?.items?.[0]?.data; + + return data ? new TextDecoder().decode(data) : ''; + }) + .filter(Boolean) + .join(';'); + + interactionCaptures.push({ + uriKey: uri.toString(), + cellSourcesJoined: newCells.map((c) => c.value).join('\n'), + outputPlainJoined + }); + } - return DeepnoteFileChangeWatcher.prototype.applyNotebookEdits.apply(this, [ - uri, - edits - ] as never); - }); + return DeepnoteFileChangeWatcher.prototype.applyNotebookEdits.apply(this, [uri, edits]); + }); }); teardown(() => { diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index baa5e49591..4cde84ebb5 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -274,6 +274,8 @@ project: const result = await serializer.serializeNotebook(mockNotebookData as any, {} as any); const serializedProject = deserializeDeepnoteFile(new TextDecoder().decode(result)); + assert.strictEqual(serializedProject.project.notebooks[0].id, 'notebook-1'); + assert.strictEqual(serializedProject.project.notebooks[1].id, 'notebook-2'); assert.strictEqual(serializedProject.project.notebooks[0].blocks?.[0].content, 'print("hello")'); assert.strictEqual(serializedProject.project.notebooks[1].blocks?.[0].content, '# Updated second notebook'); }); From 552366b28e7558d468f0a5e22bcd66f274f0c150 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 1 Apr 2026 09:43:26 +0000 Subject: [PATCH 22/47] Minor improvements --- .../deepnote/deepnoteFileChangeWatcher.unit.test.ts | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 47fe1668d8..c57c398616 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -136,7 +136,7 @@ suite('DeepnoteFileChangeWatcher', () => { }); function testFileUri(...pathSegments: string[]): Uri { - return Uri.file(join(testFixturesDir, ...pathSegments)); + return Uri.joinPath(Uri.file(testFixturesDir), ...pathSegments); } let watcher: DeepnoteFileChangeWatcher; @@ -1367,8 +1367,7 @@ project: let interactionCaptures: SnapshotInteractionCapture[]; let snapshotApplyEditStub: sinon.SinonStub; - setup(function () { - this.timeout(12_000); + setup(() => { interactionCaptures = []; reset(mockedNotebookManager); From bd4d5fd51fb497880f210e845d5f9cf98d9b65fa Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 1 Apr 2026 14:04:11 +0000 Subject: [PATCH 23/47] fix(deepnote): Enhance error handling for metadata restoration in DeepnoteFileChangeWatcher - Updated `applyNotebookEdits` calls to handle failures in restoring block IDs, logging warnings when restoration fails after execution API and replaceCells operations. - Added unit tests to verify that saving does not occur when metadata restoration fails, ensuring data integrity during notebook edits. - Improved test coverage for scenarios involving metadata restoration failures, enhancing the reliability of the DeepnoteFileChangeWatcher functionality. --- .../deepnote/deepnoteFileChangeWatcher.ts | 14 +- .../deepnoteFileChangeWatcher.unit.test.ts | 161 ++++++++++++++++++ 2 files changed, 173 insertions(+), 2 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index 940ae8a3ec..efa916b7fe 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -507,7 +507,13 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } } if (metadataEdits.length > 0) { - await this.applyNotebookEdits(notebook.uri, metadataEdits); + const metadataApplied = await this.applyNotebookEdits(notebook.uri, metadataEdits); + if (!metadataApplied) { + logger.warn( + `[FileChangeWatcher] Failed to restore block IDs via execution path: ${notebook.uri.path}` + ); + return; + } } logger.info(`[FileChangeWatcher] Updated notebook outputs via execution API: ${notebook.uri.path}`); @@ -554,7 +560,11 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic }) ); } - await this.applyNotebookEdits(notebook.uri, metadataEdits); + const metadataApplied = await this.applyNotebookEdits(notebook.uri, metadataEdits); + if (!metadataApplied) { + logger.warn(`[FileChangeWatcher] Failed to restore block IDs after replaceCells: ${notebook.uri.path}`); + return; + } // Save to sync mtime — mark as self-write first this.markSelfWrite(notebook.uri); diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index c57c398616..f14a8998f3 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -14,6 +14,7 @@ import { Uri } from 'vscode'; import { join } from '../../platform/vscode-path/path'; +import { logger } from '../../platform/logging'; import type { IDisposableRegistry } from '../../platform/common/types'; import type { DeepnoteOutput, DeepnoteProject } from '../../platform/deepnote/deepnoteTypes'; @@ -1363,6 +1364,166 @@ project: fbOnDidCreate.dispose(); }); + test('should not save when metadata restore fails after replaceCells fallback', async () => { + const mdDisposables: IDisposableRegistry = []; + const mdOnDidChange = new EventEmitter(); + const mdOnDidCreate = new EventEmitter(); + const fsWatcherMd = mock(); + when(fsWatcherMd.onDidChange).thenReturn(mdOnDidChange.event); + when(fsWatcherMd.onDidCreate).thenReturn(mdOnDidCreate.event); + when(fsWatcherMd.dispose()).thenReturn(); + when(mockedVSCodeNamespaces.workspace.createFileSystemWatcher(anything())).thenReturn( + instance(fsWatcherMd) + ); + + let mdApplyEditInvocation = 0; + let mdSaveCount = 0; + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { + mdApplyEditInvocation++; + return Promise.resolve(mdApplyEditInvocation === 1); + }); + when(mockedVSCodeNamespaces.workspace.save(anything())).thenCall((uri: Uri) => { + mdSaveCount++; + return Promise.resolve(uri); + }); + + const mockedControllerRegistration = mock(); + when(mockedControllerRegistration.getSelected(anything())).thenReturn(undefined); + + const mdWatcher = new DeepnoteFileChangeWatcher( + mdDisposables, + mockNotebookManager, + instance(mockSnapshotService), + instance(mockedControllerRegistration) + ); + mdWatcher.activate(); + + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); + const notebook = createMockNotebook({ + uri: testFileUri('metadata-fail-replace.deepnote'), + cells: [ + { + metadata: { id: 'block-1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + mdOnDidChange.fire(snapshotUri); + + await waitFor(() => mdApplyEditInvocation >= 2); + + assert.strictEqual( + mdApplyEditInvocation, + 2, + 'replaceCells and metadata restore should each invoke applyEdit once' + ); + assert.strictEqual( + mdSaveCount, + 0, + 'workspace.save must not run when metadata restore fails (would persist wrong cell IDs)' + ); + + for (const d of mdDisposables) { + d.dispose(); + } + mdOnDidChange.dispose(); + mdOnDidCreate.dispose(); + }); + + test('should warn and return when metadata restore fails after execution API with block ID fallback', async () => { + const warnStub = sinon.stub(logger, 'warn'); + + try { + const exMdDisposables: IDisposableRegistry = []; + const exMdOnDidChange = new EventEmitter(); + const exMdOnDidCreate = new EventEmitter(); + const fsWatcherExMd = mock(); + when(fsWatcherExMd.onDidChange).thenReturn(exMdOnDidChange.event); + when(fsWatcherExMd.onDidCreate).thenReturn(exMdOnDidCreate.event); + when(fsWatcherExMd.dispose()).thenReturn(); + when(mockedVSCodeNamespaces.workspace.createFileSystemWatcher(anything())).thenReturn( + instance(fsWatcherExMd) + ); + + let exMdSaveCount = 0; + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(false)); + when(mockedVSCodeNamespaces.workspace.save(anything())).thenCall((uri: Uri) => { + exMdSaveCount++; + return Promise.resolve(uri); + }); + + const mockVSCodeController = { + createNotebookCellExecution: () => ({ + start: () => {}, + replaceOutput: () => Promise.resolve(), + end: () => {} + }) + }; + const mockedControllerRegistration = mock(); + when(mockedControllerRegistration.getSelected(anything())).thenReturn({ + controller: mockVSCodeController + } as any); + + const mockedManagerEx = mock(); + when(mockedManagerEx.consumePendingNotebookResolution(anything())).thenReturn(undefined); + when(mockedManagerEx.getOriginalProject(anything())).thenReturn(validProject); + when(mockedManagerEx.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1'); + when(mockedManagerEx.queueNotebookResolution(anything(), anything())).thenReturn(); + when(mockedManagerEx.updateOriginalProject(anything(), anything())).thenReturn(); + + const exMdWatcher = new DeepnoteFileChangeWatcher( + exMdDisposables, + instance(mockedManagerEx), + instance(mockSnapshotService), + instance(mockedControllerRegistration) + ); + exMdWatcher.activate(); + + const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); + const notebook = createMockNotebook({ + uri: testFileUri('metadata-fail-exec.deepnote'), + cells: [ + { + metadata: { type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + exMdOnDidChange.fire(snapshotUri); + + await waitFor(() => warnStub.called); + + assert.include( + warnStub.firstCall.args[0] as string, + 'Failed to restore block IDs via execution path', + 'should log when metadata restore fails after execution API' + ); + assert.strictEqual( + exMdSaveCount, + 0, + 'execution API path should not save after failed metadata restore' + ); + + for (const d of exMdDisposables) { + d.dispose(); + } + exMdOnDidChange.dispose(); + exMdOnDidCreate.dispose(); + } finally { + warnStub.restore(); + } + }); + suite('snapshot and deserialization interaction', () => { let interactionCaptures: SnapshotInteractionCapture[]; let snapshotApplyEditStub: sinon.SinonStub; From ec96f534cd721402cfab7a84bb5b51a9e1f119a4 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 2 Apr 2026 13:24:43 +0000 Subject: [PATCH 24/47] feat(deepnote): Add DeepnoteNotebookInfoStatusBar for displaying notebook details in the status bar - Introduced `DeepnoteNotebookInfoStatusBar` to show the active Deepnote notebook name and provide a copy action for notebook details. - Updated service registration in both `serviceRegistry.node.ts` and `serviceRegistry.web.ts` to include the new status bar service. - Added a new command `CopyNotebookDetails` to facilitate copying notebook information to the clipboard. --- .../deepnote/deepnoteNotebookInfoStatusBar.ts | 154 ++++++++++++++++++ src/notebooks/serviceRegistry.node.ts | 5 + src/notebooks/serviceRegistry.web.ts | 5 + src/platform/common/constants.ts | 1 + 4 files changed, 165 insertions(+) create mode 100644 src/notebooks/deepnote/deepnoteNotebookInfoStatusBar.ts diff --git a/src/notebooks/deepnote/deepnoteNotebookInfoStatusBar.ts b/src/notebooks/deepnote/deepnoteNotebookInfoStatusBar.ts new file mode 100644 index 0000000000..3959cf6f91 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteNotebookInfoStatusBar.ts @@ -0,0 +1,154 @@ +import { inject, injectable } from 'inversify'; +import { + commands, + Disposable, + env, + l10n, + NotebookDocument, + StatusBarAlignment, + StatusBarItem, + window, + workspace +} from 'vscode'; + +import { DEEPNOTE_NOTEBOOK_TYPE } from '../../kernels/deepnote/types'; +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { Commands } from '../../platform/common/constants'; +import { IDisposableRegistry } from '../../platform/common/types'; + +/** + * Shows the active Deepnote notebook name in the status bar; tooltip and copy action include full debug details (metadata and URI). + */ +@injectable() +export class DeepnoteNotebookInfoStatusBar implements IExtensionSyncActivationService, Disposable { + private readonly disposables: Disposable[] = []; + + private disposed = false; + + private statusBarItem: StatusBarItem | undefined; + + constructor(@inject(IDisposableRegistry) disposableRegistry: IDisposableRegistry) { + disposableRegistry.push(this); + } + + public activate(): void { + this.statusBarItem = window.createStatusBarItem('deepnote.notebookInfo', StatusBarAlignment.Left, 99); + this.statusBarItem.name = l10n.t('Deepnote Notebook Info'); + this.statusBarItem.command = Commands.CopyNotebookDetails; + this.disposables.push(this.statusBarItem); + + this.disposables.push( + commands.registerCommand(Commands.CopyNotebookDetails, () => { + this.copyActiveNotebookDetails(); + }) + ); + + this.disposables.push(window.onDidChangeActiveNotebookEditor(() => this.updateStatusBar())); + + this.disposables.push( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook === window.activeNotebookEditor?.notebook) { + this.updateStatusBar(); + } + }) + ); + + this.updateStatusBar(); + } + + public dispose(): void { + if (this.disposed) { + return; + } + + this.disposed = true; + + while (this.disposables.length) { + const disposable = this.disposables.pop(); + + try { + disposable?.dispose(); + } catch { + // ignore + } + } + } + + private copyActiveNotebookDetails(): void { + const notebook = window.activeNotebookEditor?.notebook; + + if (!notebook) { + return; + } + + const info = this.getNotebookDebugInfo(notebook); + + if (!info) { + return; + } + + void env.clipboard.writeText(info.detailsText); + void window.showInformationMessage(l10n.t('Copied notebook details to clipboard.')); + } + + private getNotebookDebugInfo(notebook: NotebookDocument): { detailsText: string; displayName: string } | undefined { + if (notebook.notebookType !== DEEPNOTE_NOTEBOOK_TYPE) { + return undefined; + } + + const metadata = notebook.metadata as { + deepnoteNotebookId?: string; + deepnoteNotebookName?: string; + deepnoteProjectId?: string; + deepnoteProjectName?: string; + deepnoteVersion?: string; + name?: string; + }; + + const displayName = metadata.deepnoteNotebookName ?? metadata.name ?? l10n.t('Deepnote notebook'); + const uriString = notebook.uri.toString(true); + + const lines: string[] = [ + l10n.t('Notebook: {0}', displayName), + l10n.t('Notebook ID: {0}', metadata.deepnoteNotebookId ?? l10n.t('(unknown)')), + l10n.t('Project: {0}', metadata.deepnoteProjectName ?? l10n.t('(unknown)')), + l10n.t('Project ID: {0}', metadata.deepnoteProjectId ?? l10n.t('(unknown)')) + ]; + + if (metadata.deepnoteVersion !== undefined) { + lines.push(l10n.t('Deepnote version: {0}', String(metadata.deepnoteVersion))); + } + + lines.push(l10n.t('URI: {0}', uriString)); + + return { detailsText: lines.join('\n'), displayName }; + } + + private updateStatusBar(): void { + const item = this.statusBarItem; + + if (!item) { + return; + } + + const editor = window.activeNotebookEditor; + + if (!editor) { + item.hide(); + + return; + } + + const info = this.getNotebookDebugInfo(editor.notebook); + + if (!info) { + item.hide(); + + return; + } + + item.text = `$(notebook) ${info.displayName}`; + item.tooltip = [info.detailsText, l10n.t('Click to copy details')].join('\n'); + item.show(); + } +} diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index cbd8b860fe..f4ea82b469 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -41,6 +41,7 @@ import { NotebookTracebackFormatter } from './outputs/tracebackFormatter'; import { InterpreterPackageTracker } from './telemetry/interpreterPackageTracker.node'; import { INotebookEditorProvider, INotebookPythonEnvironmentService } from './types'; import { DeepnoteActivationService } from './deepnote/deepnoteActivationService'; +import { DeepnoteNotebookInfoStatusBar } from './deepnote/deepnoteNotebookInfoStatusBar'; import { DeepnoteNotebookManager } from './deepnote/deepnoteNotebookManager'; import { IDeepnoteNotebookManager } from './types'; import { IntegrationStorage } from '../platform/notebooks/deepnote/integrationStorage'; @@ -169,6 +170,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteActivationService ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + DeepnoteNotebookInfoStatusBar + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNotebookCommandListener diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index 2488ff73d7..51ea5e2b88 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/src/notebooks/serviceRegistry.web.ts @@ -36,6 +36,7 @@ import { CellOutputMimeTypeTracker } from './outputs/cellOutputMimeTypeTracker'; import { NotebookTracebackFormatter } from './outputs/tracebackFormatter'; import { INotebookEditorProvider, INotebookPythonEnvironmentService } from './types'; import { DeepnoteActivationService } from './deepnote/deepnoteActivationService'; +import { DeepnoteNotebookInfoStatusBar } from './deepnote/deepnoteNotebookInfoStatusBar'; import { DeepnoteNotebookManager } from './deepnote/deepnoteNotebookManager'; import { IDeepnoteNotebookManager } from './types'; import { DeepnoteNotebookCommandListener } from './deepnote/deepnoteNotebookCommandListener'; @@ -108,6 +109,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteActivationService ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + DeepnoteNotebookInfoStatusBar + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNotebookCommandListener diff --git a/src/platform/common/constants.ts b/src/platform/common/constants.ts index c85719faec..d60321f1cf 100644 --- a/src/platform/common/constants.ts +++ b/src/platform/common/constants.ts @@ -223,6 +223,7 @@ export namespace Commands { export const OpenDeepnoteNotebook = 'deepnote.openNotebook'; export const OpenDeepnoteFile = 'deepnote.openFile'; export const RevealInDeepnoteExplorer = 'deepnote.revealInExplorer'; + export const CopyNotebookDetails = 'deepnote.copyNotebookDetails'; export const EnableSnapshots = 'deepnote.enableSnapshots'; export const DisableSnapshots = 'deepnote.disableSnapshots'; export const ManageIntegrations = 'deepnote.manageIntegrations'; From d44c1aa3ad78b5de95470a6d8aa55825a38d9a24 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 2 Apr 2026 14:04:33 +0000 Subject: [PATCH 25/47] feat(deepnote): Add command to copy active notebook details - Introduced a new command `deepnote.copyNotebookDetails` to allow users to copy details of the active Deepnote notebook. - Updated localization file to include the title for the new command, enhancing user experience with clear labeling. --- package.json | 6 ++++++ package.nls.json | 1 + 2 files changed, 7 insertions(+) diff --git a/package.json b/package.json index dfba8b252b..c9f71a0d14 100644 --- a/package.json +++ b/package.json @@ -93,6 +93,12 @@ "category": "Deepnote", "icon": "$(reveal)" }, + { + "command": "deepnote.copyNotebookDetails", + "title": "%deepnote.commands.copyNotebookDetails.title%", + "category": "Deepnote", + "icon": "$(copy)" + }, { "command": "deepnote.enableSnapshots", "title": "%deepnote.commands.enableSnapshots.title%", diff --git a/package.nls.json b/package.nls.json index 35ee95ae66..d13aed9b17 100644 --- a/package.nls.json +++ b/package.nls.json @@ -250,6 +250,7 @@ "deepnote.commands.openNotebook.title": "Open Notebook", "deepnote.commands.openFile.title": "Open File", "deepnote.commands.revealInExplorer.title": "Reveal in Explorer", + "deepnote.commands.copyNotebookDetails.title": "Copy Active Deepnote Notebook Details", "deepnote.commands.enableSnapshots.title": "Enable Snapshots", "deepnote.commands.disableSnapshots.title": "Disable Snapshots", "deepnote.commands.manageIntegrations.title": "Manage Integrations", From 7815f42a769c7828a02a789b8f90fd01d6945dc7 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 2 Apr 2026 14:47:57 +0000 Subject: [PATCH 26/47] Add a failing test for external change rerender --- .../deepnoteFileChangeWatcher.unit.test.ts | 98 +++++++++++++++++-- 1 file changed, 88 insertions(+), 10 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index f14a8998f3..a4f4886baf 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -8,6 +8,7 @@ import { Disposable, EventEmitter, FileSystemWatcher, + NotebookCellData, NotebookCellKind, NotebookDocument, NotebookEdit, @@ -21,6 +22,7 @@ import type { DeepnoteOutput, DeepnoteProject } from '../../platform/deepnote/de import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; import type { IControllerRegistration } from '../controllers/types'; import { IDeepnoteNotebookManager } from '../types'; +import { DeepnoteDataConverter } from './deepnoteDataConverter'; import { DeepnoteFileChangeWatcher } from './deepnoteFileChangeWatcher'; import { SnapshotService } from './snapshots/snapshotService'; @@ -209,20 +211,26 @@ suite('DeepnoteFileChangeWatcher', () => { notebookType?: string; cellCount?: number; metadata?: Record; - cells?: Array<{ - metadata?: Record; - outputs: any[]; - kind?: number; - document?: { getText: () => string; languageId?: string }; - }>; + cells?: Array< + | { + metadata?: Record; + outputs: any[]; + kind?: number; + document?: { getText: () => string; languageId?: string }; + } + | NotebookCellData + >; }): NotebookDocument { const cells = (opts.cells ?? []).map((c) => ({ ...c, kind: c.kind ?? NotebookCellKind.Code, - document: { - getText: c.document?.getText ?? (() => ''), - languageId: c.document?.languageId ?? 'python' - } + document: + 'document' in c && c.document + ? { getText: c.document.getText ?? (() => ''), languageId: c.document.languageId ?? 'python' } + : { + getText: () => ('value' in c ? (c.value as string) : ''), + languageId: 'languageId' in c ? (c.languageId as string) : 'python' + } })); return { @@ -544,6 +552,76 @@ project: assert.isAtLeast(applyEditCount, 2, 'applyEdit should be called for both external changes'); }); + test('editor→external→editor→external: second external edit must reload (self-write leak regression)', async function () { + this.timeout(15_000); + const uri = testFileUri('self-write-leak.deepnote'); + + when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); + + // Initial state: editor content matches disk — use the real converter + const converter = new DeepnoteDataConverter(); + const nb1 = multiNotebookProject.project.notebooks[0]; + const notebook = createMockNotebook({ + uri, + metadata: { deepnoteProjectId: multiNotebookProject.project.id, deepnoteNotebookId: nb1.id }, + cells: converter.convertBlocksToCells(nb1.blocks) + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(multiNotebookYaml); + + // Real VS Code behavior: workspace.save() fires onDidSaveNotebookDocument. + // executeMainFileSync calls markSelfWrite before workspace.save, AND the + // onDidSaveNotebookDocument handler also calls markSelfWrite — two marks + // for one FS event. This leaks a phantom self-write count. + when(mockedVSCodeNamespaces.workspace.save(anything())).thenCall((saveUri: Uri) => { + saveCount++; + onDidSaveNotebook.fire(notebook); + return Promise.resolve(saveUri); + }); + + // Step 1: editor edit → user saves notebook 1 + onDidSaveNotebook.fire(notebook); + onDidChangeFile.fire(uri); + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); + assert.strictEqual(applyEditCount, 0, 'Step 1: editor save FS event should be suppressed'); + + // Step 2: first external edit → triggers executeMainFileSync (reload works) + const externalProject1 = structuredClone(multiNotebookProject); + externalProject1.project.notebooks[0].blocks![0].content = 'print("external-1")'; + setupMockFs(serializeDeepnoteFile(externalProject1)); + + onDidChangeFile.fire(uri); + await waitFor(() => saveCount >= 1); + assert.isAtLeast(applyEditCount, 1, 'Step 2: first external edit should reload'); + + const applyCountAfterFirstReload = applyEditCount; + + // Consume FS events from executeMainFileSync's writeFile + save + onDidChangeFile.fire(uri); + onDidChangeFile.fire(uri); + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); + + // Step 3: editor edit → user saves again + onDidSaveNotebook.fire(notebook); + onDidChangeFile.fire(uri); + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); + + // Step 4: second external edit → should reload but phantom self-write blocks it + const externalProject2 = structuredClone(multiNotebookProject); + externalProject2.project.notebooks[0].blocks![0].content = 'print("external-2")'; + setupMockFs(serializeDeepnoteFile(externalProject2)); + + onDidChangeFile.fire(uri); + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs + 500)); + + assert.isAbove( + applyEditCount, + applyCountAfterFirstReload, + 'Step 4: second external edit should reload, but phantom self-write from executeMainFileSync leaks and suppresses it' + ); + }); + test('should use atomic edit (single applyEdit for replaceCells + metadata)', async () => { const uri = testFileUri('test.deepnote'); const notebook = createMockNotebook({ uri, cellCount: 0 }); From 855f38c6b0614921d0d5a943976a913d1b5bc5c0 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 2 Apr 2026 15:56:07 +0000 Subject: [PATCH 27/47] Refactor self write mark handling --- .../deepnote/deepnoteFileChangeWatcher.ts | 20 +++++++------------ 1 file changed, 7 insertions(+), 13 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index efa916b7fe..1a5caf0459 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -387,16 +387,14 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic // Save to clear dirty state. VS Code serializes (same bytes) and sees the // mtime from our recent write, so no "content is newer" conflict. - this.markSelfWrite(fileUri); + // NOTE: onDidSaveNotebookDocument handles the self-write mark for this save. try { const saved = await workspace.save(notebook.uri); if (!saved) { - this.consumeSelfWrite(fileUri); logger.warn(`[FileChangeWatcher] Save after sync write returned undefined: ${notebook.uri.path}`); return; } } catch (saveError) { - this.consumeSelfWrite(fileUri); logger.warn(`[FileChangeWatcher] Save after sync write failed: ${notebook.uri.path}`, saveError); } } catch (serializeError) { @@ -566,14 +564,9 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic return; } - // Save to sync mtime — mark as self-write first - this.markSelfWrite(notebook.uri); - try { - await workspace.save(notebook.uri); - } catch (error) { - this.consumeSelfWrite(notebook.uri); - throw error; - } + // Save to sync mtime. + // NOTE: onDidSaveNotebookDocument handles the self-write mark for this save. + await workspace.save(notebook.uri); logger.info(`[FileChangeWatcher] Updated notebook outputs from external snapshot: ${notebook.uri.path}`); } @@ -645,8 +638,9 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } /** - * Marks a URI as about to be written by us (workspace.save). - * Call before workspace.save() to prevent the resulting fs event from triggering a reload. + * Marks a URI as about to be written by us. + * For workspace.fs.writeFile(): call this before the write. + * For workspace.save(): do NOT call this — onDidSaveNotebookDocument handles it. */ private markSelfWrite(uri: Uri): void { const key = this.normalizeFileUri(uri); From f72b0fa4814d9e6e50f48dd5bcbf5f07f989f0e8 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 2 Apr 2026 18:42:32 +0000 Subject: [PATCH 28/47] feat(deepnote): Enhance notebook ID resolution with tab-based logic - Added support for resolving notebook IDs from open tabs, improving the handling of notebook selections during reloads. - Introduced a new method `findNotebookIdsFromTabs` to extract notebook IDs from the current tab groups. - Updated `findCurrentNotebookId` to prioritize tab-based resolution when available, alongside existing resolution strategies. - Enhanced unit tests to cover various scenarios for tab-based resolution, ensuring robust functionality. --- src/notebooks/deepnote/deepnoteSerializer.ts | 51 ++++++++++++- .../deepnote/deepnoteSerializer.unit.test.ts | 75 ++++++++++++++++++- src/test/vscode-mock.ts | 1 + 3 files changed, 122 insertions(+), 5 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index a9a0c8491d..cddaf6a60e 100644 --- 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, workspace, type CancellationToken, type NotebookData, type NotebookSerializer } from 'vscode'; +import { l10n, window, workspace, type CancellationToken, type NotebookData, type NotebookSerializer } from 'vscode'; import { z } from 'zod'; import { computeHash } from '../../platform/common/crypto'; @@ -56,6 +56,8 @@ const recentSerializeTtlMs = 5_000; @injectable() export class DeepnoteNotebookSerializer implements NotebookSerializer { + private readonly consumedTabResolutions = new Map>(); + private converter = new DeepnoteDataConverter(); /** @@ -107,7 +109,8 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } const projectId = deepnoteFile.project.id; - const resolvedNotebookId = notebookId ?? this.findCurrentNotebookId(projectId); + const projectNotebookIds = deepnoteFile.project.notebooks.map((nb) => nb.id); + const resolvedNotebookId = notebookId ?? this.findCurrentNotebookId(projectId, projectNotebookIds); logger.debug(`DeepnoteSerializer: Project ID: ${projectId}, Selected notebook ID: ${resolvedNotebookId}`); @@ -200,11 +203,13 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { /** * Finds the notebook ID to deserialize without relying on active-editor state. - * Prefers a pending resolution hint, then current/open-document state. + * Prefers a pending resolution hint, then tab-based resolution (for reload), + * then current/open-document state. * @param projectId The project ID to find a notebook for + * @param projectNotebookIds Optional list of notebook IDs in the project, enables tab-based resolution * @returns The notebook ID to deserialize, or undefined if none found */ - findCurrentNotebookId(projectId: string): string | undefined { + findCurrentNotebookId(projectId: string, projectNotebookIds?: string[]): string | undefined { const pendingNotebookId = this.notebookManager.consumePendingNotebookResolution(projectId); const openNotebookIds = this.findOpenNotebookIds(projectId); const currentNotebookId = this.notebookManager.getCurrentNotebookId(projectId); @@ -213,6 +218,21 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return pendingNotebookId; } + if (projectNotebookIds && projectNotebookIds.length > 0) { + const tabNotebookIds = this.findNotebookIdsFromTabs(projectNotebookIds); + const consumedIds = this.consumedTabResolutions.get(projectId) ?? new Set(); + const remaining = tabNotebookIds.filter((id) => !consumedIds.has(id) && !openNotebookIds.includes(id)); + + if (remaining.length > 0) { + const consumed = this.consumedTabResolutions.get(projectId) ?? new Set(); + + consumed.add(remaining[0]); + this.consumedTabResolutions.set(projectId, consumed); + + return remaining[0]; + } + } + if (currentNotebookId && openNotebookIds.length === 0) { return currentNotebookId; } @@ -591,6 +611,29 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return false; } + private findNotebookIdsFromTabs(projectNotebookIds: string[]): string[] { + const projectIdSet = new Set(projectNotebookIds); + const notebookIds = new Set(); + + for (const group of window.tabGroups.all) { + for (const tab of group.tabs) { + const input = tab.input as { uri?: { query?: string }; notebookType?: string } | undefined; + + if (!input || !('uri' in input) || !('notebookType' in input) || input.notebookType !== 'deepnote') { + continue; + } + + const notebookId = new URLSearchParams(input.uri?.query ?? '').get('notebook'); + + if (notebookId && projectIdSet.has(notebookId)) { + notebookIds.add(notebookId); + } + } + } + + return [...notebookIds]; + } + private findOpenNotebookIds(projectId: string): string[] { return [ ...workspace.notebookDocuments.reduce((ids, doc) => { diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index 4cde84ebb5..919af1ff23 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -2,7 +2,7 @@ import { deserializeDeepnoteFile, serializeDeepnoteFile, type DeepnoteFile } fro import { assert } from 'chai'; import { parse as parseYaml } from 'yaml'; import { when } from 'ts-mockito'; -import type { NotebookDocument } from 'vscode'; +import { Uri, type NotebookDocument } from 'vscode'; import { DeepnoteNotebookSerializer } from './deepnoteSerializer'; import { DeepnoteNotebookManager } from './deepnoteNotebookManager'; @@ -430,6 +430,7 @@ project: // Reset only the specific mocks used in this suite when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(undefined); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); + when(mockedVSCodeNamespaces.window.tabGroups).thenReturn({ all: [] } as any); }); test('should return queued notebook resolution when available', () => { @@ -705,6 +706,78 @@ project: assert.strictEqual(result, undefined); }); + + suite('tab-based resolution', () => { + function createTabGroups(...notebookIds: string[]) { + const tabs = notebookIds.map((id) => ({ + input: { + uri: Uri.parse(`file:///test/project.deepnote?notebook=${id}`), + notebookType: 'deepnote' + } + })); + + return { all: [{ tabs }] } as any; + } + + teardown(() => { + when(mockedVSCodeNamespaces.window.tabGroups).thenReturn({ all: [] } as any); + }); + + test('should resolve different notebook IDs from tabs on sequential reload calls', () => { + when(mockedVSCodeNamespaces.window.tabGroups).thenReturn(createTabGroups('notebook-1', 'notebook-2')); + + const projectNotebookIds = ['notebook-1', 'notebook-2']; + + const first = serializer.findCurrentNotebookId('project-123', projectNotebookIds); + const second = serializer.findCurrentNotebookId('project-123', projectNotebookIds); + + assert.strictEqual(first, 'notebook-1'); + assert.strictEqual(second, 'notebook-2'); + }); + + test('should skip tab resolution when all tab notebook IDs are already open', () => { + when(mockedVSCodeNamespaces.window.tabGroups).thenReturn(createTabGroups('notebook-1', 'notebook-2')); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([ + createNotebookDoc('notebook-1'), + createNotebookDoc('notebook-2') + ]); + + const projectNotebookIds = ['notebook-1', 'notebook-2']; + const result = serializer.findCurrentNotebookId('project-123', projectNotebookIds); + + assert.strictEqual(result, undefined); + }); + + test('should filter tab notebook IDs by projectNotebookIds', () => { + when(mockedVSCodeNamespaces.window.tabGroups).thenReturn( + createTabGroups('notebook-1', 'other-project-notebook') + ); + + const projectNotebookIds = ['notebook-1', 'notebook-2']; + const result = serializer.findCurrentNotebookId('project-123', projectNotebookIds); + + assert.strictEqual(result, 'notebook-1'); + }); + + test('should prioritize pending resolution over tab resolution', () => { + when(mockedVSCodeNamespaces.window.tabGroups).thenReturn(createTabGroups('notebook-1', 'notebook-2')); + + manager.queueNotebookResolution('project-123', 'queued-notebook'); + + const projectNotebookIds = ['notebook-1', 'notebook-2']; + const result = serializer.findCurrentNotebookId('project-123', projectNotebookIds); + + assert.strictEqual(result, 'queued-notebook'); + }); + + test('should skip tab resolution when projectNotebookIds is not provided', () => { + when(mockedVSCodeNamespaces.window.tabGroups).thenReturn(createTabGroups('notebook-1', 'notebook-2')); + + const result = serializer.findCurrentNotebookId('project-123'); + + assert.strictEqual(result, undefined); + }); + }); }); suite('component integration', () => { diff --git a/src/test/vscode-mock.ts b/src/test/vscode-mock.ts index b6b0f5371f..7142c44d75 100644 --- a/src/test/vscode-mock.ts +++ b/src/test/vscode-mock.ts @@ -83,6 +83,7 @@ export function resetVSCodeMocks() { when(mockedVSCodeNamespaces.window.visibleNotebookEditors).thenReturn([]); when(mockedVSCodeNamespaces.window.activeTextEditor).thenReturn(undefined); when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(undefined); + when(mockedVSCodeNamespaces.window.tabGroups).thenReturn({ all: [] } as any); // Window dialog methods with overloads (1-5 parameters) // showInformationMessage From 8e6f679a61c4227ba1381c7e822a0700bae4cb47 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 7 Apr 2026 13:10:11 +0000 Subject: [PATCH 29/47] feat(deepnote): Refactor notebook selection handling and enhance mismatch resolution - Removed the `selectNotebookForProject` and `clearNotebookSelection` methods from `DeepnoteNotebookManager` to simplify notebook selection logic. - Introduced a new mechanism in `DeepnoteActivationService` to check and fix notebook ID mismatches when opening documents, improving the reliability of notebook loading. - Updated `DeepnoteSerializer` to prioritize current notebook IDs and handle active tab resolutions more effectively. - Enhanced unit tests to cover the new mismatch resolution logic and ensure proper functionality across various scenarios. --- build/mocha-esm-loader.js | 1 + .../deepnote/deepnoteActivationService.ts | 118 ++++++++++++- .../deepnote/deepnoteExplorerView.ts | 1 - .../deepnote/deepnoteFileChangeWatcher.ts | 3 - .../deepnoteFileChangeWatcher.unit.test.ts | 8 +- .../deepnote/deepnoteNotebookManager.ts | 38 ----- .../deepnoteNotebookManager.unit.test.ts | 160 +----------------- src/notebooks/deepnote/deepnoteSerializer.ts | 83 +++++---- .../deepnote/deepnoteSerializer.unit.test.ts | 125 +------------- src/notebooks/types.ts | 4 - src/test/mocks/vsc/extHostedTypes.ts | 10 ++ src/test/vscode-mock.ts | 1 + 12 files changed, 179 insertions(+), 373 deletions(-) diff --git a/build/mocha-esm-loader.js b/build/mocha-esm-loader.js index 7a45afa78b..f948d8aae6 100644 --- a/build/mocha-esm-loader.js +++ b/build/mocha-esm-loader.js @@ -335,6 +335,7 @@ export async function load(url, context, nextLoad) { export const EndOfLine = createClassProxy('EndOfLine'); export const PortAutoForwardAction = createClassProxy('PortAutoForwardAction'); export const PortAttributes = createClassProxy('PortAttributes'); + export const TabInputNotebook = createClassProxy('TabInputNotebook'); `, shortCircuit: true }; diff --git a/src/notebooks/deepnote/deepnoteActivationService.ts b/src/notebooks/deepnote/deepnoteActivationService.ts index 96c35288cd..8958ad7e22 100644 --- a/src/notebooks/deepnote/deepnoteActivationService.ts +++ b/src/notebooks/deepnote/deepnoteActivationService.ts @@ -1,5 +1,17 @@ import { inject, injectable, optional } from 'inversify'; -import { commands, l10n, workspace, window, type Disposable, type NotebookDocumentContentOptions } from 'vscode'; +import { + CancellationTokenSource, + commands, + l10n, + NotebookEdit, + NotebookRange, + workspace, + window, + WorkspaceEdit, + type Disposable, + type NotebookDocument, + type NotebookDocumentContentOptions +} from 'vscode'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { IExtensionContext } from '../../platform/common/types'; @@ -15,6 +27,9 @@ import { SnapshotService } from './snapshots/snapshotService'; * Service responsible for activating and configuring Deepnote notebook support in VS Code. * Registers serializers, command handlers, and manages the notebook selection workflow. */ +const MISMATCH_CHECK_DELAY_MS = 200; +const MAX_MISMATCH_RETRIES = 10; + @injectable() export class DeepnoteActivationService implements IExtensionSyncActivationService { private editProtection: DeepnoteInputBlockEditProtection; @@ -23,6 +38,10 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic private integrationManager: IIntegrationManager; + private mismatchCheckTimer: ReturnType | undefined; + + private mismatchRetryCount = 0; + private serializer: DeepnoteNotebookSerializer; private serializerRegistration?: Disposable; @@ -49,6 +68,18 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic this.editProtection = new DeepnoteInputBlockEditProtection(this.logger); this.snapshotsEnabled = this.isSnapshotsEnabled(); + this.extensionContext.subscriptions.push( + workspace.onDidOpenNotebookDocument((doc) => { + if (doc.notebookType !== 'deepnote') { + return; + } + + if (new URLSearchParams(doc.uri.query).has('notebook')) { + this.scheduleMismatchCheck(); + } + }) + ); + this.registerSerializer(); this.extensionContext.subscriptions.push(this.editProtection); this.extensionContext.subscriptions.push( @@ -70,6 +101,80 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic this.integrationManager.activate(); } + /** + * Checks all open deepnote documents for URI/metadata notebook ID mismatches + * and fixes them by re-deserializing with the correct notebook ID. + * This handles the case where VS Code's deserializeNotebook API does not + * pass the document URI, causing the wrong notebook to be loaded on reload. + */ + private async checkAndFixMismatches(): Promise { + const hasLoadingDocs = workspace.notebookDocuments.some( + (doc) => + doc.notebookType === 'deepnote' && + !doc.metadata?.deepnoteNotebookId && + new URLSearchParams(doc.uri.query).has('notebook') + ); + + if (hasLoadingDocs && this.mismatchRetryCount < MAX_MISMATCH_RETRIES) { + this.mismatchRetryCount++; + this.scheduleMismatchCheck(); + + return; + } + + this.mismatchRetryCount = 0; + + for (const doc of workspace.notebookDocuments) { + if (doc.notebookType !== 'deepnote' || doc.isClosed) { + continue; + } + + const uriNotebookId = new URLSearchParams(doc.uri.query).get('notebook'); + const metadataNotebookId = doc.metadata?.deepnoteNotebookId as string | undefined; + + if (!uriNotebookId || uriNotebookId === metadataNotebookId) { + continue; + } + + await this.fixDocumentNotebook(doc, uriNotebookId); + } + } + + private async fixDocumentNotebook(doc: NotebookDocument, correctNotebookId: string): Promise { + const fileUri = doc.uri.with({ query: '', fragment: '' }); + + let content: Uint8Array; + try { + content = await workspace.fs.readFile(fileUri); + } catch { + this.logger.warn(`[DeepnoteActivation] Cannot read file for mismatch fix: ${fileUri.path}`); + + return; + } + + const cts = new CancellationTokenSource(); + try { + const data = await this.serializer.deserializeNotebook(content, cts.token, correctNotebookId); + + const wsEdit = new WorkspaceEdit(); + wsEdit.set(doc.uri, [ + NotebookEdit.replaceCells(new NotebookRange(0, doc.cellCount), data.cells), + NotebookEdit.updateNotebookMetadata(data.metadata!) + ]); + await workspace.applyEdit(wsEdit); + await doc.save(); + + this.logger.info( + `[DeepnoteActivation] Fixed notebook mismatch for ${doc.uri.path}: ` + + `loaded ${correctNotebookId} (was ${doc.metadata?.deepnoteNotebookId})` + ); + } catch (error) { + this.logger.error(`[DeepnoteActivation] Failed to fix notebook mismatch: ${doc.uri.path}`, error); + } finally { + cts.dispose(); + } + } + private isSnapshotsEnabled(): boolean { if (this.snapshotService) { return this.snapshotService.isSnapshotsEnabled(); @@ -130,4 +235,15 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic this.serializerRegistration = workspace.registerNotebookSerializer('deepnote', this.serializer, contentOptions); this.extensionContext.subscriptions.push(this.serializerRegistration); } + + private scheduleMismatchCheck(): void { + if (this.mismatchCheckTimer !== undefined) { + clearTimeout(this.mismatchCheckTimer); + } + + this.mismatchCheckTimer = setTimeout(() => { + this.mismatchCheckTimer = undefined; + void this.checkAndFixMismatches(); + }, MISMATCH_CHECK_DELAY_MS); + } } diff --git a/src/notebooks/deepnote/deepnoteExplorerView.ts b/src/notebooks/deepnote/deepnoteExplorerView.ts index 6d450401ba..9c41083e5f 100644 --- a/src/notebooks/deepnote/deepnoteExplorerView.ts +++ b/src/notebooks/deepnote/deepnoteExplorerView.ts @@ -519,7 +519,6 @@ export class DeepnoteExplorerView { private registerNotebookOpenIntent(projectId: string, notebookId: string): void { this.manager.queueNotebookResolution(projectId, notebookId); - this.manager.selectNotebookForProject(projectId, notebookId); } private refreshExplorer(): void { diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index 1a5caf0459..8e54997f2e 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -292,9 +292,6 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic return; } - // Pass the notebook ID explicitly to avoid mutating the global selection state. - // Multiple notebooks from the same project may be open simultaneously, and - // mutating selectedNotebookByProject would cause race conditions. const targetNotebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; const tokenSource = new CancellationTokenSource(); diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index a4f4886baf..fedca12907 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -3,7 +3,7 @@ import { assert } from 'chai'; import { mkdtempSync, rmSync } from 'fs'; import { tmpdir } from 'os'; import * as sinon from 'sinon'; -import { anything, instance, mock, reset, resetCalls, verify, when } from 'ts-mockito'; +import { anything, instance, mock, reset, resetCalls, when } from 'ts-mockito'; import { Disposable, EventEmitter, @@ -164,7 +164,6 @@ suite('DeepnoteFileChangeWatcher', () => { mockedNotebookManager = mock(); when(mockedNotebookManager.consumePendingNotebookResolution(anything())).thenReturn(undefined); when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(validProject); - when(mockedNotebookManager.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1'); when(mockedNotebookManager.queueNotebookResolution(anything(), anything())).thenReturn(); when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); mockNotebookManager = instance(mockedNotebookManager); @@ -1550,7 +1549,6 @@ project: const mockedManagerEx = mock(); when(mockedManagerEx.consumePendingNotebookResolution(anything())).thenReturn(undefined); when(mockedManagerEx.getOriginalProject(anything())).thenReturn(validProject); - when(mockedManagerEx.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1'); when(mockedManagerEx.queueNotebookResolution(anything(), anything())).thenReturn(); when(mockedManagerEx.updateOriginalProject(anything(), anything())).thenReturn(); @@ -1612,7 +1610,6 @@ project: reset(mockedNotebookManager); when(mockedNotebookManager.consumePendingNotebookResolution(anything())).thenReturn(undefined); when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(validProject); - when(mockedNotebookManager.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1'); when(mockedNotebookManager.queueNotebookResolution(anything(), anything())).thenReturn(); when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); resetCalls(mockedNotebookManager); @@ -2157,7 +2154,6 @@ project: reset(mockedNotebookManager); when(mockedNotebookManager.consumePendingNotebookResolution(anything())).thenReturn(undefined); when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); - when(mockedNotebookManager.getTheSelectedNotebookForAProject(anything())).thenReturn('notebook-1'); when(mockedNotebookManager.queueNotebookResolution(anything(), anything())).thenReturn(); when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); resetCalls(mockedNotebookManager); @@ -2278,8 +2274,6 @@ project: onDidChangeFile.fire(basePath); await waitFor(() => applyEditCount >= 2); - - 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 index 68360c7df4..b2e654e673 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.ts @@ -20,15 +20,6 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { private readonly originalProjects = new Map(); private readonly pendingNotebookResolutions = new Map(); private readonly projectsWithInitNotebookRun = new Set(); - private readonly selectedNotebookByProject = new Map(); - - /** - * 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. @@ -66,15 +57,6 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { return this.originalProjects.get(projectId); } - /** - * Gets the selected notebook ID for a specific project. - * @param projectId Project identifier - * @returns Selected notebook ID or undefined if not set - */ - getTheSelectedNotebookForAProject(projectId: string): string | undefined { - return this.selectedNotebookByProject.get(projectId); - } - /** * Queues a short-lived notebook resolution hint for the next deserialize. * @@ -92,16 +74,6 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { 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 - */ - selectNotebookForProject(projectId: string, notebookId: string): void { - this.selectedNotebookByProject.set(projectId, notebookId); - } - /** * Stores the original project data and sets the initial current notebook. * This is used during deserialization to cache project data and track the active notebook. @@ -132,16 +104,6 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { this.originalProjects.set(projectId, clonedProject); } - /** - * Updates the current notebook ID for a project. - * Used when switching notebooks within the same project. - * @param projectId Project identifier - * @param notebookId New current notebook ID - */ - updateCurrentNotebookId(projectId: string, notebookId: string): void { - this.currentNotebookId.set(projectId, notebookId); - } - /** * Updates the integrations list in the project data. * This modifies the stored project to reflect changes in configured integrations. diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts index b4276c605e..8a7de52200 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts @@ -39,15 +39,6 @@ suite('DeepnoteNotebookManager', () => { assert.strictEqual(result, 'notebook-456'); }); - - test('should return updated notebook ID', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); - manager.updateCurrentNotebookId('project-123', 'notebook-789'); - - const result = manager.getCurrentNotebookId('project-123'); - - assert.strictEqual(result, 'notebook-789'); - }); }); suite('getOriginalProject', () => { @@ -66,33 +57,6 @@ suite('DeepnoteNotebookManager', () => { }); }); - suite('getTheSelectedNotebookForAProject', () => { - test('should return undefined for unknown project', () => { - const result = manager.getTheSelectedNotebookForAProject('unknown-project'); - - assert.strictEqual(result, undefined); - }); - - test('should return notebook ID after setting', () => { - manager.selectNotebookForProject('project-123', 'notebook-456'); - - const result = manager.getTheSelectedNotebookForAProject('project-123'); - - assert.strictEqual(result, 'notebook-456'); - }); - - test('should handle multiple projects independently', () => { - manager.selectNotebookForProject('project-1', 'notebook-1'); - manager.selectNotebookForProject('project-2', 'notebook-2'); - - const result1 = manager.getTheSelectedNotebookForAProject('project-1'); - const result2 = manager.getTheSelectedNotebookForAProject('project-2'); - - assert.strictEqual(result1, 'notebook-1'); - assert.strictEqual(result2, 'notebook-2'); - }); - }); - suite('consumePendingNotebookResolution', () => { test('should return undefined when no pending resolution exists', () => { const result = manager.consumePendingNotebookResolution('unknown-project'); @@ -126,70 +90,6 @@ suite('DeepnoteNotebookManager', () => { }); }); - suite('selectNotebookForProject', () => { - test('should store notebook selection for project', () => { - manager.selectNotebookForProject('project-123', 'notebook-456'); - - const selectedNotebook = manager.getTheSelectedNotebookForAProject('project-123'); - - assert.strictEqual(selectedNotebook, 'notebook-456'); - }); - - test('should overwrite existing selection', () => { - manager.selectNotebookForProject('project-123', 'notebook-456'); - manager.selectNotebookForProject('project-123', 'notebook-789'); - - const result = manager.getTheSelectedNotebookForAProject('project-123'); - - assert.strictEqual(result, 'notebook-789'); - }); - - test('should handle multiple projects independently', () => { - manager.selectNotebookForProject('project-1', 'notebook-1'); - manager.selectNotebookForProject('project-2', 'notebook-2'); - - const result1 = manager.getTheSelectedNotebookForAProject('project-1'); - const result2 = manager.getTheSelectedNotebookForAProject('project-2'); - - assert.strictEqual(result1, 'notebook-1'); - assert.strictEqual(result2, 'notebook-2'); - }); - }); - - suite('clearNotebookSelection', () => { - test('should clear selection for a project', () => { - manager.selectNotebookForProject('project-123', 'notebook-456'); - manager.clearNotebookSelection('project-123'); - - const selectedNotebook = manager.getTheSelectedNotebookForAProject('project-123'); - - assert.strictEqual(selectedNotebook, undefined); - }); - - test('should not affect other projects', () => { - manager.selectNotebookForProject('project-1', 'notebook-1'); - manager.selectNotebookForProject('project-2', 'notebook-2'); - manager.clearNotebookSelection('project-1'); - - assert.strictEqual(manager.getTheSelectedNotebookForAProject('project-1'), undefined); - assert.strictEqual(manager.getTheSelectedNotebookForAProject('project-2'), 'notebook-2'); - }); - - test('should be idempotent for unknown project', () => { - assert.doesNotThrow(() => { - manager.clearNotebookSelection('unknown-project'); - 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', () => { test('should store both project and current notebook ID', () => { manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); @@ -291,36 +191,6 @@ suite('DeepnoteNotebookManager', () => { }); }); - suite('updateCurrentNotebookId', () => { - test('should update notebook ID for existing project', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); - manager.updateCurrentNotebookId('project-123', 'notebook-789'); - - const result = manager.getCurrentNotebookId('project-123'); - - assert.strictEqual(result, 'notebook-789'); - }); - - test('should set notebook ID for new project', () => { - manager.updateCurrentNotebookId('new-project', 'notebook-123'); - - const result = manager.getCurrentNotebookId('new-project'); - - assert.strictEqual(result, 'notebook-123'); - }); - - test('should handle multiple projects independently', () => { - manager.updateCurrentNotebookId('project-1', 'notebook-1'); - manager.updateCurrentNotebookId('project-2', 'notebook-2'); - - const result1 = manager.getCurrentNotebookId('project-1'); - const result2 = manager.getCurrentNotebookId('project-2'); - - assert.strictEqual(result1, 'notebook-1'); - assert.strictEqual(result2, 'notebook-2'); - }); - }); - suite('updateProjectIntegrations', () => { test('should update integrations list for existing project and return true', () => { manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); @@ -409,9 +279,8 @@ suite('DeepnoteNotebookManager', () => { }); test('should update integrations when currentNotebookId is undefined and return true', () => { - // Store project with a notebook ID, then clear it to simulate the edge case - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); - manager.updateCurrentNotebookId('project-123', undefined as any); + // Use updateOriginalProject which doesn't set currentNotebookId + manager.updateOriginalProject('project-123', mockProject); const integrations: ProjectIntegration[] = [ { id: 'int-1', name: 'PostgreSQL', type: 'pgsql' }, @@ -435,38 +304,17 @@ suite('DeepnoteNotebookManager', () => { suite('integration scenarios', () => { test('should handle complete workflow for multiple projects', () => { manager.storeOriginalProject('project-1', mockProject, 'notebook-1'); - manager.selectNotebookForProject('project-1', 'notebook-1'); - manager.storeOriginalProject('project-2', mockProject, 'notebook-2'); - manager.selectNotebookForProject('project-2', 'notebook-2'); assert.strictEqual(manager.getCurrentNotebookId('project-1'), 'notebook-1'); assert.strictEqual(manager.getCurrentNotebookId('project-2'), 'notebook-2'); - assert.strictEqual(manager.getTheSelectedNotebookForAProject('project-1'), 'notebook-1'); - assert.strictEqual(manager.getTheSelectedNotebookForAProject('project-2'), 'notebook-2'); }); - test('should handle notebook switching within same project', () => { + test('should handle notebook switching within same project via storeOriginalProject', () => { manager.storeOriginalProject('project-123', mockProject, 'notebook-1'); - manager.selectNotebookForProject('project-123', 'notebook-1'); - - manager.updateCurrentNotebookId('project-123', 'notebook-2'); - manager.selectNotebookForProject('project-123', 'notebook-2'); + manager.storeOriginalProject('project-123', mockProject, 'notebook-2'); assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-2'); - assert.strictEqual(manager.getTheSelectedNotebookForAProject('project-123'), 'notebook-2'); - }); - - test('should maintain separation between current and selected notebook IDs', () => { - // Store original project sets current notebook - manager.storeOriginalProject('project-123', mockProject, 'notebook-original'); - - // Selecting a different notebook for the project - manager.selectNotebookForProject('project-123', 'notebook-selected'); - - // Both should be maintained independently - assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-original'); - assert.strictEqual(manager.getTheSelectedNotebookForAProject('project-123'), 'notebook-selected'); }); }); }); diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index cddaf6a60e..0653ba57dd 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -1,7 +1,15 @@ 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 { + TabInputNotebook, + l10n, + window, + workspace, + type CancellationToken, + type NotebookData, + type NotebookSerializer +} from 'vscode'; import { z } from 'zod'; import { computeHash } from '../../platform/common/crypto'; @@ -56,8 +64,6 @@ const recentSerializeTtlMs = 5_000; @injectable() export class DeepnoteNotebookSerializer implements NotebookSerializer { - private readonly consumedTabResolutions = new Map>(); - private converter = new DeepnoteDataConverter(); /** @@ -109,8 +115,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } const projectId = deepnoteFile.project.id; - const projectNotebookIds = deepnoteFile.project.notebooks.map((nb) => nb.id); - const resolvedNotebookId = notebookId ?? this.findCurrentNotebookId(projectId, projectNotebookIds); + const resolvedNotebookId = notebookId ?? this.findCurrentNotebookId(projectId); logger.debug(`DeepnoteSerializer: Project ID: ${projectId}, Selected notebook ID: ${resolvedNotebookId}`); @@ -203,13 +208,12 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { /** * Finds the notebook ID to deserialize without relying on active-editor state. - * Prefers a pending resolution hint, then tab-based resolution (for reload), - * then current/open-document state. + * Prefers a pending resolution hint, then current/open-document state, + * and falls back to the active tab's URI query param for session restore. * @param projectId The project ID to find a notebook for - * @param projectNotebookIds Optional list of notebook IDs in the project, enables tab-based resolution * @returns The notebook ID to deserialize, or undefined if none found */ - findCurrentNotebookId(projectId: string, projectNotebookIds?: string[]): string | undefined { + findCurrentNotebookId(projectId: string): string | undefined { const pendingNotebookId = this.notebookManager.consumePendingNotebookResolution(projectId); const openNotebookIds = this.findOpenNotebookIds(projectId); const currentNotebookId = this.notebookManager.getCurrentNotebookId(projectId); @@ -218,21 +222,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return pendingNotebookId; } - if (projectNotebookIds && projectNotebookIds.length > 0) { - const tabNotebookIds = this.findNotebookIdsFromTabs(projectNotebookIds); - const consumedIds = this.consumedTabResolutions.get(projectId) ?? new Set(); - const remaining = tabNotebookIds.filter((id) => !consumedIds.has(id) && !openNotebookIds.includes(id)); - - if (remaining.length > 0) { - const consumed = this.consumedTabResolutions.get(projectId) ?? new Set(); - - consumed.add(remaining[0]); - this.consumedTabResolutions.set(projectId, consumed); - - return remaining[0]; - } - } - if (currentNotebookId && openNotebookIds.length === 0) { return currentNotebookId; } @@ -262,6 +251,14 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } } + const activeTabNotebookId = this.findNotebookIdFromActiveTab(); + + if (activeTabNotebookId) { + logger.debug(`DeepnoteSerializer: Resolved notebook ID from active tab URI: ${activeTabNotebookId}`); + + return activeTabNotebookId; + } + return undefined; } @@ -404,11 +401,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { // for other open notebooks to resolve to the wrong notebook. this.notebookManager.updateOriginalProject(projectId, originalProject); this.recentSerializations.set(projectId, { notebookId, timestamp: Date.now() }); - const openNotebookIdsAtSerialize = this.findOpenNotebookIds(projectId); - - if (openNotebookIdsAtSerialize.length === 0) { - this.notebookManager.queueNotebookResolution(projectId, notebookId); - } logger.debug('SerializeNotebook: Serializing to YAML'); @@ -611,27 +603,28 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return false; } - private findNotebookIdsFromTabs(projectNotebookIds: string[]): string[] { - const projectIdSet = new Set(projectNotebookIds); - const notebookIds = new Set(); - - for (const group of window.tabGroups.all) { - for (const tab of group.tabs) { - const input = tab.input as { uri?: { query?: string }; notebookType?: string } | undefined; + /** + * Extracts the notebook ID from the active tab's URI query params. + * During session restore or tab open, VS Code may set the active tab + * before calling deserializeNotebook. The URI retains the + * `?notebook=` query param set when the tab was originally opened. + */ + private findNotebookIdFromActiveTab(): string | undefined { + const activeTab = window.tabGroups.activeTabGroup?.activeTab; - if (!input || !('uri' in input) || !('notebookType' in input) || input.notebookType !== 'deepnote') { - continue; - } + if (!activeTab || !(activeTab.input instanceof TabInputNotebook)) { + return undefined; + } - const notebookId = new URLSearchParams(input.uri?.query ?? '').get('notebook'); + const tabInput = activeTab.input; - if (notebookId && projectIdSet.has(notebookId)) { - notebookIds.add(notebookId); - } - } + if (tabInput.notebookType !== 'deepnote') { + return undefined; } - return [...notebookIds]; + const notebookId = new URLSearchParams(tabInput.uri.query).get('notebook'); + + return notebookId ?? undefined; } private findOpenNotebookIds(projectId: string): string[] { diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index 919af1ff23..7b0367c34d 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -2,7 +2,7 @@ import { deserializeDeepnoteFile, serializeDeepnoteFile, type DeepnoteFile } fro import { assert } from 'chai'; import { parse as parseYaml } from 'yaml'; import { when } from 'ts-mockito'; -import { Uri, type NotebookDocument } from 'vscode'; +import type { NotebookDocument } from 'vscode'; import { DeepnoteNotebookSerializer } from './deepnoteSerializer'; import { DeepnoteNotebookManager } from './deepnoteNotebookManager'; @@ -177,8 +177,8 @@ project: assert.include(result.cells[0].value, 'Title'); }); - test('should ignore stored selection when explicit notebookId is provided', async () => { - manager.selectNotebookForProject('project-123', 'notebook-1'); + test('should ignore queued resolution when explicit notebookId is provided', async () => { + manager.queueNotebookResolution('project-123', 'notebook-1'); const content = projectToYaml(mockProject); const result = await serializer.deserializeNotebook(content, {} as any, 'notebook-2'); @@ -253,9 +253,8 @@ project: assert.include(yamlString, 'notebook-1'); }); - test('should use current notebook ID instead of stale selected notebook when metadata notebook ID is missing', async () => { + test('should use current notebook ID when metadata notebook ID is missing', async () => { manager.storeOriginalProject('project-123', mockProject, 'notebook-2'); - manager.selectNotebookForProject('project-123', 'notebook-1'); const mockNotebookData = { cells: [ @@ -280,31 +279,6 @@ project: assert.strictEqual(serializedProject.project.notebooks[1].blocks?.[0].content, '# Updated second notebook'); }); - test('should queue the serialized notebook for the next resolution hint', async () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-1'); - - const mockNotebookData = { - cells: [ - { - kind: 1, // NotebookCellKind.Markup - value: '# Updated second notebook', - languageId: 'markdown', - metadata: {} - } - ], - metadata: { - deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'notebook-2' - } - }; - - await serializer.serializeNotebook(mockNotebookData as any, {} as any); - - manager.updateCurrentNotebookId('project-123', 'notebook-1'); - - assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-2'); - }); - suite('multi-notebook save scenarios', () => { teardown(() => { when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(undefined); @@ -450,7 +424,7 @@ project: test('should prioritize queued notebook resolution over current notebook and open documents', () => { manager.queueNotebookResolution('project-123', 'queued-notebook'); - manager.updateCurrentNotebookId('project-123', 'current-notebook'); + manager.storeOriginalProject('project-123', mockProject, 'current-notebook'); const mockNotebookDoc = { then: undefined, @@ -478,7 +452,7 @@ project: }); test('should return current notebook ID when no pending resolution exists', () => { - manager.updateCurrentNotebookId('project-123', 'current-notebook'); + manager.storeOriginalProject('project-123', mockProject, 'current-notebook'); const result = serializer.findCurrentNotebookId('project-123'); @@ -512,7 +486,7 @@ project: }); test('should prefer current notebook ID when multiple notebooks are open for a project', () => { - manager.updateCurrentNotebookId('project-123', 'notebook-b'); + manager.storeOriginalProject('project-123', mockProject, 'notebook-b'); const notebookA = { then: undefined, @@ -599,14 +573,6 @@ project: assert.strictEqual(result, undefined); }); - test('should ignore stale selected notebook state when no other resolver state exists', () => { - manager.selectNotebookForProject('project-123', 'stored-notebook'); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, undefined); - }); - test('after serializing notebook-1 with both notebooks open should return sibling notebook-2', async () => { manager.storeOriginalProject('project-123', mockProject, 'notebook-1'); @@ -706,78 +672,6 @@ project: assert.strictEqual(result, undefined); }); - - suite('tab-based resolution', () => { - function createTabGroups(...notebookIds: string[]) { - const tabs = notebookIds.map((id) => ({ - input: { - uri: Uri.parse(`file:///test/project.deepnote?notebook=${id}`), - notebookType: 'deepnote' - } - })); - - return { all: [{ tabs }] } as any; - } - - teardown(() => { - when(mockedVSCodeNamespaces.window.tabGroups).thenReturn({ all: [] } as any); - }); - - test('should resolve different notebook IDs from tabs on sequential reload calls', () => { - when(mockedVSCodeNamespaces.window.tabGroups).thenReturn(createTabGroups('notebook-1', 'notebook-2')); - - const projectNotebookIds = ['notebook-1', 'notebook-2']; - - const first = serializer.findCurrentNotebookId('project-123', projectNotebookIds); - const second = serializer.findCurrentNotebookId('project-123', projectNotebookIds); - - assert.strictEqual(first, 'notebook-1'); - assert.strictEqual(second, 'notebook-2'); - }); - - test('should skip tab resolution when all tab notebook IDs are already open', () => { - when(mockedVSCodeNamespaces.window.tabGroups).thenReturn(createTabGroups('notebook-1', 'notebook-2')); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([ - createNotebookDoc('notebook-1'), - createNotebookDoc('notebook-2') - ]); - - const projectNotebookIds = ['notebook-1', 'notebook-2']; - const result = serializer.findCurrentNotebookId('project-123', projectNotebookIds); - - assert.strictEqual(result, undefined); - }); - - test('should filter tab notebook IDs by projectNotebookIds', () => { - when(mockedVSCodeNamespaces.window.tabGroups).thenReturn( - createTabGroups('notebook-1', 'other-project-notebook') - ); - - const projectNotebookIds = ['notebook-1', 'notebook-2']; - const result = serializer.findCurrentNotebookId('project-123', projectNotebookIds); - - assert.strictEqual(result, 'notebook-1'); - }); - - test('should prioritize pending resolution over tab resolution', () => { - when(mockedVSCodeNamespaces.window.tabGroups).thenReturn(createTabGroups('notebook-1', 'notebook-2')); - - manager.queueNotebookResolution('project-123', 'queued-notebook'); - - const projectNotebookIds = ['notebook-1', 'notebook-2']; - const result = serializer.findCurrentNotebookId('project-123', projectNotebookIds); - - assert.strictEqual(result, 'queued-notebook'); - }); - - test('should skip tab resolution when projectNotebookIds is not provided', () => { - when(mockedVSCodeNamespaces.window.tabGroups).thenReturn(createTabGroups('notebook-1', 'notebook-2')); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, undefined); - }); - }); }); suite('component integration', () => { @@ -806,12 +700,7 @@ project: assert.isFunction(manager.consumePendingNotebookResolution, 'has consumePendingNotebookResolution method'); assert.isFunction(manager.getCurrentNotebookId, 'has getCurrentNotebookId method'); assert.isFunction(manager.getOriginalProject, 'has getOriginalProject method'); - assert.isFunction( - manager.getTheSelectedNotebookForAProject, - 'has getTheSelectedNotebookForAProject method' - ); assert.isFunction(manager.queueNotebookResolution, 'has queueNotebookResolution method'); - assert.isFunction(manager.selectNotebookForProject, 'has selectNotebookForProject method'); assert.isFunction(manager.storeOriginalProject, 'has storeOriginalProject method'); }); diff --git a/src/notebooks/types.ts b/src/notebooks/types.ts index 527a56ddee..2da356ee25 100644 --- a/src/notebooks/types.ts +++ b/src/notebooks/types.ts @@ -37,17 +37,13 @@ export interface ProjectIntegration { export const IDeepnoteNotebookManager = Symbol('IDeepnoteNotebookManager'); export interface IDeepnoteNotebookManager { - clearNotebookSelection(projectId: string): void; consumePendingNotebookResolution(projectId: string): string | undefined; getCurrentNotebookId(projectId: string): string | undefined; getOriginalProject(projectId: string): DeepnoteProject | undefined; - getTheSelectedNotebookForAProject(projectId: string): string | undefined; hasInitNotebookBeenRun(projectId: string): boolean; markInitNotebookAsRun(projectId: string): void; queueNotebookResolution(projectId: string, notebookId: string): void; - selectNotebookForProject(projectId: string, notebookId: string): void; storeOriginalProject(projectId: string, project: DeepnoteProject, notebookId: string): void; - updateCurrentNotebookId(projectId: string, notebookId: string): void; updateOriginalProject(projectId: string, project: DeepnoteProject): void; /** diff --git a/src/test/mocks/vsc/extHostedTypes.ts b/src/test/mocks/vsc/extHostedTypes.ts index 4a58f1c077..ab429319c4 100644 --- a/src/test/mocks/vsc/extHostedTypes.ts +++ b/src/test/mocks/vsc/extHostedTypes.ts @@ -2544,6 +2544,16 @@ export namespace vscMockExtHostedTypes { */ Default = 0 } + export class TabInputNotebook { + readonly uri: vscUri.URI; + readonly notebookType: string; + + constructor(uri: vscUri.URI, notebookType: string) { + this.uri = uri; + this.notebookType = notebookType; + } + } + export class NotebookCellData { kind: NotebookCellKind; value: string; diff --git a/src/test/vscode-mock.ts b/src/test/vscode-mock.ts index 7142c44d75..d413837dc7 100644 --- a/src/test/vscode-mock.ts +++ b/src/test/vscode-mock.ts @@ -301,6 +301,7 @@ export function resetVSCodeMocks() { (mockedVSCode as any).NotebookCellExecutionState = vscodeMocks.vscMockExtHostedTypes.NotebookCellExecutionState; (mockedVSCode as any).NotebookEditorRevealType = vscodeMocks.vscMockExtHostedTypes.NotebookEditorRevealType; // Mock ColorThemeKind enum + (mockedVSCode as any).TabInputNotebook = vscodeMocks.vscMockExtHostedTypes.TabInputNotebook; (mockedVSCode as any).ColorThemeKind = { Light: 1, Dark: 2, HighContrast: 3, HighContrastLight: 4 }; mockedVSCode.EndOfLine = vscodeMocks.vscMockExtHostedTypes.EndOfLine; } From 0396a2c1c7710b53558fef3c6bd566630d096a8a Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 7 Apr 2026 16:40:51 +0000 Subject: [PATCH 30/47] refactor(deepnote): Simplify DeepnoteActivationService and enhance notebook ID resolution - Removed unused variables and methods related to mismatch checking in `DeepnoteActivationService`, streamlining the activation process. - Updated `findCurrentNotebookId` in `DeepnoteSerializer` to improve notebook ID resolution logic, prioritizing active tab detection. - Adjusted unit tests to reflect changes in notebook ID resolution, ensuring accurate behavior when no pending resolutions exist. --- .../deepnote/deepnoteActivationService.ts | 118 +----------------- src/notebooks/deepnote/deepnoteSerializer.ts | 38 +++--- .../deepnote/deepnoteSerializer.unit.test.ts | 79 +----------- 3 files changed, 19 insertions(+), 216 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteActivationService.ts b/src/notebooks/deepnote/deepnoteActivationService.ts index 8958ad7e22..96c35288cd 100644 --- a/src/notebooks/deepnote/deepnoteActivationService.ts +++ b/src/notebooks/deepnote/deepnoteActivationService.ts @@ -1,17 +1,5 @@ import { inject, injectable, optional } from 'inversify'; -import { - CancellationTokenSource, - commands, - l10n, - NotebookEdit, - NotebookRange, - workspace, - window, - WorkspaceEdit, - type Disposable, - type NotebookDocument, - type NotebookDocumentContentOptions -} from 'vscode'; +import { commands, l10n, workspace, window, type Disposable, type NotebookDocumentContentOptions } from 'vscode'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { IExtensionContext } from '../../platform/common/types'; @@ -27,9 +15,6 @@ import { SnapshotService } from './snapshots/snapshotService'; * Service responsible for activating and configuring Deepnote notebook support in VS Code. * Registers serializers, command handlers, and manages the notebook selection workflow. */ -const MISMATCH_CHECK_DELAY_MS = 200; -const MAX_MISMATCH_RETRIES = 10; - @injectable() export class DeepnoteActivationService implements IExtensionSyncActivationService { private editProtection: DeepnoteInputBlockEditProtection; @@ -38,10 +23,6 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic private integrationManager: IIntegrationManager; - private mismatchCheckTimer: ReturnType | undefined; - - private mismatchRetryCount = 0; - private serializer: DeepnoteNotebookSerializer; private serializerRegistration?: Disposable; @@ -68,18 +49,6 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic this.editProtection = new DeepnoteInputBlockEditProtection(this.logger); this.snapshotsEnabled = this.isSnapshotsEnabled(); - this.extensionContext.subscriptions.push( - workspace.onDidOpenNotebookDocument((doc) => { - if (doc.notebookType !== 'deepnote') { - return; - } - - if (new URLSearchParams(doc.uri.query).has('notebook')) { - this.scheduleMismatchCheck(); - } - }) - ); - this.registerSerializer(); this.extensionContext.subscriptions.push(this.editProtection); this.extensionContext.subscriptions.push( @@ -101,80 +70,6 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic this.integrationManager.activate(); } - /** - * Checks all open deepnote documents for URI/metadata notebook ID mismatches - * and fixes them by re-deserializing with the correct notebook ID. - * This handles the case where VS Code's deserializeNotebook API does not - * pass the document URI, causing the wrong notebook to be loaded on reload. - */ - private async checkAndFixMismatches(): Promise { - const hasLoadingDocs = workspace.notebookDocuments.some( - (doc) => - doc.notebookType === 'deepnote' && - !doc.metadata?.deepnoteNotebookId && - new URLSearchParams(doc.uri.query).has('notebook') - ); - - if (hasLoadingDocs && this.mismatchRetryCount < MAX_MISMATCH_RETRIES) { - this.mismatchRetryCount++; - this.scheduleMismatchCheck(); - - return; - } - - this.mismatchRetryCount = 0; - - for (const doc of workspace.notebookDocuments) { - if (doc.notebookType !== 'deepnote' || doc.isClosed) { - continue; - } - - const uriNotebookId = new URLSearchParams(doc.uri.query).get('notebook'); - const metadataNotebookId = doc.metadata?.deepnoteNotebookId as string | undefined; - - if (!uriNotebookId || uriNotebookId === metadataNotebookId) { - continue; - } - - await this.fixDocumentNotebook(doc, uriNotebookId); - } - } - - private async fixDocumentNotebook(doc: NotebookDocument, correctNotebookId: string): Promise { - const fileUri = doc.uri.with({ query: '', fragment: '' }); - - let content: Uint8Array; - try { - content = await workspace.fs.readFile(fileUri); - } catch { - this.logger.warn(`[DeepnoteActivation] Cannot read file for mismatch fix: ${fileUri.path}`); - - return; - } - - const cts = new CancellationTokenSource(); - try { - const data = await this.serializer.deserializeNotebook(content, cts.token, correctNotebookId); - - const wsEdit = new WorkspaceEdit(); - wsEdit.set(doc.uri, [ - NotebookEdit.replaceCells(new NotebookRange(0, doc.cellCount), data.cells), - NotebookEdit.updateNotebookMetadata(data.metadata!) - ]); - await workspace.applyEdit(wsEdit); - await doc.save(); - - this.logger.info( - `[DeepnoteActivation] Fixed notebook mismatch for ${doc.uri.path}: ` + - `loaded ${correctNotebookId} (was ${doc.metadata?.deepnoteNotebookId})` - ); - } catch (error) { - this.logger.error(`[DeepnoteActivation] Failed to fix notebook mismatch: ${doc.uri.path}`, error); - } finally { - cts.dispose(); - } - } - private isSnapshotsEnabled(): boolean { if (this.snapshotService) { return this.snapshotService.isSnapshotsEnabled(); @@ -235,15 +130,4 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic this.serializerRegistration = workspace.registerNotebookSerializer('deepnote', this.serializer, contentOptions); this.extensionContext.subscriptions.push(this.serializerRegistration); } - - private scheduleMismatchCheck(): void { - if (this.mismatchCheckTimer !== undefined) { - clearTimeout(this.mismatchCheckTimer); - } - - this.mismatchCheckTimer = setTimeout(() => { - this.mismatchCheckTimer = undefined; - void this.checkAndFixMismatches(); - }, MISMATCH_CHECK_DELAY_MS); - } } diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 0653ba57dd..6aac61ad8d 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -207,33 +207,35 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } /** - * Finds the notebook ID to deserialize without relying on active-editor state. - * Prefers a pending resolution hint, then current/open-document state, - * and falls back to the active tab's URI query param for session restore. + * Finds the notebook ID to deserialize for VS Code-initiated deserialization. + * Priority: + * 1. Pending resolution hint (explicit intent from explorer view) + * 2. Active tab URI query param (VS Code's ground truth for which tab is loading) + * 3. Sibling detection (post-save re-read for other open tabs) * @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 { const pendingNotebookId = this.notebookManager.consumePendingNotebookResolution(projectId); - const openNotebookIds = this.findOpenNotebookIds(projectId); - const currentNotebookId = this.notebookManager.getCurrentNotebookId(projectId); if (pendingNotebookId) { + logger.debug(`DeepnoteSerializer: Resolved notebook ID from pending resolution: ${pendingNotebookId}`); return pendingNotebookId; } - if (currentNotebookId && openNotebookIds.length === 0) { - return currentNotebookId; - } + const activeTabNotebookId = this.findNotebookIdFromActiveTab(); + + if (activeTabNotebookId) { + logger.debug(`DeepnoteSerializer: Resolved notebook ID from active tab URI: ${activeTabNotebookId}`); - if (openNotebookIds.length === 1) { - return openNotebookIds[0]; + return activeTabNotebookId; } // Multiple notebooks open — VS Code may be re-reading the file for a // sibling tab after a save. Pick a sibling (any open notebook that is - // NOT the one just serialized). If no recent serialization, fall back - // to currentNotebookId for backward compatibility. + // NOT the one just serialized). + const openNotebookIds = this.findOpenNotebookIds(projectId); + if (openNotebookIds.length > 1) { const recent = this.recentSerializations.get(projectId); @@ -245,18 +247,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return sibling; } } - - if (currentNotebookId && openNotebookIds.includes(currentNotebookId)) { - return currentNotebookId; - } - } - - const activeTabNotebookId = this.findNotebookIdFromActiveTab(); - - if (activeTabNotebookId) { - logger.debug(`DeepnoteSerializer: Resolved notebook ID from active tab URI: ${activeTabNotebookId}`); - - return activeTabNotebookId; } return undefined; diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index 7b0367c34d..f0d9608f7b 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -451,83 +451,12 @@ project: assert.strictEqual(result, 'queued-notebook'); }); - test('should return current notebook ID when no pending resolution exists', () => { + test('should return undefined when no pending resolution or active tab exists', () => { manager.storeOriginalProject('project-123', mockProject, 'current-notebook'); const result = serializer.findCurrentNotebookId('project-123'); - assert.strictEqual(result, 'current-notebook'); - }); - - test('should return the only open notebook when current notebook ID is unavailable', () => { - const mockNotebookDoc = { - then: undefined, - notebookType: 'deepnote', - metadata: { - deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'notebook-from-workspace' - }, - uri: {} as any, - version: 1, - isDirty: false, - isUntitled: false, - isClosed: false, - cellCount: 0, - cellAt: () => ({}) as any, - getCells: () => [], - save: async () => true - } as NotebookDocument; - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([mockNotebookDoc]); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, 'notebook-from-workspace'); - }); - - test('should prefer current notebook ID when multiple notebooks are open for a project', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-b'); - - const notebookA = { - then: undefined, - notebookType: 'deepnote', - metadata: { - deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'notebook-a' - }, - uri: {} as any, - version: 1, - isDirty: false, - isUntitled: false, - isClosed: false, - cellCount: 0, - cellAt: () => ({}) as any, - getCells: () => [], - save: async () => true - } as NotebookDocument; - const notebookB = { - then: undefined, - notebookType: 'deepnote', - metadata: { - deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'notebook-b' - }, - uri: {} as any, - version: 1, - isDirty: false, - isUntitled: false, - isClosed: false, - cellCount: 0, - cellAt: () => ({}) as any, - getCells: () => [], - save: async () => true - } as NotebookDocument; - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookA, notebookB]); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, 'notebook-b'); + assert.strictEqual(result, undefined); }); test('should return undefined when multiple notebooks are open and no stronger signal exists', () => { @@ -629,7 +558,7 @@ project: await serializer.serializeNotebook(mockNotebookData as any, {} as any); assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-2'); - assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-1'); + assert.strictEqual(serializer.findCurrentNotebookId('project-123'), undefined); }); test('recent serialization hint expires after TTL', async () => { @@ -664,7 +593,7 @@ project: const result = serializer.findCurrentNotebookId('project-123'); - assert.strictEqual(result, 'notebook-1'); + assert.strictEqual(result, undefined); }); test('should return undefined for unknown project', () => { From 06d31979c118b7c9c2da0fd0a7ce5726a9bf3f06 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 7 Apr 2026 17:06:22 +0000 Subject: [PATCH 31/47] refactor(deepnote): Streamline DeepnoteNotebookSerializer and remove unused logic - Eliminated unnecessary variables and methods related to recent serialization tracking in `DeepnoteNotebookSerializer`, simplifying the notebook ID resolution process. - Updated the logic for finding the current notebook ID to focus on active tab detection, enhancing overall efficiency. - Removed outdated unit tests that were no longer relevant to the current implementation, ensuring test suite accuracy. --- src/notebooks/deepnote/deepnoteSerializer.ts | 51 ---- .../deepnote/deepnoteSerializer.unit.test.ts | 242 ------------------ 2 files changed, 293 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 6aac61ad8d..8bbdc0b796 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -10,7 +10,6 @@ import { type NotebookData, type NotebookSerializer } from 'vscode'; -import { z } from 'zod'; import { computeHash } from '../../platform/common/crypto'; import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; @@ -59,21 +58,10 @@ function cloneWithoutCircularRefs(obj: T, seen = new WeakSet()): T { * Serializer for converting between Deepnote YAML files and VS Code notebook format. * Handles reading/writing .deepnote files and manages project state persistence. */ -/** Window (ms) during which a post-serialize deserialize excludes the serialized notebook. */ -const recentSerializeTtlMs = 5_000; - @injectable() export class DeepnoteNotebookSerializer implements NotebookSerializer { private converter = new DeepnoteDataConverter(); - /** - * Tracks the most recently serialized notebook per project. - * When VS Code calls $dataToNotebook after a save, it's re-reading the - * file for a SIBLING tab (the saved notebook already has its content). - * This tracker lets findCurrentNotebookId pick a sibling instead. - */ - private readonly recentSerializations = new Map(); - constructor( @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, @inject(SnapshotService) @optional() private readonly snapshotService?: SnapshotService @@ -211,7 +199,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { * Priority: * 1. Pending resolution hint (explicit intent from explorer view) * 2. Active tab URI query param (VS Code's ground truth for which tab is loading) - * 3. Sibling detection (post-save re-read for other open tabs) * @param projectId The project ID to find a notebook for * @returns The notebook ID to deserialize, or undefined if none found */ @@ -219,7 +206,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { const pendingNotebookId = this.notebookManager.consumePendingNotebookResolution(projectId); if (pendingNotebookId) { - logger.debug(`DeepnoteSerializer: Resolved notebook ID from pending resolution: ${pendingNotebookId}`); return pendingNotebookId; } @@ -231,24 +217,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return activeTabNotebookId; } - // Multiple notebooks open — VS Code may be re-reading the file for a - // sibling tab after a save. Pick a sibling (any open notebook that is - // NOT the one just serialized). - const openNotebookIds = this.findOpenNotebookIds(projectId); - - if (openNotebookIds.length > 1) { - const recent = this.recentSerializations.get(projectId); - - if (recent && Date.now() - recent.timestamp < recentSerializeTtlMs) { - this.recentSerializations.delete(projectId); - const sibling = openNotebookIds.find((id) => id !== recent.notebookId); - - if (sibling) { - return sibling; - } - } - } - return undefined; } @@ -390,7 +358,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { // currentNotebookId here would cause VS Code's follow-up deserialize calls // for other open notebooks to resolve to the wrong notebook. this.notebookManager.updateOriginalProject(projectId, originalProject); - this.recentSerializations.set(projectId, { notebookId, timestamp: Date.now() }); logger.debug('SerializeNotebook: Serializing to YAML'); @@ -617,24 +584,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return notebookId ?? undefined; } - private findOpenNotebookIds(projectId: string): string[] { - return [ - ...workspace.notebookDocuments.reduce((ids, doc) => { - if (doc.notebookType !== 'deepnote' || doc.metadata.deepnoteProjectId !== projectId) { - return ids; - } - - const parsed = z.string().safeParse(doc.metadata.deepnoteNotebookId); - - if (parsed.success) { - ids.add(parsed.data); - } - - return ids; - }, new Set()) - ]; - } - /** * 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 index f0d9608f7b..139673a4b5 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -73,26 +73,6 @@ suite('DeepnoteNotebookSerializer', () => { return new TextEncoder().encode(yamlString); } - function createNotebookDoc(notebookId: string, projectId: string = 'project-123'): NotebookDocument { - return { - cellAt: () => ({}) as any, - cellCount: 0, - getCells: () => [], - isClosed: false, - isDirty: false, - isUntitled: false, - metadata: { - deepnoteNotebookId: notebookId, - deepnoteProjectId: projectId - }, - notebookType: 'deepnote', - save: async () => true, - then: undefined, - uri: {} as any, - version: 1 - } as NotebookDocument; - } - suite('deserializeNotebook', () => { test('should deserialize valid project with queued notebook resolution', async () => { manager.queueNotebookResolution('project-123', 'notebook-1'); @@ -311,91 +291,6 @@ project: assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-2'); }); - test('serializing notebook-1 then deserializing without explicit ID should resolve to sibling notebook-2', async () => { - const yamlBytes = projectToYaml(mockProject); - - manager.queueNotebookResolution('project-123', 'notebook-1'); - await serializer.deserializeNotebook(yamlBytes, {} as any); - manager.queueNotebookResolution('project-123', 'notebook-2'); - await serializer.deserializeNotebook(yamlBytes, {} as any); - - const notebookDoc1 = createNotebookDoc('notebook-1'); - const notebookDoc2 = createNotebookDoc('notebook-2'); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookDoc1, notebookDoc2]); - - const mockNotebookData = { - cells: [ - { - kind: 2, - languageId: 'python', - metadata: {}, - value: 'print("saved from nb1")' - } - ], - metadata: { - deepnoteNotebookId: 'notebook-1', - deepnoteProjectId: 'project-123' - } - }; - - const saved = await serializer.serializeNotebook(mockNotebookData as any, {} as any); - const afterSave = await serializer.deserializeNotebook(saved, {} as any); - - assert.strictEqual(afterSave.metadata?.deepnoteNotebookId, 'notebook-2'); - }); - - test('sequential saves: contextless deserialize returns the sibling after each serialize', async () => { - const yamlBytes = projectToYaml(mockProject); - - manager.queueNotebookResolution('project-123', 'notebook-1'); - await serializer.deserializeNotebook(yamlBytes, {} as any); - manager.queueNotebookResolution('project-123', 'notebook-2'); - await serializer.deserializeNotebook(yamlBytes, {} as any); - - const notebookDoc1 = createNotebookDoc('notebook-1'); - const notebookDoc2 = createNotebookDoc('notebook-2'); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookDoc1, notebookDoc2]); - - const dataNb1 = { - cells: [ - { - kind: 2, - languageId: 'python', - metadata: {}, - value: 'print("round-a")' - } - ], - metadata: { - deepnoteNotebookId: 'notebook-1', - deepnoteProjectId: 'project-123' - } - }; - let bytes = await serializer.serializeNotebook(dataNb1 as any, {} as any); - let meta = await serializer.deserializeNotebook(bytes, {} as any); - - assert.strictEqual(meta.metadata?.deepnoteNotebookId, 'notebook-2'); - - const dataNb2 = { - cells: [ - { - kind: 1, - languageId: 'markdown', - metadata: {}, - value: '# round-b' - } - ], - metadata: { - deepnoteNotebookId: 'notebook-2', - deepnoteProjectId: 'project-123' - } - }; - bytes = await serializer.serializeNotebook(dataNb2 as any, {} as any); - meta = await serializer.deserializeNotebook(bytes, {} as any); - - assert.strictEqual(meta.metadata?.deepnoteNotebookId, 'notebook-1'); - }); }); }); @@ -459,143 +354,6 @@ project: assert.strictEqual(result, undefined); }); - test('should return undefined when multiple notebooks are open and no stronger signal exists', () => { - const notebookA = { - then: undefined, - notebookType: 'deepnote', - metadata: { - deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'notebook-a' - }, - uri: {} as any, - version: 1, - isDirty: false, - isUntitled: false, - isClosed: false, - cellCount: 0, - cellAt: () => ({}) as any, - getCells: () => [], - save: async () => true - } as NotebookDocument; - const notebookB = { - then: undefined, - notebookType: 'deepnote', - metadata: { - deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'notebook-b' - }, - uri: {} as any, - version: 1, - isDirty: false, - isUntitled: false, - isClosed: false, - cellCount: 0, - cellAt: () => ({}) as any, - getCells: () => [], - save: async () => true - } as NotebookDocument; - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookA, notebookB]); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, undefined); - }); - - test('after serializing notebook-1 with both notebooks open should return sibling notebook-2', async () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-1'); - - const notebookDoc1 = createNotebookDoc('notebook-1'); - const notebookDoc2 = createNotebookDoc('notebook-2'); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookDoc1, notebookDoc2]); - - const mockNotebookData = { - cells: [ - { - kind: 2, - languageId: 'python', - metadata: {}, - value: 'print("findOpen")' - } - ], - metadata: { - deepnoteNotebookId: 'notebook-1', - deepnoteProjectId: 'project-123' - } - }; - - await serializer.serializeNotebook(mockNotebookData as any, {} as any); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, 'notebook-2'); - }); - - test('recent serialization hint is one-shot', async () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-1'); - - const notebookDoc1 = createNotebookDoc('notebook-1'); - const notebookDoc2 = createNotebookDoc('notebook-2'); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookDoc1, notebookDoc2]); - - const mockNotebookData = { - cells: [ - { - kind: 2, - languageId: 'python', - metadata: {}, - value: 'print("one-shot")' - } - ], - metadata: { - deepnoteNotebookId: 'notebook-1', - deepnoteProjectId: 'project-123' - } - }; - - await serializer.serializeNotebook(mockNotebookData as any, {} as any); - - assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'notebook-2'); - assert.strictEqual(serializer.findCurrentNotebookId('project-123'), undefined); - }); - - test('recent serialization hint expires after TTL', async () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-1'); - - const notebookDoc1 = createNotebookDoc('notebook-1'); - const notebookDoc2 = createNotebookDoc('notebook-2'); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebookDoc1, notebookDoc2]); - - const mockNotebookData = { - cells: [ - { - kind: 2, - languageId: 'python', - metadata: {}, - value: 'print("ttl")' - } - ], - metadata: { - deepnoteNotebookId: 'notebook-1', - deepnoteProjectId: 'project-123' - } - }; - - await serializer.serializeNotebook(mockNotebookData as any, {} as any); - - (serializer as any).recentSerializations.set('project-123', { - notebookId: 'notebook-1', - timestamp: Date.now() - 6000 - }); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, undefined); - }); - test('should return undefined for unknown project', () => { const result = serializer.findCurrentNotebookId('unknown-project'); From f4b79d3dfd389beae8602c5c1afda8dee808734a Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 7 Apr 2026 17:52:38 +0000 Subject: [PATCH 32/47] feat(deepnote): Add logging for pending notebook resolutions in DeepnoteNotebookManager - Introduced debug logging in `DeepnoteNotebookManager` to track valid pending notebook resolutions, enhancing traceability during resolution processes. - Simplified block iteration in `DeepnoteNotebookSerializer` by removing unnecessary null checks, improving code clarity and performance. --- src/notebooks/deepnote/deepnoteNotebookManager.ts | 6 ++++++ src/notebooks/deepnote/deepnoteSerializer.ts | 6 +++--- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.ts b/src/notebooks/deepnote/deepnoteNotebookManager.ts index b2e654e673..7d80ce3391 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.ts @@ -2,6 +2,7 @@ import { injectable } from 'inversify'; import { IDeepnoteNotebookManager, ProjectIntegration } from '../types'; import type { DeepnoteProject } from '../../platform/deepnote/deepnoteTypes'; +import { logger } from '../../platform/logging'; const pendingNotebookResolutionTtlMs = 60_000; @@ -152,6 +153,11 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { private getValidPendingNotebookResolutions(projectId: string): PendingNotebookResolution[] { const cutoffTime = Date.now() - pendingNotebookResolutionTtlMs; + const allPendingResolutions = this.pendingNotebookResolutions.get(projectId) ?? []; + logger.debug( + `DeepnoteNotebookManager: getValidPendingNotebookResolutions: projectId=${projectId}, allPendingResolutions=${allPendingResolutions.length}` + ); + logger.debug(JSON.stringify(allPendingResolutions, null, 2)); const pendingResolutions = (this.pendingNotebookResolutions.get(projectId) ?? []).filter( (resolution) => resolution.queuedAt >= cutoffTime ); diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 8bbdc0b796..ba3927e4f8 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -120,12 +120,12 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } // Log block IDs from source file - for (let i = 0; i < (selectedNotebook.blocks ?? []).length; i++) { - const block = selectedNotebook.blocks![i]; + for (let i = 0; i < selectedNotebook.blocks.length; i++) { + const block = selectedNotebook.blocks[i]; logger.trace(`DeserializeNotebook: block[${i}] id=${block.id} from source file`); } - let cells = this.converter.convertBlocksToCells(selectedNotebook.blocks ?? []); + let cells = this.converter.convertBlocksToCells(selectedNotebook.blocks); logger.debug(`DeepnoteSerializer: Converted ${cells.length} cells from notebook blocks`); From c83af5532f6ebb1a9c352fcbb09702e930352d80 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 7 Apr 2026 18:01:13 +0000 Subject: [PATCH 33/47] Reformat code --- src/notebooks/deepnote/deepnoteSerializer.unit.test.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index 139673a4b5..df2973ccdf 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -290,7 +290,6 @@ project: assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-2'); }); - }); }); From 40f6c6703ce80400bb8503378288333a384f7f1e Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 7 Apr 2026 18:33:32 +0000 Subject: [PATCH 34/47] Fix test --- .../deepnote/deepnoteFileChangeWatcher.unit.test.ts | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index fedca12907..be90fbfa65 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -510,12 +510,22 @@ project: }); test('should not suppress real changes after auto-save', async function () { - this.timeout(5000); + this.timeout(10_000); const uri = testFileUri('test.deepnote'); // First change: notebook has no cells, YAML has one cell -> different -> reload const notebook = createMockNotebook({ uri, cellCount: 0, cells: [] }); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + // Override save mock to fire onDidSaveNotebook (matching real VS Code behavior). + // The onDidSaveNotebookDocument handler calls markSelfWrite, producing the + // second self-write marker that corresponds to the serializer's save-triggered write. + when(mockedVSCodeNamespaces.workspace.save(anything())).thenCall((saveUri: Uri) => { + saveCount++; + onDidSaveNotebook.fire(notebook); + return Promise.resolve(saveUri); + }); + setupMockFs(validYaml); onDidChangeFile.fire(uri); From 46cbcda38c9a71a8ead921aab1efa7e6e265e3bc Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 9 Apr 2026 07:41:22 +0000 Subject: [PATCH 35/47] feat(deepnote): Enhance notebook deserialization and verification logic - Added a new event listener in `DeepnoteActivationService` to verify deserialized notebooks upon opening. - Improved `DeepnoteNotebookSerializer` to handle cases where no notebook ID is resolved, returning an empty state instead of throwing an error. - Updated unit tests to reflect changes in notebook deserialization behavior, ensuring accurate handling of scenarios with no available notebooks. --- .../deepnote/deepnoteActivationService.ts | 7 +- src/notebooks/deepnote/deepnoteSerializer.ts | 185 ++++++++++++++---- .../deepnote/deepnoteSerializer.unit.test.ts | 51 ++--- 3 files changed, 178 insertions(+), 65 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteActivationService.ts b/src/notebooks/deepnote/deepnoteActivationService.ts index 96c35288cd..edd3474a00 100644 --- a/src/notebooks/deepnote/deepnoteActivationService.ts +++ b/src/notebooks/deepnote/deepnoteActivationService.ts @@ -1,5 +1,5 @@ import { inject, injectable, optional } from 'inversify'; -import { commands, l10n, workspace, window, type Disposable, type NotebookDocumentContentOptions } from 'vscode'; +import { commands, l10n, window, workspace, type Disposable, type NotebookDocumentContentOptions } from 'vscode'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { IExtensionContext } from '../../platform/common/types'; @@ -51,6 +51,11 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic this.registerSerializer(); this.extensionContext.subscriptions.push(this.editProtection); + this.extensionContext.subscriptions.push( + workspace.onDidOpenNotebookDocument((doc) => { + void this.serializer.verifyDeserializedNotebook(doc); + }) + ); this.extensionContext.subscriptions.push( workspace.onDidChangeConfiguration((event) => { if (event.affectsConfiguration('deepnote.snapshots.enabled')) { diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index ba3927e4f8..531c645957 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -2,12 +2,15 @@ import type { DeepnoteBlock, DeepnoteFile, DeepnoteSnapshot } from '@deepnote/bl import { deserializeDeepnoteFile, isExecutableBlock, serializeDeepnoteSnapshot } from '@deepnote/blocks'; import { inject, injectable, optional } from 'inversify'; import { - TabInputNotebook, + CancellationTokenSource, l10n, - window, + NotebookEdit, + NotebookRange, workspace, + WorkspaceEdit, type CancellationToken, type NotebookData, + type NotebookDocument, type NotebookSerializer } from 'vscode'; @@ -54,6 +57,8 @@ function cloneWithoutCircularRefs(obj: T, seen = new WeakSet()): T { } } +const LAST_SERIALIZED_TTL_MS = 10_000; + /** * Serializer for converting between Deepnote YAML files and VS Code notebook format. * Handles reading/writing .deepnote files and manages project state persistence. @@ -61,6 +66,8 @@ function cloneWithoutCircularRefs(obj: T, seen = new WeakSet()): T { @injectable() export class DeepnoteNotebookSerializer implements NotebookSerializer { private converter = new DeepnoteDataConverter(); + private lastSerializedNotebookId: string | undefined; + private lastSerializedTimestamp = 0; constructor( @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, @@ -87,13 +94,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { throw new Error('Serialization cancelled'); } - // Initialize vega-lite for output conversion (lazy-loaded ESM module) - await this.converter.initialize(); - - if (token?.isCancellationRequested) { - throw new Error('Serialization cancelled'); - } - try { const contentString = new TextDecoder('utf-8').decode(content); const deepnoteFile = deserializeDeepnoteFile(contentString); @@ -107,18 +107,38 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { logger.debug(`DeepnoteSerializer: Project ID: ${projectId}, Selected notebook ID: ${resolvedNotebookId}`); + if (!resolvedNotebookId) { + logger.debug( + 'DeepnoteSerializer: No notebook ID resolved, returning empty state for post-open verification' + ); + + return { + cells: [], + metadata: { + deepnoteProjectId: projectId, + deepnoteProjectName: deepnoteFile.project.name, + deepnoteVersion: deepnoteFile.version + } + }; + } + if (deepnoteFile.project.notebooks.length === 0) { throw new Error('Deepnote project contains no notebooks.'); } - const selectedNotebook = resolvedNotebookId - ? deepnoteFile.project.notebooks.find((nb) => nb.id === resolvedNotebookId) - : this.findDefaultNotebook(deepnoteFile); + const selectedNotebook = deepnoteFile.project.notebooks.find((nb) => nb.id === resolvedNotebookId); if (!selectedNotebook) { throw new Error(l10n.t('No notebook selected or found')); } + // Initialize vega-lite for output conversion (lazy-loaded ESM module) + await this.converter.initialize(); + + if (token?.isCancellationRequested) { + throw new Error('Serialization cancelled'); + } + // Log block IDs from source file for (let i = 0; i < selectedNotebook.blocks.length; i++) { const block = selectedNotebook.blocks[i]; @@ -195,12 +215,15 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } /** - * Finds the notebook ID to deserialize for VS Code-initiated deserialization. + * Resolves the notebook ID for a deserialization call. + * * Priority: * 1. Pending resolution hint (explicit intent from explorer view) - * 2. Active tab URI query param (VS Code's ground truth for which tab is loading) - * @param projectId The project ID to find a notebook for - * @returns The notebook ID to deserialize, or undefined if none found + * 2. Already-open document metadata (re-deserialization after file change) + * — skips the notebook that was just serialized, since that's the one + * that triggered the file change, not the one being re-deserialized. + * 3. undefined — initial open; verifyDeserializedNotebook will resolve + * from the document URI after open. */ findCurrentNotebookId(projectId: string): string | undefined { const pendingNotebookId = this.notebookManager.consumePendingNotebookResolution(projectId); @@ -209,15 +232,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return pendingNotebookId; } - const activeTabNotebookId = this.findNotebookIdFromActiveTab(); - - if (activeTabNotebookId) { - logger.debug(`DeepnoteSerializer: Resolved notebook ID from active tab URI: ${activeTabNotebookId}`); - - return activeTabNotebookId; - } - - return undefined; + return this.findNotebookIdFromOpenDocuments(projectId); } /** @@ -228,6 +243,81 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return this.converter; } + /** + * Parses the file content and returns the ID of the default notebook + * (alphabetically first, excluding Init when other notebooks exist). + * Used by the post-open verification handler for direct file opens + * that have no `?notebook=` query param. + */ + resolveDefaultNotebookId(content: Uint8Array): string | undefined { + try { + const contentString = new TextDecoder('utf-8').decode(content); + const deepnoteFile = deserializeDeepnoteFile(contentString); + + return this.findDefaultNotebook(deepnoteFile)?.id; + } catch { + return undefined; + } + } + + /** + * Ensures an opened notebook document shows the correct notebook. + * + * The serializer returns an empty state when it cannot determine which + * notebook to display (no pending resolution). This method is called + * after the document is created, reads the real notebook ID from + * the URI `?notebook=` query param (or picks the default notebook + * for direct file opens), and patches the document in place. + */ + async verifyDeserializedNotebook(doc: NotebookDocument): Promise { + if (doc.notebookType !== 'deepnote') { + return; + } + + const expectedNotebookId = new URLSearchParams(doc.uri.query).get('notebook'); + const actualNotebookId = doc.metadata?.deepnoteNotebookId as string | undefined; + + if (actualNotebookId && (!expectedNotebookId || expectedNotebookId === actualNotebookId)) { + return; + } + + const cts = new CancellationTokenSource(); + + try { + const fileUri = doc.uri.with({ query: '', fragment: '' }); + const content = await workspace.fs.readFile(fileUri); + + const targetNotebookId = expectedNotebookId ?? this.resolveDefaultNotebookId(content); + + if (!targetNotebookId || targetNotebookId === actualNotebookId) { + return; + } + + logger.info( + `Notebook verification: resolving notebook ${targetNotebookId} (was ${actualNotebookId ?? 'empty'}).` + ); + + const correctData = await this.deserializeNotebook(content, cts.token, targetNotebookId); + + const edit = new WorkspaceEdit(); + + edit.set(doc.uri, [ + NotebookEdit.replaceCells(new NotebookRange(0, doc.cellCount), correctData.cells), + NotebookEdit.updateNotebookMetadata(correctData.metadata ?? {}) + ]); + + const applied = await workspace.applyEdit(edit); + + if (applied) { + await doc.save(); + } + } catch (error) { + logger.error('Failed to verify/correct notebook content', error); + } finally { + cts.dispose(); + } + } + /** * Serializes VS Code notebook data back to Deepnote YAML format. * Converts cells to blocks, updates project data, and generates YAML. @@ -271,6 +361,9 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { throw new Error('Cannot determine which notebook to save'); } + this.lastSerializedNotebookId = notebookId; + this.lastSerializedTimestamp = Date.now(); + logger.debug(`SerializeNotebook: Notebook ID: ${notebookId}`); const notebook = originalProject.project.notebooks.find((nb: { id: string }) => nb.id === notebookId); @@ -561,27 +654,39 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } /** - * Extracts the notebook ID from the active tab's URI query params. - * During session restore or tab open, VS Code may set the active tab - * before calling deserializeNotebook. The URI retains the - * `?notebook=` query param set when the tab was originally opened. + * Returns a notebook ID from an already-open document for re-deserialization. + * + * When VS Code re-reads a file (e.g. after save or external change), it + * calls deserializeNotebook again for each open document backed by that + * file. The document already exists in workspace.notebookDocuments with + * the correct deepnoteNotebookId in its metadata. + * + * Skips the notebook that was most recently serialized — that document + * triggered the file change and is not the one being re-deserialized. */ - private findNotebookIdFromActiveTab(): string | undefined { - const activeTab = window.tabGroups.activeTabGroup?.activeTab; + private findNotebookIdFromOpenDocuments(projectId: string): string | undefined { + const recentSerialize = + this.lastSerializedNotebookId && Date.now() - this.lastSerializedTimestamp < LAST_SERIALIZED_TTL_MS; - if (!activeTab || !(activeTab.input instanceof TabInputNotebook)) { - return undefined; - } + for (const doc of workspace.notebookDocuments) { + if (doc.notebookType !== 'deepnote' || doc.metadata?.deepnoteProjectId !== projectId) { + continue; + } - const tabInput = activeTab.input; + const notebookId = doc.metadata?.deepnoteNotebookId as string | undefined; - if (tabInput.notebookType !== 'deepnote') { - return undefined; - } + if (!notebookId) { + continue; + } - const notebookId = new URLSearchParams(tabInput.uri.query).get('notebook'); + if (recentSerialize && notebookId === this.lastSerializedNotebookId) { + continue; + } + + return notebookId; + } - return notebookId ?? undefined; + return undefined; } /** diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index df2973ccdf..02b6c369ce 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -129,7 +129,7 @@ project: ); }); - test('should throw error when no notebooks found', async () => { + test('should return empty state when no notebooks found and no ID resolved', async () => { const contentWithoutNotebooks = new TextEncoder().encode(` version: '1.0.0' metadata: @@ -141,10 +141,11 @@ project: settings: {} `); - await assert.isRejected( - serializer.deserializeNotebook(contentWithoutNotebooks, {} as any), - /no notebooks|notebooks.*must contain at least 1/i - ); + const result = await serializer.deserializeNotebook(contentWithoutNotebooks, {} as any); + + assert.deepStrictEqual(result.cells, []); + assert.strictEqual(result.metadata?.deepnoteProjectId, 'project-123'); + assert.strictEqual(result.metadata?.deepnoteNotebookId, undefined); }); test('should deserialize the specified notebook when notebookId is passed', async () => { @@ -750,7 +751,7 @@ project: }); suite('default notebook selection', () => { - test('should not select Init notebook when other notebooks are available', async () => { + test('should not select Init notebook when other notebooks are available', () => { const projectData: DeepnoteFile = { version: '1.0.0', metadata: { @@ -800,14 +801,12 @@ project: }; const content = projectToYaml(projectData); - const result = await serializer.deserializeNotebook(content, {} as any); + const result = serializer.resolveDefaultNotebookId(content); - // Should select the Main notebook, not the Init notebook - assert.strictEqual(result.metadata?.deepnoteNotebookId, 'main-notebook'); - assert.strictEqual(result.metadata?.deepnoteNotebookName, 'Main'); + assert.strictEqual(result, 'main-notebook'); }); - test('should select Init notebook when it is the only notebook', async () => { + test('should select Init notebook when it is the only notebook', () => { const projectData: DeepnoteFile = { version: '1.0.0', metadata: { @@ -841,14 +840,12 @@ project: }; const content = projectToYaml(projectData); - const result = await serializer.deserializeNotebook(content, {} as any); + const result = serializer.resolveDefaultNotebookId(content); - // Should select the Init notebook since it's the only one - assert.strictEqual(result.metadata?.deepnoteNotebookId, 'init-notebook'); - assert.strictEqual(result.metadata?.deepnoteNotebookName, 'Init'); + assert.strictEqual(result, 'init-notebook'); }); - test('should select alphabetically first notebook when no initNotebookId', async () => { + test('should select alphabetically first notebook when no initNotebookId', () => { const projectData: DeepnoteFile = { version: '1.0.0', metadata: { @@ -913,14 +910,12 @@ project: }; const content = projectToYaml(projectData); - const result = await serializer.deserializeNotebook(content, {} as any); + const result = serializer.resolveDefaultNotebookId(content); - // Should select the alphabetically first notebook - assert.strictEqual(result.metadata?.deepnoteNotebookId, 'alpha-notebook'); - assert.strictEqual(result.metadata?.deepnoteNotebookName, 'Alpha Notebook'); + assert.strictEqual(result, 'alpha-notebook'); }); - test('should sort Init notebook last when multiple notebooks exist', async () => { + test('should sort Init notebook last when multiple notebooks exist', () => { const projectData: DeepnoteFile = { version: '1.0.0', metadata: { @@ -986,11 +981,19 @@ project: }; const content = projectToYaml(projectData); - const result = await serializer.deserializeNotebook(content, {} as any); + const result = serializer.resolveDefaultNotebookId(content); // Should select Alpha, not Init even though "Init" comes before "Alpha" alphabetically when in upper case - assert.strictEqual(result.metadata?.deepnoteNotebookId, 'alpha-notebook'); - assert.strictEqual(result.metadata?.deepnoteNotebookName, 'Alpha'); + assert.strictEqual(result, 'alpha-notebook'); + }); + + test('should return empty state from deserializeNotebook when no notebook ID is resolved', async () => { + const content = projectToYaml(mockProject); + const result = await serializer.deserializeNotebook(content, {} as any); + + assert.deepStrictEqual(result.cells, []); + assert.strictEqual(result.metadata?.deepnoteProjectId, 'project-123'); + assert.strictEqual(result.metadata?.deepnoteNotebookId, undefined); }); }); From 477b79055c705ad7327948fa3c9e876a8a72ae5a Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 13 Apr 2026 06:16:12 +0000 Subject: [PATCH 36/47] feat(deepnote-file): Migration to single notebook deepnote file --- .../deepnote/deepnoteActivationService.ts | 32 +- .../deepnote/deepnoteAutoSplitter.ts | 291 ++++++++++++ .../deepnote/deepnoteExplorerView.ts | 64 +-- .../deepnoteExplorerView.unit.test.ts | 16 +- .../deepnote/deepnoteFileChangeWatcher.ts | 25 +- .../deepnoteFileChangeWatcher.unit.test.ts | 14 +- .../deepnote/deepnoteNotebookManager.ts | 128 +---- .../deepnoteNotebookManager.unit.test.ts | 122 ++--- src/notebooks/deepnote/deepnoteSerializer.ts | 210 +-------- .../deepnote/deepnoteSerializer.unit.test.ts | 446 ++---------------- .../deepnote/deepnoteTreeDataProvider.ts | 282 +++++------ .../deepnoteTreeDataProvider.unit.test.ts | 49 +- src/notebooks/deepnote/deepnoteTreeItem.ts | 68 ++- .../deepnote/deepnoteTreeItem.unit.test.ts | 45 +- .../deepnote/snapshots/snapshotFiles.ts | 41 +- .../deepnote/snapshots/snapshotService.ts | 58 ++- src/notebooks/types.ts | 5 +- 17 files changed, 791 insertions(+), 1105 deletions(-) create mode 100644 src/notebooks/deepnote/deepnoteAutoSplitter.ts diff --git a/src/notebooks/deepnote/deepnoteActivationService.ts b/src/notebooks/deepnote/deepnoteActivationService.ts index edd3474a00..b81db6a7ba 100644 --- a/src/notebooks/deepnote/deepnoteActivationService.ts +++ b/src/notebooks/deepnote/deepnoteActivationService.ts @@ -1,3 +1,4 @@ +import { deserializeDeepnoteFile } from '@deepnote/blocks'; import { inject, injectable, optional } from 'inversify'; import { commands, l10n, window, workspace, type Disposable, type NotebookDocumentContentOptions } from 'vscode'; @@ -5,10 +6,11 @@ import { IExtensionSyncActivationService } from '../../platform/activation/types import { IExtensionContext } from '../../platform/common/types'; import { ILogger } from '../../platform/logging/types'; import { IDeepnoteNotebookManager } from '../types'; -import { DeepnoteNotebookSerializer } from './deepnoteSerializer'; +import { DeepnoteAutoSplitter } from './deepnoteAutoSplitter'; import { DeepnoteExplorerView } from './deepnoteExplorerView'; -import { IIntegrationManager } from './integrations/types'; import { DeepnoteInputBlockEditProtection } from './deepnoteInputBlockEditProtection'; +import { DeepnoteNotebookSerializer } from './deepnoteSerializer'; +import { IIntegrationManager } from './integrations/types'; import { SnapshotService } from './snapshots/snapshotService'; /** @@ -17,6 +19,8 @@ import { SnapshotService } from './snapshots/snapshotService'; */ @injectable() export class DeepnoteActivationService implements IExtensionSyncActivationService { + private autoSplitter: DeepnoteAutoSplitter; + private editProtection: DeepnoteInputBlockEditProtection; private explorerView: DeepnoteExplorerView; @@ -44,8 +48,9 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic * Called during extension activation to set up Deepnote integration. */ public activate() { + this.autoSplitter = new DeepnoteAutoSplitter(); this.serializer = new DeepnoteNotebookSerializer(this.notebookManager, this.snapshotService); - this.explorerView = new DeepnoteExplorerView(this.extensionContext, this.notebookManager, this.logger); + this.explorerView = new DeepnoteExplorerView(this.extensionContext, this.logger); this.editProtection = new DeepnoteInputBlockEditProtection(this.logger); this.snapshotsEnabled = this.isSnapshotsEnabled(); @@ -53,7 +58,7 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic this.extensionContext.subscriptions.push(this.editProtection); this.extensionContext.subscriptions.push( workspace.onDidOpenNotebookDocument((doc) => { - void this.serializer.verifyDeserializedNotebook(doc); + void this.checkAndSplitIfNeeded(doc); }) ); this.extensionContext.subscriptions.push( @@ -75,6 +80,25 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic this.integrationManager.activate(); } + private async checkAndSplitIfNeeded(doc: import('vscode').NotebookDocument): Promise { + if (doc.notebookType !== 'deepnote') { + return; + } + + try { + const fileUri = doc.uri; + const content = await workspace.fs.readFile(fileUri); + const deepnoteFile = deserializeDeepnoteFile(new TextDecoder().decode(content)); + const result = await this.autoSplitter.splitIfNeeded(fileUri, deepnoteFile); + + if (result.wasSplit) { + this.explorerView.refreshTree(); + } + } catch (error) { + this.logger.error('Failed to check/split notebook file', error); + } + } + private isSnapshotsEnabled(): boolean { if (this.snapshotService) { return this.snapshotService.isSnapshotsEnabled(); diff --git a/src/notebooks/deepnote/deepnoteAutoSplitter.ts b/src/notebooks/deepnote/deepnoteAutoSplitter.ts new file mode 100644 index 0000000000..31b88b6cd0 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteAutoSplitter.ts @@ -0,0 +1,291 @@ +import { serializeDeepnoteFile, type DeepnoteFile } from '@deepnote/blocks'; +import { l10n, RelativePattern, Uri, window, workspace } from 'vscode'; + +import { computeHash } from '../../platform/common/crypto'; +import { logger } from '../../platform/logging'; +import { slugifyProjectName } from './snapshots/snapshotFiles'; + +/** + * Splits multi-notebook .deepnote files into separate files (one notebook per file). + * All split files share the same project ID and metadata. + */ +export class DeepnoteAutoSplitter { + /** + * Checks if a file has more than one non-init notebook and splits it if so. + * The first non-init notebook (by array order) stays in the original file. + * Each extra notebook gets its own file with the same project metadata. + * + * @returns Info about whether a split happened and which new files were created + */ + async splitIfNeeded(fileUri: Uri, deepnoteFile: DeepnoteFile): Promise<{ wasSplit: boolean; newFiles: Uri[] }> { + const initNotebookId = deepnoteFile.project.initNotebookId; + + const nonInitNotebooks = deepnoteFile.project.notebooks.filter((nb) => nb.id !== initNotebookId); + + if (nonInitNotebooks.length <= 1) { + return { wasSplit: false, newFiles: [] }; + } + + const initNotebook = initNotebookId + ? deepnoteFile.project.notebooks.find((nb) => nb.id === initNotebookId) + : undefined; + + // First non-init notebook stays in original file + const primaryNotebook = nonInitNotebooks[0]; + const extraNotebooks = nonInitNotebooks.slice(1); + + logger.info( + `[AutoSplitter] Splitting ${nonInitNotebooks.length} notebooks from project ${deepnoteFile.project.id}` + ); + + const parentDir = Uri.joinPath(fileUri, '..'); + const originalStem = this.getFileStem(fileUri); + const newFiles: Uri[] = []; + + // Create a new file for each extra notebook + for (const notebook of extraNotebooks) { + const notebookSlug = this.slugifyNotebookName(notebook.name); + const newFileName = `${originalStem}_${notebookSlug}.deepnote`; + const newFileUri = Uri.joinPath(parentDir, newFileName); + + const notebooks = initNotebook ? [structuredClone(initNotebook), notebook] : [notebook]; + + const newProject: DeepnoteFile = { + metadata: deepnoteFile.metadata + ? structuredClone(deepnoteFile.metadata) + : { createdAt: new Date().toISOString() }, + project: { + ...deepnoteFile.project, + notebooks + }, + version: deepnoteFile.version + }; + + if (initNotebook && initNotebookId) { + newProject.project.initNotebookId = initNotebookId; + } + + // Recompute snapshotHash for the new file + (newProject.metadata as Record).snapshotHash = await this.computeSnapshotHash(newProject); + + const yaml = serializeDeepnoteFile(newProject); + await workspace.fs.writeFile(newFileUri, new TextEncoder().encode(yaml)); + + newFiles.push(newFileUri); + + logger.info(`[AutoSplitter] Created ${newFileName} for notebook "${notebook.name}"`); + } + + // Update original file to keep only primary notebook (+ init) + const originalNotebooks = initNotebook ? [structuredClone(initNotebook), primaryNotebook] : [primaryNotebook]; + deepnoteFile.project.notebooks = originalNotebooks; + + if (deepnoteFile.metadata) { + (deepnoteFile.metadata as Record).snapshotHash = await this.computeSnapshotHash( + deepnoteFile + ); + } + + const updatedYaml = serializeDeepnoteFile(deepnoteFile); + await workspace.fs.writeFile(fileUri, new TextEncoder().encode(updatedYaml)); + + // Split snapshot files too + await this.splitSnapshots( + fileUri, + deepnoteFile.project.id, + deepnoteFile.project.name, + primaryNotebook.id, + extraNotebooks.map((nb) => nb.id) + ); + + // Notify the user + const fileNames = newFiles.map((f) => f.path.split('/').pop()).join(', '); + + void window.showInformationMessage( + l10n.t( + 'This project had {0} notebooks. They have been split into separate files: {1}', + nonInitNotebooks.length, + fileNames + ) + ); + + return { wasSplit: true, newFiles }; + } + + /** + * Splits existing snapshot files so each notebook gets its own snapshot. + * Old format: {slug}_{projectId}_{variant}.snapshot.deepnote + * New format: {slug}_{projectId}_{notebookId}_{variant}.snapshot.deepnote + */ + private async splitSnapshots( + _projectFileUri: Uri, + projectId: string, + projectName: string, + primaryNotebookId: string, + extraNotebookIds: string[] + ): Promise { + const workspaceFolders = workspace.workspaceFolders; + + if (!workspaceFolders || workspaceFolders.length === 0) { + return; + } + + const snapshotGlob = `**/snapshots/*_${projectId}_*.snapshot.deepnote`; + let allSnapshotFiles: Uri[] = []; + + for (const folder of workspaceFolders) { + const pattern = new RelativePattern(folder, snapshotGlob); + const files = await workspace.findFiles(pattern, null, 100); + allSnapshotFiles = allSnapshotFiles.concat(files); + } + + if (allSnapshotFiles.length === 0) { + logger.debug(`[AutoSplitter] No snapshots found for project ${projectId}`); + return; + } + + let slug: string; + + try { + slug = slugifyProjectName(projectName); + } catch { + logger.warn(`[AutoSplitter] Cannot slugify project name, skipping snapshot split`); + return; + } + + const allNotebookIds = [primaryNotebookId, ...extraNotebookIds]; + + for (const snapshotUri of allSnapshotFiles) { + try { + await this.splitSingleSnapshot(snapshotUri, slug, projectId, allNotebookIds); + } catch (error) { + logger.warn(`[AutoSplitter] Failed to split snapshot ${snapshotUri.path}`, error); + } + } + } + + private async splitSingleSnapshot( + snapshotUri: Uri, + slug: string, + projectId: string, + notebookIds: string[] + ): Promise { + const content = await workspace.fs.readFile(snapshotUri); + const { deserializeDeepnoteFile: parseFile } = await import('@deepnote/blocks'); + const snapshotData = parseFile(new TextDecoder().decode(content)); + + const snapshotDir = Uri.joinPath(snapshotUri, '..'); + + // Extract variant from existing filename + const basename = snapshotUri.path.split('/').pop() ?? ''; + const variant = this.extractVariantFromSnapshotFilename(basename, projectId); + + if (!variant) { + logger.debug(`[AutoSplitter] Could not extract variant from ${basename}, skipping`); + return; + } + + // Create per-notebook snapshot files + for (const notebookId of notebookIds) { + const notebookData = structuredClone(snapshotData); + + // Keep only this notebook (and init) + notebookData.project.notebooks = notebookData.project.notebooks.filter( + (nb) => + nb.id === notebookId || + (notebookData.project.initNotebookId && nb.id === notebookData.project.initNotebookId) + ); + + // Recompute hash + if (notebookData.metadata) { + (notebookData.metadata as Record).snapshotHash = await this.computeSnapshotHash( + notebookData + ); + } + + const newFilename = `${slug}_${projectId}_${notebookId}_${variant}.snapshot.deepnote`; + const newUri = Uri.joinPath(snapshotDir, newFilename); + const yaml = serializeDeepnoteFile(notebookData); + + await workspace.fs.writeFile(newUri, new TextEncoder().encode(yaml)); + + logger.debug(`[AutoSplitter] Created notebook snapshot: ${newFilename}`); + } + + // Delete the old snapshot file (it's been replaced by per-notebook files) + try { + await workspace.fs.delete(snapshotUri); + logger.debug(`[AutoSplitter] Deleted old snapshot: ${basename}`); + } catch { + logger.warn(`[AutoSplitter] Failed to delete old snapshot: ${basename}`); + } + } + + /** + * Extracts the variant portion from a snapshot filename. + * Old format: {slug}_{projectId}_{variant}.snapshot.deepnote + */ + private extractVariantFromSnapshotFilename(basename: string, projectId: string): string | undefined { + const suffix = '.snapshot.deepnote'; + + if (!basename.endsWith(suffix)) { + return undefined; + } + + const withoutSuffix = basename.slice(0, -suffix.length); + const projectIdIndex = withoutSuffix.indexOf(`_${projectId}_`); + + if (projectIdIndex === -1) { + return undefined; + } + + return withoutSuffix.slice(projectIdIndex + 1 + projectId.length + 1); + } + + /** + * Computes snapshotHash using the same algorithm as DeepnoteNotebookSerializer. + */ + private async computeSnapshotHash(project: DeepnoteFile): Promise { + const contentHashes: string[] = []; + + for (const notebook of project.project.notebooks) { + for (const block of notebook.blocks ?? []) { + if (block.contentHash) { + contentHashes.push(block.contentHash); + } + } + } + + contentHashes.sort(); + + const hashInput = { + contentHashes, + environmentHash: project.environment?.hash ?? null, + integrations: (project.project.integrations ?? []) + .map((i: { id: string; name: string; type: string }) => ({ id: i.id, name: i.name, type: i.type })) + .sort((a: { id: string }, b: { id: string }) => a.id.localeCompare(b.id)), + version: project.version + }; + + const hashData = JSON.stringify(hashInput); + const hash = await computeHash(hashData, 'SHA-256'); + + return `sha256:${hash}`; + } + + private getFileStem(uri: Uri): string { + const basename = uri.path.split('/').pop() ?? ''; + const dotIndex = basename.indexOf('.'); + + return dotIndex > 0 ? basename.slice(0, dotIndex) : basename; + } + + private slugifyNotebookName(name: string): string { + try { + return slugifyProjectName(name); + } catch { + // Fallback for names that produce empty slugs + return 'notebook'; + } + } +} diff --git a/src/notebooks/deepnote/deepnoteExplorerView.ts b/src/notebooks/deepnote/deepnoteExplorerView.ts index 9c41083e5f..a8fc8173a5 100644 --- a/src/notebooks/deepnote/deepnoteExplorerView.ts +++ b/src/notebooks/deepnote/deepnoteExplorerView.ts @@ -4,7 +4,7 @@ import { serializeDeepnoteFile, type DeepnoteBlock, type DeepnoteFile } from '@d import { convertDeepnoteToJupyterNotebooks, convertIpynbFilesToDeepnoteFile } from '@deepnote/convert'; import { IExtensionContext } from '../../platform/common/types'; -import { IDeepnoteNotebookManager } from '../types'; + import { DeepnoteTreeDataProvider } from './deepnoteTreeDataProvider'; import { type DeepnoteTreeItem, DeepnoteTreeItemType, type DeepnoteTreeItemContext } from './deepnoteTreeItem'; import { uuidUtils } from '../../platform/common/uuid'; @@ -24,7 +24,6 @@ export class DeepnoteExplorerView { constructor( @inject(IExtensionContext) private readonly extensionContext: IExtensionContext, - @inject(IDeepnoteNotebookManager) private readonly manager: IDeepnoteNotebookManager, @inject(ILogger) logger: ILogger ) { this.treeDataProvider = new DeepnoteTreeDataProvider(logger); @@ -77,7 +76,7 @@ export class DeepnoteExplorerView { projectData.project.notebooks.push(newNotebook); // Save and open the new notebook - await this.saveProjectAndOpenNotebook(fileUri, projectData, newNotebook.id); + await this.saveProjectAndOpenNotebook(fileUri, projectData); return { id: newNotebook.id, name: notebookName }; } @@ -259,10 +258,8 @@ export class DeepnoteExplorerView { await this.treeDataProvider.refreshNotebook(treeItem.context.projectId); - // Optionally open the duplicated notebook - this.registerNotebookOpenIntent(treeItem.context.projectId, newNotebook.id); - const notebookUri = fileUri.with({ query: `notebook=${newNotebook.id}` }); - const document = await workspace.openNotebookDocument(notebookUri); + // Open the duplicated notebook + const document = await workspace.openNotebookDocument(fileUri); await window.showNotebookDocument(document, { preserveFocus: false, preview: false @@ -488,11 +485,7 @@ export class DeepnoteExplorerView { * @param projectData The project data to save * @param notebookId The notebook ID to open */ - private async saveProjectAndOpenNotebook( - fileUri: Uri, - projectData: DeepnoteFile, - notebookId: string - ): Promise { + private async saveProjectAndOpenNotebook(fileUri: Uri, projectData: DeepnoteFile): Promise { // Update metadata timestamp if (!projectData.metadata) { projectData.metadata = { createdAt: new Date().toISOString() }; @@ -504,21 +497,19 @@ export class DeepnoteExplorerView { const encoder = new TextEncoder(); await workspace.fs.writeFile(fileUri, encoder.encode(updatedYaml)); - // Refresh the tree view - use granular refresh for notebooks + // Refresh the tree view await this.treeDataProvider.refreshNotebook(projectData.project.id); - // Open the new notebook - this.registerNotebookOpenIntent(projectData.project.id, notebookId); - const notebookUri = fileUri.with({ query: `notebook=${notebookId}` }); - const document = await workspace.openNotebookDocument(notebookUri); + // Open the notebook + const document = await workspace.openNotebookDocument(fileUri); await window.showNotebookDocument(document, { preserveFocus: false, preview: false }); } - private registerNotebookOpenIntent(projectId: string, notebookId: string): void { - this.manager.queueNotebookResolution(projectId, notebookId); + public refreshTree(): void { + this.treeDataProvider.refresh(); } private refreshExplorer(): void { @@ -526,29 +517,12 @@ export class DeepnoteExplorerView { } private async openNotebook(context: DeepnoteTreeItemContext): Promise { - console.log(`Opening notebook: ${context.notebookId} in project: ${context.projectId}.`); - - if (!context.notebookId) { - await window.showWarningMessage(l10n.t('Cannot open: missing notebook id.')); - - return; - } + console.log(`Opening notebook in project: ${context.projectId}.`); try { - // Create a unique URI by adding the notebook ID as a query parameter - // This ensures VS Code treats each notebook as a separate document - const fileUri = Uri.file(context.filePath).with({ query: `notebook=${context.notebookId}` }); - - console.log(`Selecting notebook in manager.`); - - this.registerNotebookOpenIntent(context.projectId, context.notebookId); - - console.log(`Opening notebook document.`, fileUri); - + const fileUri = Uri.file(context.filePath); const document = await workspace.openNotebookDocument(fileUri); - console.log(`Showing notebook document.`); - await window.showNotebookDocument(document, { preview: false, preserveFocus: false @@ -594,7 +568,7 @@ export class DeepnoteExplorerView { // Try to reveal the notebook in the explorer try { - const treeItem = await this.treeDataProvider.findTreeItem(projectId, notebookId); + const treeItem = await this.treeDataProvider.findTreeItem(projectId); if (treeItem) { await this.treeView.reveal(treeItem, { select: true, focus: true, expand: true }); @@ -705,10 +679,7 @@ export class DeepnoteExplorerView { this.treeDataProvider.refresh(); - this.registerNotebookOpenIntent(projectId, notebookId); - - const notebookUri = fileUri.with({ query: `notebook=${notebookId}` }); - const document = await workspace.openNotebookDocument(notebookUri); + const document = await workspace.openNotebookDocument(fileUri); await window.showNotebookDocument(document, { preserveFocus: false, @@ -729,12 +700,7 @@ export class DeepnoteExplorerView { } const document = activeEditor.notebook; - - // Get the file URI (strip query params if present) - let fileUri = document.uri; - if (fileUri.query) { - fileUri = fileUri.with({ query: '' }); - } + const fileUri = document.uri; try { // Use shared helper to create and add notebook diff --git a/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts b/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts index 5e5e2de826..a1f12f9604 100644 --- a/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts @@ -6,7 +6,7 @@ import { Uri, workspace } from 'vscode'; import { stringify as yamlStringify } from 'yaml'; import { DeepnoteExplorerView } from './deepnoteExplorerView'; -import { DeepnoteNotebookManager } from './deepnoteNotebookManager'; + import { DeepnoteTreeItem, DeepnoteTreeItemType, type DeepnoteTreeItemContext } from './deepnoteTreeItem'; import type { IExtensionContext } from '../../platform/common/types'; import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; @@ -42,7 +42,6 @@ function createUuidMock(uuids: string[]): sinon.SinonStub { suite('DeepnoteExplorerView', () => { let explorerView: DeepnoteExplorerView; let mockExtensionContext: IExtensionContext; - let manager: DeepnoteNotebookManager; let mockLogger: ILogger; setup(() => { @@ -50,9 +49,8 @@ suite('DeepnoteExplorerView', () => { subscriptions: [] } as any; - manager = new DeepnoteNotebookManager(); mockLogger = createMockLogger(); - explorerView = new DeepnoteExplorerView(mockExtensionContext, manager, mockLogger); + explorerView = new DeepnoteExplorerView(mockExtensionContext, mockLogger); }); suite('constructor', () => { @@ -186,12 +184,10 @@ suite('DeepnoteExplorerView', () => { const context1 = { subscriptions: [] } as any; const context2 = { subscriptions: [] } as any; - const manager1 = new DeepnoteNotebookManager(); - const manager2 = new DeepnoteNotebookManager(); const logger1 = createMockLogger(); const logger2 = createMockLogger(); - const view1 = new DeepnoteExplorerView(context1, manager1, logger1); - const view2 = new DeepnoteExplorerView(context2, manager2, logger2); + const view1 = new DeepnoteExplorerView(context1, logger1); + const view2 = new DeepnoteExplorerView(context2, logger2); // Verify each view has its own context assert.strictEqual((view1 as any).extensionContext, context1); @@ -219,7 +215,6 @@ suite('DeepnoteExplorerView', () => { suite('DeepnoteExplorerView - Empty State Commands', () => { let explorerView: DeepnoteExplorerView; let mockContext: IExtensionContext; - let mockManager: DeepnoteNotebookManager; let sandbox: sinon.SinonSandbox; let uuidStubs: sinon.SinonStub[] = []; @@ -232,9 +227,8 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { subscriptions: [] } as unknown as IExtensionContext; - mockManager = new DeepnoteNotebookManager(); const mockLogger = createMockLogger(); - explorerView = new DeepnoteExplorerView(mockContext, mockManager, mockLogger); + explorerView = new DeepnoteExplorerView(mockContext, mockLogger); }); teardown(() => { diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index 8e54997f2e..76fa5158b4 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -20,7 +20,11 @@ import { IControllerRegistration } from '../controllers/types'; import { IDeepnoteNotebookManager } from '../types'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; import { DeepnoteNotebookSerializer } from './deepnoteSerializer'; -import { extractProjectIdFromSnapshotUri, isSnapshotFile } from './snapshots/snapshotFiles'; +import { + extractNotebookIdFromSnapshotUri, + extractProjectIdFromSnapshotUri, + isSnapshotFile +} from './snapshots/snapshotFiles'; import { SnapshotService } from './snapshots/snapshotService'; const debounceTimeInMilliseconds = 500; @@ -260,12 +264,15 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } /** - * Enqueue a snapshot-output-update for all notebooks matching this project. + * Enqueue a snapshot-output-update for all notebooks matching this project (and notebook if specified). * Does NOT replace a pending main-file-sync. */ - private enqueueSnapshotOutputUpdate(projectId: string): void { + private enqueueSnapshotOutputUpdate(projectId: string, notebookId?: string): void { const affectedNotebooks = workspace.notebookDocuments.filter( - (doc) => doc.notebookType === 'deepnote' && doc.metadata?.deepnoteProjectId === projectId + (doc) => + doc.notebookType === 'deepnote' && + doc.metadata?.deepnoteProjectId === projectId && + (!notebookId || doc.metadata?.deepnoteNotebookId === notebookId) ); for (const notebook of affectedNotebooks) { @@ -292,12 +299,10 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic return; } - const targetNotebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; - const tokenSource = new CancellationTokenSource(); let newData; try { - newData = await this.serializer.deserializeNotebook(content, tokenSource.token, targetNotebookId); + newData = await this.serializer.deserializeNotebook(content, tokenSource.token); } catch (error) { logger.warn(`[FileChangeWatcher] Failed to parse changed file: ${fileUri.path}`, error); return; @@ -619,7 +624,9 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic return; } - const key = `snapshot:${projectId}`; + const notebookId = extractNotebookIdFromSnapshotUri(uri); + + const key = notebookId ? `snapshot:${projectId}:${notebookId}` : `snapshot:${projectId}`; const existing = this.debounceTimers.get(key); if (existing) { clearTimeout(existing); @@ -629,7 +636,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic key, setTimeout(() => { this.debounceTimers.delete(key); - this.enqueueSnapshotOutputUpdate(projectId); + this.enqueueSnapshotOutputUpdate(projectId, notebookId); }, debounceTimeInMilliseconds) ); } diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index be90fbfa65..7877b95d04 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -162,9 +162,7 @@ suite('DeepnoteFileChangeWatcher', () => { mockDisposables = []; mockedNotebookManager = mock(); - when(mockedNotebookManager.consumePendingNotebookResolution(anything())).thenReturn(undefined); when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(validProject); - when(mockedNotebookManager.queueNotebookResolution(anything(), anything())).thenReturn(); when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); mockNotebookManager = instance(mockedNotebookManager); @@ -1557,9 +1555,7 @@ project: } as any); const mockedManagerEx = mock(); - when(mockedManagerEx.consumePendingNotebookResolution(anything())).thenReturn(undefined); when(mockedManagerEx.getOriginalProject(anything())).thenReturn(validProject); - when(mockedManagerEx.queueNotebookResolution(anything(), anything())).thenReturn(); when(mockedManagerEx.updateOriginalProject(anything(), anything())).thenReturn(); const exMdWatcher = new DeepnoteFileChangeWatcher( @@ -1618,9 +1614,7 @@ project: interactionCaptures = []; reset(mockedNotebookManager); - when(mockedNotebookManager.consumePendingNotebookResolution(anything())).thenReturn(undefined); when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(validProject); - when(mockedNotebookManager.queueNotebookResolution(anything(), anything())).thenReturn(); when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); resetCalls(mockedNotebookManager); @@ -1964,7 +1958,8 @@ project: ); }); - test('multi-notebook: snapshot outputs then external YAML update keeps per-notebook sources', async function () { + // Multi-notebook test removed — multi-notebook support has been replaced by auto-splitting into separate files + test.skip('multi-notebook: snapshot outputs then external YAML update keeps per-notebook sources', async function () { this.timeout(12_000); when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); @@ -2157,14 +2152,13 @@ project: }); }); - suite('multi-notebook file sync', () => { + // Multi-notebook file sync tests removed — multi-notebook support has been replaced by auto-splitting into separate files + suite.skip('multi-notebook file sync', () => { let workspaceSetCaptures: NotebookEditCapture[] = []; setup(() => { reset(mockedNotebookManager); - when(mockedNotebookManager.consumePendingNotebookResolution(anything())).thenReturn(undefined); when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); - when(mockedNotebookManager.queueNotebookResolution(anything(), anything())).thenReturn(); when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); resetCalls(mockedNotebookManager); workspaceSetCaptures = []; diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.ts b/src/notebooks/deepnote/deepnoteNotebookManager.ts index 7d80ce3391..26dad36a01 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.ts @@ -2,53 +2,16 @@ import { injectable } from 'inversify'; import { IDeepnoteNotebookManager, ProjectIntegration } from '../types'; import type { DeepnoteProject } from '../../platform/deepnote/deepnoteTypes'; -import { logger } from '../../platform/logging'; - -const pendingNotebookResolutionTtlMs = 60_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. + * Centralized manager for tracking Deepnote project state. + * Manages per-project data caching and init notebook tracking. */ @injectable() export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { - private readonly currentNotebookId = new Map(); private readonly originalProjects = new Map(); - private readonly pendingNotebookResolutions = new Map(); private readonly projectsWithInitNotebookRun = new Set(); - /** - * 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 - * @returns Current notebook ID or undefined if not set - */ - getCurrentNotebookId(projectId: string): string | undefined { - return this.currentNotebookId.get(projectId); - } - /** * Retrieves the original project data for a given project ID. * @param projectId Project identifier @@ -59,44 +22,36 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { } /** - * 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 + * Checks if the init notebook has already been run for a project. + * @param projectId Project identifier + * @returns True if init notebook has been run, false otherwise */ - queueNotebookResolution(projectId: string, notebookId: string): void { - const pendingResolutions = this.getValidPendingNotebookResolutions(projectId); - - pendingResolutions.push({ - notebookId, - queuedAt: Date.now() - }); + hasInitNotebookBeenRun(projectId: string): boolean { + return this.projectsWithInitNotebookRun.has(projectId); + } - this.pendingNotebookResolutions.set(projectId, pendingResolutions); + /** + * Marks the init notebook as having been run for a project. + * @param projectId Project identifier + */ + markInitNotebookAsRun(projectId: string): void { + this.projectsWithInitNotebookRun.add(projectId); } /** - * Stores the original project data and sets the initial current notebook. - * This is used during deserialization to cache project data and track the active notebook. + * Stores the original project data. + * This is used during deserialization to cache project data. * @param projectId Project identifier * @param project Original project data to store - * @param notebookId Initial notebook ID to set as current */ - storeOriginalProject(projectId: string, project: DeepnoteProject, notebookId: string): void { - // Deep clone to prevent mutations from affecting stored state - // This is critical for multi-notebook projects where multiple notebooks - // share the same stored project reference - // Using structuredClone to handle circular references (e.g., in output metadata) + storeOriginalProject(projectId: string, project: DeepnoteProject): void { const clonedProject = structuredClone(project); - this.originalProjects.set(projectId, clonedProject); - this.currentNotebookId.set(projectId, notebookId); } /** - * Updates the stored project data without changing the current notebook selection. - * Used during serialization where we need to cache the updated project state - * but must not alter notebook routing for other open notebooks. + * Updates the stored project data. + * Used during serialization where we need to cache the updated project state. * @param projectId Project identifier * @param project Updated project data to store */ @@ -122,51 +77,8 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { const updatedProject = structuredClone(project); updatedProject.project.integrations = integrations; - - const currentNotebookId = this.currentNotebookId.get(projectId); - - if (currentNotebookId) { - this.storeOriginalProject(projectId, updatedProject, currentNotebookId); - } else { - this.originalProjects.set(projectId, updatedProject); - } + this.originalProjects.set(projectId, updatedProject); return true; } - - /** - * Checks if the init notebook has already been run for a project. - * @param projectId Project identifier - * @returns True if init notebook has been run, false otherwise - */ - hasInitNotebookBeenRun(projectId: string): boolean { - return this.projectsWithInitNotebookRun.has(projectId); - } - - /** - * Marks the init notebook as having been run for a project. - * @param projectId Project identifier - */ - markInitNotebookAsRun(projectId: string): void { - this.projectsWithInitNotebookRun.add(projectId); - } - - private getValidPendingNotebookResolutions(projectId: string): PendingNotebookResolution[] { - const cutoffTime = Date.now() - pendingNotebookResolutionTtlMs; - const allPendingResolutions = this.pendingNotebookResolutions.get(projectId) ?? []; - logger.debug( - `DeepnoteNotebookManager: getValidPendingNotebookResolutions: projectId=${projectId}, allPendingResolutions=${allPendingResolutions.length}` - ); - logger.debug(JSON.stringify(allPendingResolutions, null, 2)); - 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 index 8a7de52200..7d3b226843 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts @@ -25,22 +25,6 @@ suite('DeepnoteNotebookManager', () => { manager = new DeepnoteNotebookManager(); }); - suite('getCurrentNotebookId', () => { - test('should return undefined for unknown project', () => { - const result = manager.getCurrentNotebookId('unknown-project'); - - assert.strictEqual(result, undefined); - }); - - test('should return notebook ID after storing project', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); - - const result = manager.getCurrentNotebookId('project-123'); - - assert.strictEqual(result, 'notebook-456'); - }); - }); - suite('getOriginalProject', () => { test('should return undefined for unknown project', () => { const result = manager.getOriginalProject('unknown-project'); @@ -49,7 +33,7 @@ suite('DeepnoteNotebookManager', () => { }); test('should return original project after storing', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); + manager.storeOriginalProject('project-123', mockProject); const result = manager.getOriginalProject('project-123'); @@ -57,48 +41,13 @@ 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('storeOriginalProject', () => { - test('should store both project and current notebook ID', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); + test('should store project data', () => { + manager.storeOriginalProject('project-123', mockProject); const storedProject = manager.getOriginalProject('project-123'); - const currentNotebookId = manager.getCurrentNotebookId('project-123'); assert.deepStrictEqual(storedProject, mockProject); - assert.strictEqual(currentNotebookId, 'notebook-456'); }); test('should overwrite existing project data', () => { @@ -110,19 +59,17 @@ suite('DeepnoteNotebookManager', () => { } }; - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); - manager.storeOriginalProject('project-123', updatedProject, 'notebook-789'); + manager.storeOriginalProject('project-123', mockProject); + manager.storeOriginalProject('project-123', updatedProject); const storedProject = manager.getOriginalProject('project-123'); - const currentNotebookId = manager.getCurrentNotebookId('project-123'); assert.deepStrictEqual(storedProject, updatedProject); - assert.strictEqual(currentNotebookId, 'notebook-789'); }); }); suite('updateOriginalProject', () => { - test('should update project data without changing currentNotebookId', () => { + test('should update project data', () => { const updatedProject: DeepnoteProject = { ...mockProject, project: { @@ -131,14 +78,12 @@ suite('DeepnoteNotebookManager', () => { } }; - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); + manager.storeOriginalProject('project-123', mockProject); manager.updateOriginalProject('project-123', updatedProject); const storedProject = manager.getOriginalProject('project-123'); - const currentNotebookId = manager.getCurrentNotebookId('project-123'); assert.deepStrictEqual(storedProject, updatedProject); - assert.strictEqual(currentNotebookId, 'notebook-456'); }); test('should deep-clone project data so mutations to input do not affect stored state', () => { @@ -150,7 +95,7 @@ suite('DeepnoteNotebookManager', () => { } }; - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); + manager.storeOriginalProject('project-123', mockProject); manager.updateOriginalProject('project-123', updatedProject); updatedProject.project.name = 'After Mutation'; @@ -160,7 +105,7 @@ suite('DeepnoteNotebookManager', () => { assert.strictEqual(storedProject?.project.name, 'Before Mutation'); }); - test('should overwrite existing project data while preserving currentNotebookId', () => { + test('should overwrite existing project data on successive updates', () => { const firstUpdate: DeepnoteProject = { ...mockProject, project: { ...mockProject.project, name: 'First Update' } @@ -170,15 +115,14 @@ suite('DeepnoteNotebookManager', () => { project: { ...mockProject.project, name: 'Second Update' } }; - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); + manager.storeOriginalProject('project-123', mockProject); manager.updateOriginalProject('project-123', firstUpdate); manager.updateOriginalProject('project-123', secondUpdate); - assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-456'); assert.strictEqual(manager.getOriginalProject('project-123')?.project.name, 'Second Update'); }); - test('should store project when no currentNotebookId has been set', () => { + test('should store project when no prior data exists', () => { const projectOnly: DeepnoteProject = { ...mockProject, project: { ...mockProject.project, name: 'No Notebook Id Yet' } @@ -186,14 +130,13 @@ suite('DeepnoteNotebookManager', () => { manager.updateOriginalProject('project-123', projectOnly); - assert.strictEqual(manager.getCurrentNotebookId('project-123'), undefined); assert.deepStrictEqual(manager.getOriginalProject('project-123'), projectOnly); }); }); suite('updateProjectIntegrations', () => { test('should update integrations list for existing project and return true', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); + manager.storeOriginalProject('project-123', mockProject); const integrations: ProjectIntegration[] = [ { id: 'int-1', name: 'PostgreSQL', type: 'pgsql' }, @@ -217,7 +160,7 @@ suite('DeepnoteNotebookManager', () => { } }; - manager.storeOriginalProject('project-123', projectWithIntegrations, 'notebook-456'); + manager.storeOriginalProject('project-123', projectWithIntegrations); const newIntegrations: ProjectIntegration[] = [ { id: 'new-int-1', name: 'New Integration 1', type: 'pgsql' }, @@ -241,7 +184,7 @@ suite('DeepnoteNotebookManager', () => { } }; - manager.storeOriginalProject('project-123', projectWithIntegrations, 'notebook-456'); + manager.storeOriginalProject('project-123', projectWithIntegrations); const result = manager.updateProjectIntegrations('project-123', []); @@ -263,7 +206,7 @@ suite('DeepnoteNotebookManager', () => { }); test('should preserve other project properties and return true', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); + manager.storeOriginalProject('project-123', mockProject); const integrations: ProjectIntegration[] = [{ id: 'int-1', name: 'PostgreSQL', type: 'pgsql' }]; @@ -278,7 +221,7 @@ suite('DeepnoteNotebookManager', () => { assert.deepStrictEqual(updatedProject?.metadata, mockProject.metadata); }); - test('should update integrations when currentNotebookId is undefined and return true', () => { + test('should update integrations when project was stored via updateOriginalProject and return true', () => { // Use updateOriginalProject which doesn't set currentNotebookId manager.updateOriginalProject('project-123', mockProject); @@ -301,20 +244,33 @@ suite('DeepnoteNotebookManager', () => { }); }); - suite('integration scenarios', () => { - test('should handle complete workflow for multiple projects', () => { - manager.storeOriginalProject('project-1', mockProject, 'notebook-1'); - manager.storeOriginalProject('project-2', mockProject, 'notebook-2'); + suite('hasInitNotebookBeenRun', () => { + test('should return false for unknown project', () => { + assert.strictEqual(manager.hasInitNotebookBeenRun('unknown-project'), false); + }); + + test('should return true after marking init notebook as run', () => { + manager.markInitNotebookAsRun('project-123'); - assert.strictEqual(manager.getCurrentNotebookId('project-1'), 'notebook-1'); - assert.strictEqual(manager.getCurrentNotebookId('project-2'), 'notebook-2'); + assert.strictEqual(manager.hasInitNotebookBeenRun('project-123'), true); }); + }); - test('should handle notebook switching within same project via storeOriginalProject', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-1'); - manager.storeOriginalProject('project-123', mockProject, 'notebook-2'); + suite('markInitNotebookAsRun', () => { + test('should mark init notebook as run for a project', () => { + manager.markInitNotebookAsRun('project-123'); + + assert.strictEqual(manager.hasInitNotebookBeenRun('project-123'), true); + }); + }); + + suite('integration scenarios', () => { + test('should handle complete workflow for multiple projects', () => { + manager.storeOriginalProject('project-1', mockProject); + manager.storeOriginalProject('project-2', mockProject); - assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-2'); + assert.deepStrictEqual(manager.getOriginalProject('project-1'), mockProject); + assert.deepStrictEqual(manager.getOriginalProject('project-2'), mockProject); }); }); }); diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 531c645957..83d89ec145 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -1,18 +1,7 @@ import type { DeepnoteBlock, DeepnoteFile, DeepnoteSnapshot } from '@deepnote/blocks'; import { deserializeDeepnoteFile, isExecutableBlock, serializeDeepnoteSnapshot } from '@deepnote/blocks'; import { inject, injectable, optional } from 'inversify'; -import { - CancellationTokenSource, - l10n, - NotebookEdit, - NotebookRange, - workspace, - WorkspaceEdit, - type CancellationToken, - type NotebookData, - type NotebookDocument, - type NotebookSerializer -} from 'vscode'; +import { workspace, type CancellationToken, type NotebookData, type NotebookSerializer } from 'vscode'; import { computeHash } from '../../platform/common/crypto'; import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; @@ -57,8 +46,6 @@ function cloneWithoutCircularRefs(obj: T, seen = new WeakSet()): T { } } -const LAST_SERIALIZED_TTL_MS = 10_000; - /** * Serializer for converting between Deepnote YAML files and VS Code notebook format. * Handles reading/writing .deepnote files and manages project state persistence. @@ -66,8 +53,6 @@ const LAST_SERIALIZED_TTL_MS = 10_000; @injectable() export class DeepnoteNotebookSerializer implements NotebookSerializer { private converter = new DeepnoteDataConverter(); - private lastSerializedNotebookId: string | undefined; - private lastSerializedTimestamp = 0; constructor( @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, @@ -83,11 +68,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { * @param token Cancellation token (unused) * @returns Promise resolving to notebook data */ - async deserializeNotebook( - content: Uint8Array, - token: CancellationToken, - notebookId?: string - ): Promise { + async deserializeNotebook(content: Uint8Array, token: CancellationToken): Promise { logger.debug('DeepnoteSerializer: Deserializing Deepnote notebook'); if (token?.isCancellationRequested) { @@ -103,33 +84,12 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } const projectId = deepnoteFile.project.id; - const resolvedNotebookId = notebookId ?? this.findCurrentNotebookId(projectId); - - logger.debug(`DeepnoteSerializer: Project ID: ${projectId}, Selected notebook ID: ${resolvedNotebookId}`); - - if (!resolvedNotebookId) { - logger.debug( - 'DeepnoteSerializer: No notebook ID resolved, returning empty state for post-open verification' - ); + const selectedNotebook = this.findDefaultNotebook(deepnoteFile); - return { - cells: [], - metadata: { - deepnoteProjectId: projectId, - deepnoteProjectName: deepnoteFile.project.name, - deepnoteVersion: deepnoteFile.version - } - }; - } - - if (deepnoteFile.project.notebooks.length === 0) { - throw new Error('Deepnote project contains no notebooks.'); - } - - const selectedNotebook = deepnoteFile.project.notebooks.find((nb) => nb.id === resolvedNotebookId); + logger.debug(`DeepnoteSerializer: Project ID: ${projectId}, Selected notebook ID: ${selectedNotebook?.id}`); if (!selectedNotebook) { - throw new Error(l10n.t('No notebook selected or found')); + throw new Error('Deepnote project contains no notebooks.'); } // Initialize vega-lite for output conversion (lazy-loaded ESM module) @@ -158,7 +118,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { if (this.snapshotService?.isSnapshotsEnabled()) { logger.debug(`[Snapshot] Snapshots enabled, reading snapshot for project ${projectId}`); try { - const snapshotOutputs = await this.snapshotService.readSnapshot(projectId); + const snapshotOutputs = await this.snapshotService.readSnapshot(projectId, selectedNotebook.id); if (snapshotOutputs && snapshotOutputs.size > 0) { logger.debug(`[Snapshot] Merging ${snapshotOutputs.size} block outputs from snapshot`); @@ -190,7 +150,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { ); } - this.notebookManager.storeOriginalProject(deepnoteFile.project.id, deepnoteFile, selectedNotebook.id); + this.notebookManager.storeOriginalProject(deepnoteFile.project.id, deepnoteFile); logger.debug(`DeepnoteSerializer: Stored project ${projectId} in notebook manager`); return { @@ -214,27 +174,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } } - /** - * Resolves the notebook ID for a deserialization call. - * - * Priority: - * 1. Pending resolution hint (explicit intent from explorer view) - * 2. Already-open document metadata (re-deserialization after file change) - * — skips the notebook that was just serialized, since that's the one - * that triggered the file change, not the one being re-deserialized. - * 3. undefined — initial open; verifyDeserializedNotebook will resolve - * from the document URI after open. - */ - findCurrentNotebookId(projectId: string): string | undefined { - const pendingNotebookId = this.notebookManager.consumePendingNotebookResolution(projectId); - - if (pendingNotebookId) { - return pendingNotebookId; - } - - return this.findNotebookIdFromOpenDocuments(projectId); - } - /** * Gets the data converter instance for cell/block conversion. * @returns DeepnoteDataConverter instance @@ -243,81 +182,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { return this.converter; } - /** - * Parses the file content and returns the ID of the default notebook - * (alphabetically first, excluding Init when other notebooks exist). - * Used by the post-open verification handler for direct file opens - * that have no `?notebook=` query param. - */ - resolveDefaultNotebookId(content: Uint8Array): string | undefined { - try { - const contentString = new TextDecoder('utf-8').decode(content); - const deepnoteFile = deserializeDeepnoteFile(contentString); - - return this.findDefaultNotebook(deepnoteFile)?.id; - } catch { - return undefined; - } - } - - /** - * Ensures an opened notebook document shows the correct notebook. - * - * The serializer returns an empty state when it cannot determine which - * notebook to display (no pending resolution). This method is called - * after the document is created, reads the real notebook ID from - * the URI `?notebook=` query param (or picks the default notebook - * for direct file opens), and patches the document in place. - */ - async verifyDeserializedNotebook(doc: NotebookDocument): Promise { - if (doc.notebookType !== 'deepnote') { - return; - } - - const expectedNotebookId = new URLSearchParams(doc.uri.query).get('notebook'); - const actualNotebookId = doc.metadata?.deepnoteNotebookId as string | undefined; - - if (actualNotebookId && (!expectedNotebookId || expectedNotebookId === actualNotebookId)) { - return; - } - - const cts = new CancellationTokenSource(); - - try { - const fileUri = doc.uri.with({ query: '', fragment: '' }); - const content = await workspace.fs.readFile(fileUri); - - const targetNotebookId = expectedNotebookId ?? this.resolveDefaultNotebookId(content); - - if (!targetNotebookId || targetNotebookId === actualNotebookId) { - return; - } - - logger.info( - `Notebook verification: resolving notebook ${targetNotebookId} (was ${actualNotebookId ?? 'empty'}).` - ); - - const correctData = await this.deserializeNotebook(content, cts.token, targetNotebookId); - - const edit = new WorkspaceEdit(); - - edit.set(doc.uri, [ - NotebookEdit.replaceCells(new NotebookRange(0, doc.cellCount), correctData.cells), - NotebookEdit.updateNotebookMetadata(correctData.metadata ?? {}) - ]); - - const applied = await workspace.applyEdit(edit); - - if (applied) { - await doc.save(); - } - } catch (error) { - logger.error('Failed to verify/correct notebook content', error); - } finally { - cts.dispose(); - } - } - /** * Serializes VS Code notebook data back to Deepnote YAML format. * Converts cells to blocks, updates project data, and generates YAML. @@ -354,16 +218,12 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { logger.debug('SerializeNotebook: Got and cloned original project'); - const notebookId = - data.metadata?.deepnoteNotebookId || this.notebookManager.getCurrentNotebookId(projectId); + const notebookId = data.metadata?.deepnoteNotebookId; if (!notebookId) { throw new Error('Cannot determine which notebook to save'); } - this.lastSerializedNotebookId = notebookId; - this.lastSerializedTimestamp = Date.now(); - logger.debug(`SerializeNotebook: Notebook ID: ${notebookId}`); const notebook = originalProject.project.notebooks.find((nb: { id: string }) => nb.id === notebookId); @@ -654,61 +514,23 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } /** - * Returns a notebook ID from an already-open document for re-deserialization. - * - * When VS Code re-reads a file (e.g. after save or external change), it - * calls deserializeNotebook again for each open document backed by that - * file. The document already exists in workspace.notebookDocuments with - * the correct deepnoteNotebookId in its metadata. - * - * Skips the notebook that was most recently serialized — that document - * triggered the file change and is not the one being re-deserialized. - */ - private findNotebookIdFromOpenDocuments(projectId: string): string | undefined { - const recentSerialize = - this.lastSerializedNotebookId && Date.now() - this.lastSerializedTimestamp < LAST_SERIALIZED_TTL_MS; - - for (const doc of workspace.notebookDocuments) { - if (doc.notebookType !== 'deepnote' || doc.metadata?.deepnoteProjectId !== projectId) { - continue; - } - - const notebookId = doc.metadata?.deepnoteNotebookId as string | undefined; - - if (!notebookId) { - continue; - } - - if (recentSerialize && notebookId === this.lastSerializedNotebookId) { - continue; - } - - return notebookId; - } - - return undefined; - } - - /** - * Finds the default notebook to open when no selection is made. - * @param file - * @returns + * Finds the default notebook — the first non-init notebook in array order. + * Falls back to the first notebook if only init exists. */ private findDefaultNotebook(file: DeepnoteFile): DeepnoteNotebook | undefined { if (file.project.notebooks.length === 0) { return undefined; } - const sortedNotebooks = file.project.notebooks.slice().sort((a, b) => a.name.localeCompare(b.name)); - const sortedNotebooksWithoutInit = file.project.initNotebookId - ? sortedNotebooks.filter((nb) => nb.id !== file.project.initNotebookId) - : sortedNotebooks; + const notebooksWithoutInit = file.project.initNotebookId + ? file.project.notebooks.filter((nb) => nb.id !== file.project.initNotebookId) + : file.project.notebooks; - if (sortedNotebooksWithoutInit.length > 0) { - return sortedNotebooksWithoutInit[0]; + if (notebooksWithoutInit.length > 0) { + return notebooksWithoutInit[0]; } - return sortedNotebooks[0]; + return file.project.notebooks[0]; } /** diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index 02b6c369ce..a5028c27a6 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -1,13 +1,10 @@ import { deserializeDeepnoteFile, serializeDeepnoteFile, type DeepnoteFile } from '@deepnote/blocks'; import { assert } from 'chai'; import { parse as parseYaml } from 'yaml'; -import { when } from 'ts-mockito'; -import type { NotebookDocument } from 'vscode'; import { DeepnoteNotebookSerializer } from './deepnoteSerializer'; import { DeepnoteNotebookManager } from './deepnoteNotebookManager'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; -import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; suite('DeepnoteNotebookSerializer', () => { let serializer: DeepnoteNotebookSerializer; @@ -74,9 +71,7 @@ suite('DeepnoteNotebookSerializer', () => { } suite('deserializeNotebook', () => { - test('should deserialize valid project with queued notebook resolution', async () => { - manager.queueNotebookResolution('project-123', 'notebook-1'); - + test('should deserialize valid project', async () => { const yamlContent = ` version: '1.0.0' metadata: @@ -106,9 +101,6 @@ project: assert.isDefined(result); assert.isDefined(result.cells); assert.isArray(result.cells); - assert.strictEqual(result.cells.length, 1); - assert.strictEqual(result.metadata?.deepnoteProjectId, 'project-123'); - assert.strictEqual(result.metadata?.deepnoteNotebookId, 'notebook-1'); }); test('should throw error for empty content', async () => { @@ -129,7 +121,7 @@ project: ); }); - test('should return empty state when no notebooks found and no ID resolved', async () => { + test('should throw error when no notebooks found', async () => { const contentWithoutNotebooks = new TextEncoder().encode(` version: '1.0.0' metadata: @@ -141,39 +133,20 @@ project: settings: {} `); - const result = await serializer.deserializeNotebook(contentWithoutNotebooks, {} as any); - - assert.deepStrictEqual(result.cells, []); - assert.strictEqual(result.metadata?.deepnoteProjectId, 'project-123'); - assert.strictEqual(result.metadata?.deepnoteNotebookId, undefined); - }); - - test('should deserialize the specified notebook when notebookId is passed', async () => { - const content = projectToYaml(mockProject); - const result = await serializer.deserializeNotebook(content, {} as any, 'notebook-2'); - - assert.strictEqual(result.metadata?.deepnoteNotebookId, 'notebook-2'); - assert.strictEqual(result.metadata?.deepnoteNotebookName, 'Second Notebook'); - assert.strictEqual(result.cells.length, 1); - assert.include(result.cells[0].value, 'Title'); - }); - - test('should ignore queued resolution when explicit notebookId is provided', async () => { - manager.queueNotebookResolution('project-123', 'notebook-1'); - const content = projectToYaml(mockProject); - const result = await serializer.deserializeNotebook(content, {} as any, 'notebook-2'); - - assert.strictEqual(result.metadata?.deepnoteNotebookId, 'notebook-2'); - assert.strictEqual(result.metadata?.deepnoteNotebookName, 'Second Notebook'); + await assert.isRejected( + serializer.deserializeNotebook(contentWithoutNotebooks, {} as any), + /Failed to parse Deepnote file/ + ); }); - test('should fall back to findCurrentNotebookId when notebookId is undefined', async () => { - manager.queueNotebookResolution('project-123', 'notebook-1'); + test('should deserialize default notebook when no explicit notebook ID is provided', async () => { const content = projectToYaml(mockProject); const result = await serializer.deserializeNotebook(content, {} as any); + assert.strictEqual(result.metadata?.deepnoteProjectId, 'project-123'); assert.strictEqual(result.metadata?.deepnoteNotebookId, 'notebook-1'); assert.strictEqual(result.metadata?.deepnoteNotebookName, 'First Notebook'); + assert.isAbove(result.cells.length, 0); }); }); @@ -207,7 +180,7 @@ project: test('should serialize notebook when original project exists', async () => { // First store the original project - manager.storeOriginalProject('project-123', mockProject, 'notebook-1'); + manager.storeOriginalProject('project-123', mockProject); const mockNotebookData = { cells: [ @@ -234,8 +207,8 @@ project: assert.include(yamlString, 'notebook-1'); }); - test('should use current notebook ID when metadata notebook ID is missing', async () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-2'); + test('should throw error when metadata notebook ID is missing', async () => { + manager.storeOriginalProject('project-123', mockProject); const mockNotebookData = { cells: [ @@ -251,113 +224,10 @@ project: } }; - const result = await serializer.serializeNotebook(mockNotebookData as any, {} as any); - const serializedProject = deserializeDeepnoteFile(new TextDecoder().decode(result)); - - assert.strictEqual(serializedProject.project.notebooks[0].id, 'notebook-1'); - assert.strictEqual(serializedProject.project.notebooks[1].id, 'notebook-2'); - assert.strictEqual(serializedProject.project.notebooks[0].blocks?.[0].content, 'print("hello")'); - assert.strictEqual(serializedProject.project.notebooks[1].blocks?.[0].content, '# Updated second notebook'); - }); - - suite('multi-notebook save scenarios', () => { - teardown(() => { - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(undefined); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); - }); - - test('serializing notebook-1 should not change currentNotebookId', async () => { - manager.queueNotebookResolution('project-123', 'notebook-2'); - await serializer.deserializeNotebook(projectToYaml(mockProject), {} as any); - - assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-2'); - - const mockNotebookData = { - cells: [ - { - kind: 2, - languageId: 'python', - metadata: {}, - value: 'print("edited nb1")' - } - ], - metadata: { - deepnoteNotebookId: 'notebook-1', - deepnoteProjectId: 'project-123' - } - }; - - await serializer.serializeNotebook(mockNotebookData as any, {} as any); - - assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-2'); - }); - }); - }); - - suite('findCurrentNotebookId', () => { - teardown(() => { - // Reset only the specific mocks used in this suite - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(undefined); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); - when(mockedVSCodeNamespaces.window.tabGroups).thenReturn({ all: [] } as any); - }); - - test('should return queued notebook resolution when available', () => { - manager.queueNotebookResolution('project-123', 'queued-notebook'); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, 'queued-notebook'); - }); - - test('should consume queued notebook resolution only once', () => { - manager.queueNotebookResolution('project-123', 'queued-notebook'); - - assert.strictEqual(serializer.findCurrentNotebookId('project-123'), 'queued-notebook'); - assert.strictEqual(serializer.findCurrentNotebookId('project-123'), undefined); - }); - - test('should prioritize queued notebook resolution over current notebook and open documents', () => { - manager.queueNotebookResolution('project-123', 'queued-notebook'); - manager.storeOriginalProject('project-123', mockProject, 'current-notebook'); - - const mockNotebookDoc = { - then: undefined, - notebookType: 'deepnote', - metadata: { - deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'open-notebook' - }, - uri: {} as any, - version: 1, - isDirty: false, - isUntitled: false, - isClosed: false, - cellCount: 0, - cellAt: () => ({}) as any, - getCells: () => [], - save: async () => true - } as NotebookDocument; - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([mockNotebookDoc]); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, 'queued-notebook'); - }); - - test('should return undefined when no pending resolution or active tab exists', () => { - manager.storeOriginalProject('project-123', mockProject, 'current-notebook'); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, undefined); - }); - - test('should return undefined for unknown project', () => { - const result = serializer.findCurrentNotebookId('unknown-project'); - - assert.strictEqual(result, undefined); + await assert.isRejected( + serializer.serializeNotebook(mockNotebookData as any, {} as any), + /Cannot determine which notebook to save/ + ); }); }); @@ -384,16 +254,9 @@ project: }); test('should handle manager state operations', () => { - assert.isFunction(manager.consumePendingNotebookResolution, 'has consumePendingNotebookResolution method'); - assert.isFunction(manager.getCurrentNotebookId, 'has getCurrentNotebookId method'); assert.isFunction(manager.getOriginalProject, 'has getOriginalProject method'); - assert.isFunction(manager.queueNotebookResolution, 'has queueNotebookResolution method'); assert.isFunction(manager.storeOriginalProject, 'has storeOriginalProject method'); }); - - test('should have findCurrentNotebookId method', () => { - assert.isFunction(serializer.findCurrentNotebookId, 'has findCurrentNotebookId method'); - }); }); suite('data structure handling', () => { @@ -465,7 +328,7 @@ project: } }; - manager.storeOriginalProject('project-circular', projectWithCircularRef, 'notebook-1'); + manager.storeOriginalProject('project-circular', projectWithCircularRef); const notebookData = { cells: [ @@ -533,7 +396,7 @@ project: }; // Store the project - manager.storeOriginalProject('project-id-test', projectData, 'notebook-1'); + manager.storeOriginalProject('project-id-test', projectData); // Create cells with the EXACT metadata structure that deserializeNotebook produces // This simulates what VS Code should preserve from deserialization @@ -619,7 +482,7 @@ project: } }; - manager.storeOriginalProject('project-recover-ids', projectData, 'notebook-1'); + manager.storeOriginalProject('project-recover-ids', projectData); // Cells WITHOUT id metadata (simulating what VS Code might provide if it strips metadata) // But content matches the original block @@ -686,7 +549,7 @@ project: } }; - manager.storeOriginalProject('project-new-content', projectData, 'notebook-1'); + manager.storeOriginalProject('project-new-content', projectData); // Cell with different content than any original block const notebookData = { @@ -750,253 +613,6 @@ project: }); }); - suite('default notebook selection', () => { - test('should not select Init notebook when other notebooks are available', () => { - const projectData: DeepnoteFile = { - version: '1.0.0', - metadata: { - createdAt: '2023-01-01T00:00:00Z', - modifiedAt: '2023-01-02T00:00:00Z' - }, - project: { - id: 'project-with-init', - name: 'Project with Init', - initNotebookId: 'init-notebook', - notebooks: [ - { - id: 'init-notebook', - name: 'Init', - blocks: [ - { - id: 'block-init', - content: 'print("init")', - sortingKey: 'a0', - metadata: {}, - blockGroup: '1', - type: 'code' - } - ], - executionMode: 'block', - isModule: false - }, - { - id: 'main-notebook', - name: 'Main', - blocks: [ - { - id: 'block-main', - content: 'print("main")', - sortingKey: 'a0', - metadata: {}, - blockGroup: '1', - type: 'code' - } - ], - executionMode: 'block', - isModule: false - } - ], - settings: {} - } - }; - - const content = projectToYaml(projectData); - const result = serializer.resolveDefaultNotebookId(content); - - assert.strictEqual(result, 'main-notebook'); - }); - - test('should select Init notebook when it is the only notebook', () => { - const projectData: DeepnoteFile = { - version: '1.0.0', - metadata: { - createdAt: '2023-01-01T00:00:00Z', - modifiedAt: '2023-01-02T00:00:00Z' - }, - project: { - id: 'project-only-init', - name: 'Project with only Init', - initNotebookId: 'init-notebook', - notebooks: [ - { - id: 'init-notebook', - name: 'Init', - blocks: [ - { - id: 'block-init', - content: 'print("init")', - sortingKey: 'a0', - blockGroup: '1', - metadata: {}, - type: 'code' - } - ], - executionMode: 'block', - isModule: false - } - ], - settings: {} - } - }; - - const content = projectToYaml(projectData); - const result = serializer.resolveDefaultNotebookId(content); - - assert.strictEqual(result, 'init-notebook'); - }); - - test('should select alphabetically first notebook when no initNotebookId', () => { - const projectData: DeepnoteFile = { - version: '1.0.0', - metadata: { - createdAt: '2023-01-01T00:00:00Z', - modifiedAt: '2023-01-02T00:00:00Z' - }, - project: { - id: 'project-alphabetical', - name: 'Project Alphabetical', - notebooks: [ - { - id: 'zebra-notebook', - name: 'Zebra Notebook', - blocks: [ - { - id: 'block-z', - content: 'print("zebra")', - sortingKey: 'a0', - blockGroup: '1', - metadata: {}, - type: 'code' - } - ], - executionMode: 'block', - isModule: false - }, - { - id: 'alpha-notebook', - name: 'Alpha Notebook', - blocks: [ - { - id: 'block-a', - content: 'print("alpha")', - sortingKey: 'a0', - blockGroup: '1', - metadata: {}, - type: 'code' - } - ], - executionMode: 'block', - isModule: false - }, - { - id: 'bravo-notebook', - name: 'Bravo Notebook', - blocks: [ - { - id: 'block-b', - content: 'print("bravo")', - sortingKey: 'a0', - blockGroup: '1', - metadata: {}, - type: 'code' - } - ], - executionMode: 'block', - isModule: false - } - ], - settings: {} - } - }; - - const content = projectToYaml(projectData); - const result = serializer.resolveDefaultNotebookId(content); - - assert.strictEqual(result, 'alpha-notebook'); - }); - - test('should sort Init notebook last when multiple notebooks exist', () => { - const projectData: DeepnoteFile = { - version: '1.0.0', - metadata: { - createdAt: '2023-01-01T00:00:00Z', - modifiedAt: '2023-01-02T00:00:00Z' - }, - project: { - id: 'project-multiple', - name: 'Project with Multiple', - initNotebookId: 'init-notebook', - notebooks: [ - { - id: 'charlie-notebook', - name: 'Charlie', - blocks: [ - { - id: 'block-c', - content: 'print("charlie")', - sortingKey: 'a0', - blockGroup: '1', - metadata: {}, - type: 'code' - } - ], - executionMode: 'block', - isModule: false - }, - { - id: 'init-notebook', - name: 'Init', - blocks: [ - { - id: 'block-init', - content: 'print("init")', - sortingKey: 'a0', - blockGroup: '1', - metadata: {}, - type: 'code' - } - ], - executionMode: 'block', - isModule: false - }, - { - id: 'alpha-notebook', - name: 'Alpha', - blocks: [ - { - id: 'block-a', - content: 'print("alpha")', - sortingKey: 'a0', - blockGroup: '1', - metadata: {}, - type: 'code' - } - ], - executionMode: 'block', - isModule: false - } - ], - settings: {} - } - }; - - const content = projectToYaml(projectData); - const result = serializer.resolveDefaultNotebookId(content); - - // Should select Alpha, not Init even though "Init" comes before "Alpha" alphabetically when in upper case - assert.strictEqual(result, 'alpha-notebook'); - }); - - test('should return empty state from deserializeNotebook when no notebook ID is resolved', async () => { - const content = projectToYaml(mockProject); - const result = await serializer.deserializeNotebook(content, {} as any); - - assert.deepStrictEqual(result.cells, []); - assert.strictEqual(result.metadata?.deepnoteProjectId, 'project-123'); - assert.strictEqual(result.metadata?.deepnoteNotebookId, undefined); - }); - }); - suite('detectContentChanges', () => { test('should detect no changes when content is identical', () => { const project: DeepnoteFile = { @@ -1486,7 +1102,7 @@ project: } }; - manager.storeOriginalProject('project-snapshot-hash', projectData, 'notebook-1'); + manager.storeOriginalProject('project-snapshot-hash', projectData); const notebookData = { cells: [ @@ -1561,13 +1177,13 @@ project: }; // Serialize twice - manager.storeOriginalProject('project-deterministic', structuredClone(projectData), 'notebook-1'); + manager.storeOriginalProject('project-deterministic', structuredClone(projectData)); const result1 = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed1 = parseYaml(new TextDecoder().decode(result1)) as DeepnoteFile & { metadata: { snapshotHash?: string }; }; - manager.storeOriginalProject('project-deterministic', structuredClone(projectData), 'notebook-1'); + manager.storeOriginalProject('project-deterministic', structuredClone(projectData)); const result2 = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed2 = parseYaml(new TextDecoder().decode(result2)) as DeepnoteFile & { metadata: { snapshotHash?: string }; @@ -1662,7 +1278,7 @@ project: // Serialize 5 times and collect all hashes for (let i = 0; i < 5; i++) { - manager.storeOriginalProject('project-multi-serialize', structuredClone(projectData), 'notebook-1'); + manager.storeOriginalProject('project-multi-serialize', structuredClone(projectData)); const result = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed = parseYaml(new TextDecoder().decode(result)) as DeepnoteFile & { metadata: { snapshotHash?: string }; @@ -1711,7 +1327,7 @@ project: } }; - manager.storeOriginalProject('project-content-change', projectData1, 'notebook-1'); + manager.storeOriginalProject('project-content-change', projectData1); const notebookData1 = { cells: [ @@ -1789,7 +1405,7 @@ project: } }; - manager.storeOriginalProject('project-version-change', projectData1, 'notebook-1'); + manager.storeOriginalProject('project-version-change', projectData1); const notebookData = { cells: [ @@ -1813,7 +1429,7 @@ project: // Change version const projectData2: DeepnoteFile = { ...structuredClone(projectData1), version: '2.0' }; - manager.storeOriginalProject('project-version-change', projectData2, 'notebook-1'); + manager.storeOriginalProject('project-version-change', projectData2); const result2 = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed2 = parseYaml(new TextDecoder().decode(result2)) as DeepnoteFile & { @@ -1855,7 +1471,7 @@ project: } }; - manager.storeOriginalProject('project-integrations-change', projectData1, 'notebook-1'); + manager.storeOriginalProject('project-integrations-change', projectData1); const notebookData = { cells: [ @@ -1880,7 +1496,7 @@ project: // Add integrations const projectData2 = structuredClone(projectData1); projectData2.project.integrations = [{ id: 'int-1', name: 'PostgreSQL', type: 'postgres' }]; - manager.storeOriginalProject('project-integrations-change', projectData2, 'notebook-1'); + manager.storeOriginalProject('project-integrations-change', projectData2); const result2 = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed2 = parseYaml(new TextDecoder().decode(result2)) as DeepnoteFile & { @@ -1922,7 +1538,7 @@ project: } }; - manager.storeOriginalProject('project-env-hash', projectData1, 'notebook-1'); + manager.storeOriginalProject('project-env-hash', projectData1); const notebookData = { cells: [ @@ -1947,7 +1563,7 @@ project: // Add environment hash const projectData2 = structuredClone(projectData1); projectData2.environment = { hash: 'env-hash-123' }; - manager.storeOriginalProject('project-env-hash', projectData2, 'notebook-1'); + manager.storeOriginalProject('project-env-hash', projectData2); const result2 = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed2 = parseYaml(new TextDecoder().decode(result2)) as DeepnoteFile & { diff --git a/src/notebooks/deepnote/deepnoteTreeDataProvider.ts b/src/notebooks/deepnote/deepnoteTreeDataProvider.ts index be9e471e27..96c4bc2d1e 100644 --- a/src/notebooks/deepnote/deepnoteTreeDataProvider.ts +++ b/src/notebooks/deepnote/deepnoteTreeDataProvider.ts @@ -13,7 +13,12 @@ import { l10n } from 'vscode'; -import { DeepnoteTreeItem, DeepnoteTreeItemType, DeepnoteTreeItemContext } from './deepnoteTreeItem'; +import { + DeepnoteTreeItem, + DeepnoteTreeItemType, + DeepnoteTreeItemContext, + type ProjectGroupData +} from './deepnoteTreeItem'; import type { DeepnoteProject, DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; import { readDeepnoteProjectFile } from './deepnoteProjectUtils'; import { ILogger } from '../../platform/logging/types'; @@ -28,9 +33,17 @@ export function compareTreeItemsByLabel(a: DeepnoteTreeItem, b: DeepnoteTreeItem return labelA.toLowerCase().localeCompare(labelB.toLowerCase()); } +/** + * Loaded project file data + */ +interface LoadedProjectFile { + filePath: string; + project: DeepnoteProject; +} + /** * Tree data provider for the Deepnote explorer view. - * Manages the tree structure displaying Deepnote project files and their notebooks. + * Groups files by project ID, showing project groups at the top level. */ export class DeepnoteTreeDataProvider implements TreeDataProvider { private _onDidChangeTreeData: EventEmitter = new EventEmitter< @@ -67,82 +80,24 @@ export class DeepnoteTreeDataProvider implements TreeDataProvider { - // Get the cached tree item BEFORE clearing caches - const cacheKey = `project:${filePath}`; - const cachedTreeItem = this.treeItemCache.get(cacheKey); - - // Clear the project data cache to force reload this.cachedProjects.delete(filePath); - - if (cachedTreeItem) { - // Reload the project data and update the cached tree item - try { - const fileUri = Uri.file(filePath); - const project = await this.loadDeepnoteProject(fileUri); - if (project) { - // Update the tree item's data - cachedTreeItem.data = project; - - // Update visual fields (label, description, tooltip) to reflect changes - cachedTreeItem.updateVisualFields(); - } - } catch (error) { - this.logger?.error(`Failed to reload project ${filePath}`, error); - } - - // Fire change event with the existing cached tree item - this._onDidChangeTreeData.fire(cachedTreeItem); - } else { - // If not found in cache, do a full refresh - this._onDidChangeTreeData.fire(); - } + // Full refresh since project grouping may have changed + this._onDidChangeTreeData.fire(); } /** * Refresh notebooks for a specific project - * @param projectId The project ID whose notebooks should be refreshed */ public async refreshNotebook(projectId: string): Promise { - // Find the cached tree item by scanning the cache - let cachedTreeItem: DeepnoteTreeItem | undefined; - let filePath: string | undefined; - - for (const [key, item] of this.treeItemCache.entries()) { - if (key.startsWith('project:') && item.context.projectId === projectId) { - cachedTreeItem = item; - filePath = item.context.filePath; - break; + // Clear all cached projects that match this project ID to force reload + for (const [path, project] of this.cachedProjects.entries()) { + if (project.project.id === projectId) { + this.cachedProjects.delete(path); } } - - if (cachedTreeItem && filePath) { - // Clear the project data cache to force reload - this.cachedProjects.delete(filePath); - - // Reload the project data and update the cached tree item - try { - const fileUri = Uri.file(filePath); - const project = await this.loadDeepnoteProject(fileUri); - if (project) { - // Update the tree item's data - cachedTreeItem.data = project; - - // Update visual fields (label, description, tooltip) to reflect changes - cachedTreeItem.updateVisualFields(); - } - } catch (error) { - this.logger?.error(`Failed to reload project ${filePath}`, error); - } - - // Fire change event with the existing cached tree item to refresh its children - this._onDidChangeTreeData.fire(cachedTreeItem); - } else { - // If not found in cache, do a full refresh - this._onDidChangeTreeData.fire(); - } + this._onDidChangeTreeData.fire(); } public getTreeItem(element: DeepnoteTreeItem): TreeItem { @@ -150,8 +105,11 @@ export class DeepnoteTreeDataProvider implements TreeDataProvider { - // If element is provided, we can return children regardless of workspace if (element) { + if (element.type === DeepnoteTreeItemType.ProjectGroup) { + return this.getFilesForProjectGroup(element); + } + if (element.type === DeepnoteTreeItemType.ProjectFile) { return this.getNotebooksForProject(element); } @@ -159,7 +117,7 @@ export class DeepnoteTreeDataProvider implements TreeDataProvider { try { - await this.getDeepnoteProjectFiles(); + await this.loadAllProjectFiles(); } finally { this.isInitialScanComplete = true; this.initialScanPromise = undefined; @@ -198,8 +156,11 @@ export class DeepnoteTreeDataProvider implements TreeDataProvider { - const deepnoteFiles: DeepnoteTreeItem[] = []; + /** + * Load all .deepnote project files from the workspace + */ + private async loadAllProjectFiles(): Promise { + const results: LoadedProjectFile[] = []; for (const workspaceFolder of workspace.workspaceFolders || []) { const pattern = new RelativePattern(workspaceFolder, '**/*.deepnote'); @@ -209,64 +170,128 @@ export class DeepnoteTreeDataProvider implements TreeDataProvider { + const allFiles = await this.loadAllProjectFiles(); + + // Group by project ID + const groupsByProjectId = new Map(); + + for (const file of allFiles) { + const projectId = file.project.project.id; + const existing = groupsByProjectId.get(projectId) ?? []; + existing.push(file); + groupsByProjectId.set(projectId, existing); + } + + const groups: DeepnoteTreeItem[] = []; + + for (const [projectId, files] of groupsByProjectId.entries()) { + const projectName = files[0].project.project.name; + + // If only one file with one non-init notebook, show directly as a project file (no group nesting) + if (files.length === 1) { + const file = files[0]; + const initNotebookId = file.project.project.initNotebookId; + const nonInitNotebooks = file.project.project.notebooks?.filter((nb) => nb.id !== initNotebookId) ?? []; + + if (nonInitNotebooks.length <= 1) { const context: DeepnoteTreeItemContext = { - filePath: file.path, - projectId: project.project.id + filePath: file.filePath, + projectId }; - // Check if we have a cached tree item for this project - const cacheKey = `project:${file.path}`; - let treeItem = this.treeItemCache.get(cacheKey); - - if (!treeItem) { - // Create new tree item only if not cached - const hasNotebooks = project.project.notebooks && project.project.notebooks.length > 0; - const collapsibleState = hasNotebooks - ? TreeItemCollapsibleState.Collapsed - : TreeItemCollapsibleState.None; - - treeItem = new DeepnoteTreeItem( - DeepnoteTreeItemType.ProjectFile, - context, - project, - collapsibleState - ); - - this.treeItemCache.set(cacheKey, treeItem); - } else { - // Update the cached tree item's data - treeItem.data = project; - } - - deepnoteFiles.push(treeItem); - } catch (error) { - this.logger?.error(`Failed to load Deepnote project from ${file.path}`, error); + const treeItem = new DeepnoteTreeItem( + DeepnoteTreeItemType.ProjectFile, + context, + file.project, + TreeItemCollapsibleState.None + ); + groups.push(treeItem); + continue; } } + + // Multiple files or multi-notebook file: create a group + const groupData: ProjectGroupData = { + projectId, + projectName, + files: files.map((f) => ({ filePath: f.filePath, project: f.project })) + }; + + const context: DeepnoteTreeItemContext = { + filePath: files[0].filePath, + projectId + }; + + const groupItem = new DeepnoteTreeItem( + DeepnoteTreeItemType.ProjectGroup, + context, + groupData, + TreeItemCollapsibleState.Collapsed + ); + groups.push(groupItem); + } + + groups.sort(compareTreeItemsByLabel); + + return groups; + } + + /** + * Get file items for a project group + */ + private getFilesForProjectGroup(groupItem: DeepnoteTreeItem): DeepnoteTreeItem[] { + const groupData = groupItem.data as ProjectGroupData; + const fileItems: DeepnoteTreeItem[] = []; + + for (const file of groupData.files) { + const initNotebookId = file.project.project.initNotebookId; + const nonInitNotebooks = file.project.project.notebooks?.filter((nb) => nb.id !== initNotebookId) ?? []; + const hasMultipleNotebooks = nonInitNotebooks.length > 1; + + const context: DeepnoteTreeItemContext = { + filePath: file.filePath, + projectId: groupData.projectId + }; + + const treeItem = new DeepnoteTreeItem( + DeepnoteTreeItemType.ProjectFile, + context, + file.project, + hasMultipleNotebooks ? TreeItemCollapsibleState.Collapsed : TreeItemCollapsibleState.None + ); + fileItems.push(treeItem); } - // Sort projects alphabetically by name (case-insensitive) - deepnoteFiles.sort(compareTreeItemsByLabel); + fileItems.sort(compareTreeItemsByLabel); - return deepnoteFiles; + return fileItems; } - private async getNotebooksForProject(projectItem: DeepnoteTreeItem): Promise { + /** + * Get notebook items for a project file (shown for multi-notebook files before splitting) + */ + private getNotebooksForProject(projectItem: DeepnoteTreeItem): DeepnoteTreeItem[] { const project = projectItem.data as DeepnoteProject; const notebooks = project.project.notebooks || []; - // Sort notebooks alphabetically by name (case-insensitive) - const sortedNotebooks = [...notebooks].sort((a, b) => { - const nameA = a.name || ''; - const nameB = b.name || ''; - return nameA.toLowerCase().localeCompare(nameB.toLowerCase()); - }); - - return sortedNotebooks.map((notebook: DeepnoteNotebook) => { + return notebooks.map((notebook: DeepnoteNotebook) => { const context: DeepnoteTreeItemContext = { filePath: projectItem.context.filePath, projectId: projectItem.context.projectId, @@ -312,7 +337,6 @@ export class DeepnoteTreeDataProvider implements TreeDataProvider { - const projectFiles = await this.getDeepnoteProjectFiles(); + public async findTreeItem(projectId: string): Promise { + const groups = await this.getProjectGroups(); - for (const projectItem of projectFiles) { - if (projectItem.context.projectId === projectId) { - if (!notebookId) { - return projectItem; - } - - const notebooks = await this.getNotebooksForProject(projectItem); - for (const notebookItem of notebooks) { - if (notebookItem.context.notebookId === notebookId) { - return notebookItem; - } - } + for (const item of groups) { + if (item.context.projectId === projectId) { + return item; } } diff --git a/src/notebooks/deepnote/deepnoteTreeDataProvider.unit.test.ts b/src/notebooks/deepnote/deepnoteTreeDataProvider.unit.test.ts index 4944422590..4c10e5408b 100644 --- a/src/notebooks/deepnote/deepnoteTreeDataProvider.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteTreeDataProvider.unit.test.ts @@ -379,8 +379,8 @@ suite('DeepnoteTreeDataProvider', () => { treeItemCache.set(cacheKey, mockTreeItem); // Verify initial state - assert.strictEqual(mockTreeItem.label, 'Original Name'); - assert.strictEqual(mockTreeItem.description, '1 notebook'); + assert.strictEqual(mockTreeItem.label, 'test-project.deepnote'); + assert.strictEqual(mockTreeItem.description, '0 cells'); // Update the project data (simulating rename and adding notebooks) const updatedProject: DeepnoteProject = { @@ -407,19 +407,20 @@ suite('DeepnoteTreeDataProvider', () => { mockTreeItem.updateVisualFields(); } else { // Manually update visual fields for testing purposes - mockTreeItem.label = updatedProject.project.name || 'Untitled Project'; + const fileName = mockTreeItem.context.filePath.split('/').pop() ?? ''; + mockTreeItem.label = fileName || updatedProject.project.name || 'Untitled Project'; mockTreeItem.tooltip = `Deepnote Project: ${updatedProject.project.name}\nFile: ${mockTreeItem.context.filePath}`; - const notebookCount = updatedProject.project.notebooks?.length || 0; - mockTreeItem.description = `${notebookCount} notebook${notebookCount !== 1 ? 's' : ''}`; + const blockCount = + updatedProject.project.notebooks?.reduce( + (sum: number, nb: any) => sum + (nb.blocks?.length ?? 0), + 0 + ) ?? 0; + mockTreeItem.description = `${blockCount} cell${blockCount !== 1 ? 's' : ''}`; } // Verify visual fields were updated - assert.strictEqual(mockTreeItem.label, 'Renamed Project', 'Label should reflect new project name'); - assert.strictEqual( - mockTreeItem.description, - '2 notebooks', - 'Description should reflect new notebook count' - ); + assert.strictEqual(mockTreeItem.label, 'test-project.deepnote', 'Label should reflect filename'); + assert.strictEqual(mockTreeItem.description, '0 cells', 'Description should reflect cell count'); assert.include( mockTreeItem.tooltip as string, 'Renamed Project', @@ -522,16 +523,16 @@ suite('DeepnoteTreeDataProvider', () => { ) ); - // Verify items are initially unsorted - assert.strictEqual(treeItems[0].label, 'Zebra Project'); + // Verify items are initially unsorted (label is filename derived from project name) + assert.strictEqual(treeItems[0].label, 'Zebra Project.deepnote'); // Sort using the exported comparator const sortedItems = [...treeItems].sort(compareTreeItemsByLabel); // Verify alphabetical order - assert.strictEqual(sortedItems[0].label, 'Apple Project'); - assert.strictEqual(sortedItems[1].label, 'Middle Project'); - assert.strictEqual(sortedItems[2].label, 'Zebra Project'); + assert.strictEqual(sortedItems[0].label, 'Apple Project.deepnote'); + assert.strictEqual(sortedItems[1].label, 'Middle Project.deepnote'); + assert.strictEqual(sortedItems[2].label, 'Zebra Project.deepnote'); }); test('should sort notebooks alphabetically by name within a project', async () => { @@ -584,11 +585,11 @@ suite('DeepnoteTreeDataProvider', () => { const notebookItems = await provider.getChildren(mockProjectItem); - // Verify notebooks are sorted alphabetically + // Verify notebooks are returned in original order assert.strictEqual(notebookItems.length, 3, 'Should have 3 notebooks'); - assert.strictEqual(notebookItems[0].label, 'Apple Notebook', 'First notebook should be Apple Notebook'); - assert.strictEqual(notebookItems[1].label, 'Middle Notebook', 'Second notebook should be Middle Notebook'); - assert.strictEqual(notebookItems[2].label, 'Zebra Notebook', 'Third notebook should be Zebra Notebook'); + assert.strictEqual(notebookItems[0].label, 'Zebra Notebook', 'First notebook should be Zebra Notebook'); + assert.strictEqual(notebookItems[1].label, 'Apple Notebook', 'Second notebook should be Apple Notebook'); + assert.strictEqual(notebookItems[2].label, 'Middle Notebook', 'Third notebook should be Middle Notebook'); }); test('should sort notebooks case-insensitively', async () => { @@ -641,11 +642,11 @@ suite('DeepnoteTreeDataProvider', () => { const notebookItems = await provider.getChildren(mockProjectItem); - // Verify case-insensitive sorting + // Verify notebooks are returned in original order assert.strictEqual(notebookItems.length, 3, 'Should have 3 notebooks'); - assert.strictEqual(notebookItems[0].label, 'Apple Notebook', 'First should be Apple Notebook'); - assert.strictEqual(notebookItems[1].label, 'MIDDLE Notebook', 'Second should be MIDDLE Notebook'); - assert.strictEqual(notebookItems[2].label, 'zebra notebook', 'Third should be zebra notebook'); + assert.strictEqual(notebookItems[0].label, 'zebra notebook', 'First should be zebra notebook'); + assert.strictEqual(notebookItems[1].label, 'Apple Notebook', 'Second should be Apple Notebook'); + assert.strictEqual(notebookItems[2].label, 'MIDDLE Notebook', 'Third should be MIDDLE Notebook'); }); }); }); diff --git a/src/notebooks/deepnote/deepnoteTreeItem.ts b/src/notebooks/deepnote/deepnoteTreeItem.ts index 0763f0e1bb..15577142da 100644 --- a/src/notebooks/deepnote/deepnoteTreeItem.ts +++ b/src/notebooks/deepnote/deepnoteTreeItem.ts @@ -1,10 +1,12 @@ -import { TreeItem, TreeItemCollapsibleState, Uri, ThemeIcon } from 'vscode'; +import { TreeItem, TreeItemCollapsibleState, ThemeIcon } from 'vscode'; + import type { DeepnoteProject, DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; /** * Represents different types of items in the Deepnote tree view */ export enum DeepnoteTreeItemType { + ProjectGroup = 'projectGroup', ProjectFile = 'projectFile', Notebook = 'notebook', Loading = 'loading' @@ -20,25 +22,39 @@ export interface DeepnoteTreeItemContext { } /** - * Tree item representing a Deepnote project file or notebook in the explorer view + * Data associated with a ProjectGroup tree item + */ +export interface ProjectGroupData { + readonly projectId: string; + readonly projectName: string; + readonly files: Array<{ filePath: string; project: DeepnoteProject }>; +} + +/** + * Tree item representing a Deepnote project group, file, or notebook in the explorer view */ export class DeepnoteTreeItem extends TreeItem { constructor( public readonly type: DeepnoteTreeItemType, public readonly context: DeepnoteTreeItemContext, - public data: DeepnoteProject | DeepnoteNotebook | null, + public data: DeepnoteProject | DeepnoteNotebook | ProjectGroupData | null, collapsibleState: TreeItemCollapsibleState ) { super('', collapsibleState); this.contextValue = this.type; - // Inline method calls to avoid ES module TreeItem extension issues if (this.type === DeepnoteTreeItemType.Loading) { this.label = 'Loading…'; this.tooltip = 'Loading…'; this.description = ''; this.iconPath = new ThemeIcon('loading~spin'); + } else if (this.type === DeepnoteTreeItemType.ProjectGroup) { + const groupData = this.data as ProjectGroupData; + this.label = groupData.projectName || 'Untitled Project'; + this.tooltip = `Deepnote Project: ${groupData.projectName}\n${groupData.files.length} file(s)`; + this.description = `${groupData.files.length} file${groupData.files.length !== 1 ? 's' : ''}`; + this.iconPath = new ThemeIcon('notebook'); } else { // getTooltip() inline if (this.type === DeepnoteTreeItemType.ProjectFile) { @@ -51,7 +67,7 @@ export class DeepnoteTreeItem extends TreeItem { // getIcon() inline if (this.type === DeepnoteTreeItemType.ProjectFile) { - this.iconPath = new ThemeIcon('notebook'); + this.iconPath = new ThemeIcon('file-code'); } else { this.iconPath = new ThemeIcon('file-code'); } @@ -59,7 +75,8 @@ export class DeepnoteTreeItem extends TreeItem { // getLabel() inline if (this.type === DeepnoteTreeItemType.ProjectFile) { const project = this.data as DeepnoteProject; - this.label = project.project.name || 'Untitled Project'; + const fileName = this.context.filePath.split('/').pop() ?? ''; + this.label = fileName || project.project.name || 'Untitled Project'; } else { const notebook = this.data as DeepnoteNotebook; this.label = notebook.name || 'Untitled Notebook'; @@ -68,8 +85,10 @@ export class DeepnoteTreeItem extends TreeItem { // getDescription() inline if (this.type === DeepnoteTreeItemType.ProjectFile) { const project = this.data as DeepnoteProject; - const notebookCount = project.project.notebooks?.length || 0; - this.description = `${notebookCount} notebook${notebookCount !== 1 ? 's' : ''}`; + const initNotebookId = project.project.initNotebookId; + const nonInitNotebooks = project.project.notebooks?.filter((nb) => nb.id !== initNotebookId) ?? []; + const blockCount = nonInitNotebooks.reduce((sum, nb) => sum + (nb.blocks?.length ?? 0), 0); + this.description = `${blockCount} cell${blockCount !== 1 ? 's' : ''}`; } else { const notebook = this.data as DeepnoteNotebook; const blockCount = notebook.blocks?.length || 0; @@ -77,11 +96,17 @@ export class DeepnoteTreeItem extends TreeItem { } } + // ProjectFile items open the file directly (no query param) + if (this.type === DeepnoteTreeItemType.ProjectFile) { + this.command = { + command: 'deepnote.openNotebook', + title: 'Open Notebook', + arguments: [this.context] + }; + } + + // Notebook items also open the file directly (no query param) if (this.type === DeepnoteTreeItemType.Notebook) { - // getNotebookUri() inline - if (this.context.notebookId) { - this.resourceUri = Uri.parse(`deepnote-notebook://${this.context.filePath}#${this.context.notebookId}`); - } this.command = { command: 'deepnote.openNotebook', title: 'Open Notebook', @@ -103,15 +128,24 @@ export class DeepnoteTreeItem extends TreeItem { return; } + if (this.type === DeepnoteTreeItemType.ProjectGroup) { + const groupData = this.data as ProjectGroupData; + this.label = groupData.projectName || 'Untitled Project'; + this.tooltip = `Deepnote Project: ${groupData.projectName}\n${groupData.files.length} file(s)`; + this.description = `${groupData.files.length} file${groupData.files.length !== 1 ? 's' : ''}`; + return; + } + if (this.type === DeepnoteTreeItemType.ProjectFile) { const project = this.data as DeepnoteProject; - - this.label = project.project.name || 'Untitled Project'; + const fileName = this.context.filePath.split('/').pop() ?? ''; + this.label = fileName || project.project.name || 'Untitled Project'; this.tooltip = `Deepnote Project: ${project.project.name}\nFile: ${this.context.filePath}`; - const notebookCount = project.project.notebooks?.length || 0; - - this.description = `${notebookCount} notebook${notebookCount !== 1 ? 's' : ''}`; + const initNotebookId = project.project.initNotebookId; + const nonInitNotebooks = project.project.notebooks?.filter((nb) => nb.id !== initNotebookId) ?? []; + const blockCount = nonInitNotebooks.reduce((sum, nb) => sum + (nb.blocks?.length ?? 0), 0); + this.description = `${blockCount} cell${blockCount !== 1 ? 's' : ''}`; } else { const notebook = this.data as DeepnoteNotebook; diff --git a/src/notebooks/deepnote/deepnoteTreeItem.unit.test.ts b/src/notebooks/deepnote/deepnoteTreeItem.unit.test.ts index bbadeceee4..cbc2a99bf3 100644 --- a/src/notebooks/deepnote/deepnoteTreeItem.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteTreeItem.unit.test.ts @@ -61,8 +61,8 @@ suite('DeepnoteTreeItem', () => { assert.strictEqual(item.type, DeepnoteTreeItemType.ProjectFile); assert.deepStrictEqual(item.context, context); assert.strictEqual(item.collapsibleState, TreeItemCollapsibleState.Collapsed); - assert.strictEqual(item.label, 'Test Project'); - assert.strictEqual(item.description, '1 notebook'); + assert.strictEqual(item.label, 'project.deepnote'); + assert.strictEqual(item.description, '0 cells'); }); test('should create notebook item with basic properties', () => { @@ -117,19 +117,20 @@ suite('DeepnoteTreeItem', () => { TreeItemCollapsibleState.Collapsed ); - assert.strictEqual(item.label, 'Test Project'); + assert.strictEqual(item.label, 'my-project.deepnote'); assert.strictEqual(item.type, DeepnoteTreeItemType.ProjectFile); assert.strictEqual(item.collapsibleState, TreeItemCollapsibleState.Collapsed); assert.strictEqual(item.contextValue, 'projectFile'); assert.strictEqual(item.tooltip, 'Deepnote Project: Test Project\nFile: /workspace/my-project.deepnote'); - assert.strictEqual(item.description, '1 notebook'); + assert.strictEqual(item.description, '0 cells'); - // Should have notebook icon for project files + // Should have file-code icon for project files assert.instanceOf(item.iconPath, ThemeIcon); - assert.strictEqual((item.iconPath as ThemeIcon).id, 'notebook'); + assert.strictEqual((item.iconPath as ThemeIcon).id, 'file-code'); - // Should not have command for project files - assert.isUndefined(item.command); + // Should have command for project files + assert.isDefined(item.command); + assert.strictEqual(item.command!.command, 'deepnote.openNotebook'); }); test('should handle project with multiple notebooks', () => { @@ -175,7 +176,7 @@ suite('DeepnoteTreeItem', () => { TreeItemCollapsibleState.Collapsed ); - assert.strictEqual(item.description, '3 notebooks'); + assert.strictEqual(item.description, '0 cells'); }); test('should handle project with no notebooks', () => { @@ -199,7 +200,7 @@ suite('DeepnoteTreeItem', () => { TreeItemCollapsibleState.Collapsed ); - assert.strictEqual(item.description, '0 notebooks'); + assert.strictEqual(item.description, '0 cells'); }); test('should handle unnamed project', () => { @@ -223,7 +224,7 @@ suite('DeepnoteTreeItem', () => { TreeItemCollapsibleState.Collapsed ); - assert.strictEqual(item.label, 'Untitled Project'); + assert.strictEqual(item.label, 'project.deepnote'); }); }); @@ -259,12 +260,8 @@ suite('DeepnoteTreeItem', () => { assert.strictEqual(item.command!.title, 'Open Notebook'); assert.deepStrictEqual(item.command!.arguments, [context]); - // Should have resource URI - assert.isDefined(item.resourceUri); - assert.strictEqual( - item.resourceUri!.toString(), - 'deepnote-notebook:/workspace/project.deepnote#notebook-789' - ); + // Should not have resource URI + assert.isUndefined(item.resourceUri); }); test('should handle notebook with multiple blocks', () => { @@ -418,7 +415,7 @@ suite('DeepnoteTreeItem', () => { }); suite('command configuration', () => { - test('should not create command for project files', () => { + test('should create command for project files', () => { const context: DeepnoteTreeItemContext = { filePath: '/test/project.deepnote', projectId: 'project-123' @@ -431,7 +428,10 @@ suite('DeepnoteTreeItem', () => { TreeItemCollapsibleState.Collapsed ); - assert.isUndefined(item.command); + assert.isDefined(item.command); + assert.strictEqual(item.command!.command, 'deepnote.openNotebook'); + assert.strictEqual(item.command!.title, 'Open Notebook'); + assert.deepStrictEqual(item.command!.arguments, [context]); }); test('should create correct command for notebooks', () => { @@ -457,7 +457,7 @@ suite('DeepnoteTreeItem', () => { }); suite('icon configuration', () => { - test('should use notebook icon for project files', () => { + test('should use file-code icon for project files', () => { const context: DeepnoteTreeItemContext = { filePath: '/test/project.deepnote', projectId: 'project-123' @@ -471,7 +471,7 @@ suite('DeepnoteTreeItem', () => { ); assert.instanceOf(item.iconPath, ThemeIcon); - assert.strictEqual((item.iconPath as ThemeIcon).id, 'notebook'); + assert.strictEqual((item.iconPath as ThemeIcon).id, 'file-code'); }); test('should use file-code icon for notebooks', () => { @@ -746,7 +746,8 @@ suite('DeepnoteTreeItem', () => { items.forEach((item, index) => { assert.strictEqual(item.context.filePath, contexts[index].filePath); assert.strictEqual(item.context.projectId, contexts[index].projectId); - assert.isUndefined(item.command); // Project files don't have commands + assert.isDefined(item.command); // Project files have commands + assert.strictEqual(item.command!.command, 'deepnote.openNotebook'); }); }); }); diff --git a/src/notebooks/deepnote/snapshots/snapshotFiles.ts b/src/notebooks/deepnote/snapshots/snapshotFiles.ts index 6f6ff9f5fd..109411da9f 100644 --- a/src/notebooks/deepnote/snapshots/snapshotFiles.ts +++ b/src/notebooks/deepnote/snapshots/snapshotFiles.ts @@ -5,8 +5,21 @@ import { InvalidProjectNameError } from '../../../platform/errors/invalidProject /** File suffix for snapshot files */ export const SNAPSHOT_FILE_SUFFIX = '.snapshot.deepnote'; -/** Regex pattern for extracting project ID from snapshot filenames. */ -const SNAPSHOT_FILENAME_PATTERN = new RegExp(`^[a-z0-9-]+_(.+)_[^_]+${SNAPSHOT_FILE_SUFFIX.replace(/\./g, '\\.')}$`); +/** + * Regex pattern for extracting project ID and notebook ID from snapshot filenames. + * New format: {slug}_{projectId}_{notebookId}_{variant}.snapshot.deepnote + */ +const SNAPSHOT_FILENAME_PATTERN = new RegExp( + `^[a-z0-9-]+_(.+)_([a-f0-9-]+)_[^_]+${SNAPSHOT_FILE_SUFFIX.replace(/\./g, '\\.')}$` +); + +/** + * Legacy pattern for old-format snapshots without notebook ID. + * Old format: {slug}_{projectId}_{variant}.snapshot.deepnote + */ +const LEGACY_SNAPSHOT_FILENAME_PATTERN = new RegExp( + `^[a-z0-9-]+_(.+)_[^_]+${SNAPSHOT_FILE_SUFFIX.replace(/\./g, '\\.')}$` +); /** * Checks if a URI represents a snapshot file @@ -17,14 +30,34 @@ export function isSnapshotFile(uri: Uri): boolean { /** * Extracts the project ID from a snapshot file URI. - * Snapshot filenames follow: `${slug}_${projectId}_${variant}.snapshot.deepnote` + * Supports both new format ({slug}_{projectId}_{notebookId}_{variant}) and + * legacy format ({slug}_{projectId}_{variant}). * @returns The project ID, or undefined if the URI is not a valid snapshot file */ export function extractProjectIdFromSnapshotUri(uri: Uri): string | undefined { const basename = uri.path.split('/').pop() ?? ''; + + // Try new format first + const newMatch = basename.match(SNAPSHOT_FILENAME_PATTERN); + if (newMatch) { + return newMatch[1]; + } + + // Fall back to legacy format + const legacyMatch = basename.match(LEGACY_SNAPSHOT_FILENAME_PATTERN); + return legacyMatch?.[1]; +} + +/** + * Extracts the notebook ID from a snapshot file URI. + * Only works with new format: {slug}_{projectId}_{notebookId}_{variant}.snapshot.deepnote + * @returns The notebook ID, or undefined if the URI uses legacy format or is invalid + */ +export function extractNotebookIdFromSnapshotUri(uri: Uri): string | undefined { + const basename = uri.path.split('/').pop() ?? ''; const match = basename.match(SNAPSHOT_FILENAME_PATTERN); - return match?.[1]; + return match?.[2]; } /** diff --git a/src/notebooks/deepnote/snapshots/snapshotService.ts b/src/notebooks/deepnote/snapshots/snapshotService.ts index 6e4b99e96e..4ed8095628 100644 --- a/src/notebooks/deepnote/snapshots/snapshotService.ts +++ b/src/notebooks/deepnote/snapshots/snapshotService.ts @@ -202,9 +202,17 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync projectId: string, projectName: string, projectData: DeepnoteFile, - notebookUri?: string + notebookUri?: string, + notebookId?: string ): Promise { - const prepared = await this.prepareSnapshotData(projectUri, projectId, projectName, projectData, notebookUri); + const prepared = await this.prepareSnapshotData( + projectUri, + projectId, + projectName, + projectData, + notebookUri, + notebookId + ); if (!prepared) { logger.debug(`[Snapshot] No changes detected, skipping snapshot creation`); @@ -214,7 +222,7 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync const { latestPath, content } = prepared; const timestamp = generateTimestamp(); - const timestampedPath = this.buildSnapshotPath(projectUri, projectId, projectName, timestamp); + const timestampedPath = this.buildSnapshotPath(projectUri, projectId, projectName, timestamp, notebookId); // Write to timestamped file first (safe - doesn't touch existing files) try { @@ -390,8 +398,8 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync }; } - async readSnapshot(projectId: string): Promise | undefined> { - logger.debug(`[Snapshot] readSnapshot called for projectId=${projectId}`); + async readSnapshot(projectId: string, notebookId?: string): Promise | undefined> { + logger.debug(`[Snapshot] readSnapshot called for projectId=${projectId}, notebookId=${notebookId}`); const workspaceFolders = workspace.workspaceFolders; if (!workspaceFolders || workspaceFolders.length === 0) { @@ -404,8 +412,10 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync logger.debug(`[Snapshot] Searching ${workspaceFolders.length} workspace folder(s) for snapshots`); - // 1. Try to find a 'latest' snapshot file - const latestGlob = `**/snapshots/*_${projectId}_latest.snapshot.deepnote`; + // 1. Try to find a 'latest' snapshot file (new format with notebookId first, then legacy) + const latestGlob = notebookId + ? `**/snapshots/*_${projectId}_${notebookId}_latest.snapshot.deepnote` + : `**/snapshots/*_${projectId}_latest.snapshot.deepnote`; for (const folder of workspaceFolders) { logger.debug(`[Snapshot] Searching for latest snapshot with glob: ${latestGlob} in ${folder.uri.path}`); @@ -430,7 +440,9 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync logger.debug(`[Snapshot] No latest snapshot found, looking for timestamped files`); // 2. Find timestamped snapshots across all workspace folders - const timestampedGlob = `**/snapshots/*_${projectId}_*.snapshot.deepnote`; + const timestampedGlob = notebookId + ? `**/snapshots/*_${projectId}_${notebookId}_*.snapshot.deepnote` + : `**/snapshots/*_${projectId}_*.snapshot.deepnote`; let allTimestampedFiles: Uri[] = []; for (const folder of workspaceFolders) { @@ -494,11 +506,14 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync projectUri: Uri, projectId: string, projectName: string, - variant: 'latest' | string + variant: 'latest' | string, + notebookId?: string ): Uri { const parentDir = Uri.joinPath(projectUri, '..'); const slug = slugifyProjectName(projectName); - const filename = `${slug}_${projectId}_${variant}.snapshot.deepnote`; + const filename = notebookId + ? `${slug}_${projectId}_${notebookId}_${variant}.snapshot.deepnote` + : `${slug}_${projectId}_${variant}.snapshot.deepnote`; return Uri.joinPath(parentDir, 'snapshots', filename); } @@ -800,7 +815,8 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync projectId, originalProject.project.name, snapshotProject, - notebookUri + notebookUri, + notebookId ); if (snapshotUri) { @@ -814,7 +830,8 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync projectId, originalProject.project.name, snapshotProject, - notebookUri + notebookUri, + notebookId ); if (snapshotUri) { @@ -866,12 +883,13 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync projectId: string, projectName: string, projectData: DeepnoteFile, - notebookUri?: string + notebookUri?: string, + notebookId?: string ): Promise<{ latestPath: Uri; content: Uint8Array } | undefined> { let latestPath: Uri; try { - latestPath = this.buildSnapshotPath(projectUri, projectId, projectName, 'latest'); + latestPath = this.buildSnapshotPath(projectUri, projectId, projectName, 'latest', notebookId); } catch (error) { if (error instanceof InvalidProjectNameError) { logger.warn('[Snapshot] Skipping snapshots due to invalid project name', error); @@ -1011,9 +1029,17 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync projectId: string, projectName: string, projectData: DeepnoteFile, - notebookUri?: string + notebookUri?: string, + notebookId?: string ): Promise { - const prepared = await this.prepareSnapshotData(projectUri, projectId, projectName, projectData, notebookUri); + const prepared = await this.prepareSnapshotData( + projectUri, + projectId, + projectName, + projectData, + notebookUri, + notebookId + ); if (!prepared) { logger.debug(`[Snapshot] No changes detected, skipping latest snapshot update`); diff --git a/src/notebooks/types.ts b/src/notebooks/types.ts index 2da356ee25..de3c3e664b 100644 --- a/src/notebooks/types.ts +++ b/src/notebooks/types.ts @@ -37,13 +37,10 @@ export interface ProjectIntegration { export const IDeepnoteNotebookManager = Symbol('IDeepnoteNotebookManager'); export interface IDeepnoteNotebookManager { - consumePendingNotebookResolution(projectId: string): string | undefined; - getCurrentNotebookId(projectId: string): string | undefined; getOriginalProject(projectId: string): DeepnoteProject | undefined; hasInitNotebookBeenRun(projectId: string): boolean; markInitNotebookAsRun(projectId: string): void; - queueNotebookResolution(projectId: string, notebookId: string): void; - storeOriginalProject(projectId: string, project: DeepnoteProject, notebookId: string): void; + storeOriginalProject(projectId: string, project: DeepnoteProject): void; updateOriginalProject(projectId: string, project: DeepnoteProject): void; /** From 1140fa0e6548996956efdb32bc7ed99cc2a5088c Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 16 Apr 2026 17:31:44 +0000 Subject: [PATCH 37/47] feat(deepnote): Enhance snapshot output update handling with notebook ID support - Added optional `notebookId` parameter to `executeSnapshotOutputUpdate` method to support notebook-scoped snapshots. - Updated logic to read snapshots based on the presence of `notebookId`, improving output refresh accuracy for open notebooks. - Enhanced unit tests to validate behavior when handling notebook-scoped snapshots, ensuring correct application of outputs based on matching notebook IDs. --- .../deepnote/deepnoteFileChangeWatcher.ts | 28 +++- .../deepnoteFileChangeWatcher.unit.test.ts | 149 ++++++++++++++++++ 2 files changed, 171 insertions(+), 6 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index 76fa5158b4..e883987d01 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -42,6 +42,11 @@ interface PendingOperation { type: OperationType; /** For snapshot-output-update: the project ID to read outputs from. */ projectId?: string; + /** + * For snapshot-output-update: notebook ID from the snapshot filename. + * Omitted for legacy project-scoped snapshot files. + */ + notebookId?: string; } /** @@ -234,7 +239,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic if (op.type === 'main-file-sync') { await this.executeMainFileSync(notebook, fileUri ?? notebook.uri.with({ query: '', fragment: '' })); } else if (op.type === 'snapshot-output-update' && op.projectId) { - await this.executeSnapshotOutputUpdate(notebook, op.projectId); + await this.executeSnapshotOutputUpdate(notebook, op.projectId, op.notebookId); } } catch (error) { logger.error(`[FileChangeWatcher] Operation ${op.type} failed for ${nbKey}`, error); @@ -282,7 +287,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic if (pending?.type === 'main-file-sync') { continue; } - this.pendingOperations.set(nbKey, { type: 'snapshot-output-update', projectId }); + this.pendingOperations.set(nbKey, { type: 'snapshot-output-update', projectId, notebookId }); void this.drainQueue(nbKey, notebook); } } @@ -413,12 +418,19 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic * Prefers the notebook execution API (outputs set this way respect transientOutputs * and do not mark the notebook dirty). Falls back to replaceCells when no kernel is active. */ - private async executeSnapshotOutputUpdate(notebook: NotebookDocument, projectId: string): Promise { + private async executeSnapshotOutputUpdate( + notebook: NotebookDocument, + projectId: string, + notebookId?: string + ): Promise { if (!this.snapshotService) { return; } - const snapshotOutputs = await this.snapshotService.readSnapshot(projectId); + const snapshotOutputs = + notebookId === undefined + ? await this.snapshotService.readSnapshot(projectId) + : await this.snapshotService.readSnapshot(projectId, notebookId); if (!snapshotOutputs || snapshotOutputs.size === 0) { return; } @@ -433,8 +445,12 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } const liveCells = notebook.getCells(); - const notebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; - const originalBlocks = notebookId ? notebookBlocksMap.get(notebookId) : undefined; + const docNotebookId = + notebook.metadata?.deepnoteNotebookId !== undefined && + typeof notebook.metadata.deepnoteNotebookId === 'string' + ? notebook.metadata.deepnoteNotebookId + : undefined; + const originalBlocks = docNotebookId ? notebookBlocksMap.get(docNotebookId) : undefined; // Collect cells that need output updates const cellUpdates: Array<{ diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 7877b95d04..2c31132d04 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -1606,6 +1606,155 @@ project: } }); + test('refreshes open notebook when only notebook-scoped latest snapshot exists on disk', async () => { + // Catches: queued snapshot work dropped the notebook ID from the filename and called readSnapshot(projectId), so notebook-scoped latest files returned no outputs and the editor never refreshed. + const scopedNotebookId = 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa'; + const snapshotUri = testFileUri('snapshots', `proj_project-1_${scopedNotebookId}_latest.snapshot.deepnote`); + + resetCalls(mockSnapshotService); + when(mockSnapshotService.isSnapshotsEnabled()).thenReturn(true); + when(mockSnapshotService.readSnapshot(anything(), anything())).thenCall((pid: string, nid?: string) => { + readSnapshotCallCount++; + if (typeof nid === 'undefined') { + return Promise.resolve(undefined); + } + if (pid === 'project-1' && nid === scopedNotebookId) { + return Promise.resolve(snapshotOutputs); + } + + return Promise.resolve(undefined); + }); + + const notebook = createMockNotebook({ + uri: testFileUri('scoped-snapshot-only.deepnote'), + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: scopedNotebookId + }, + cells: [ + { + metadata: { id: 'block-1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + snapshotApplyEditCount = 0; + snapshotOnDidChange.fire(snapshotUri); + + await waitFor(() => snapshotApplyEditCount > 0); + + assert.isAtLeast(readSnapshotCallCount, 1, 'readSnapshot should be called'); + assert.isAtLeast( + snapshotApplyEditCount, + 1, + 'applyEdit should refresh outputs from notebook-scoped snapshot' + ); + }); + + test('applies notebook-scoped latest snapshot only to the open notebook whose ID matches the filename', async () => { + // Catches: readSnapshot omitted the notebook ID from the snapshot path, so sibling notebooks in the same project could be updated from the wrong snapshot scope or the targeted notebook read no outputs. + const nb1Id = 'bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb'; + const nb2Id = 'cccccccc-cccc-4ccc-8ccc-cccccccccccc'; + + resetCalls(mockSnapshotService); + when(mockSnapshotService.isSnapshotsEnabled()).thenReturn(true); + + const readArgs: Array<{ projectId: string; notebookId?: string }> = []; + const nb1Outputs = new Map([ + [ + 'block-nb1', + [ + { + output_type: 'execute_result', + data: { 'text/plain': 'ForNb1Only' }, + execution_count: 1 + } as DeepnoteOutput + ] + ] + ]); + const nb2Outputs = new Map([ + [ + 'block-nb2', + [ + { + output_type: 'execute_result', + data: { 'text/plain': 'ForNb2Only' }, + execution_count: 1 + } as DeepnoteOutput + ] + ] + ]); + + when(mockSnapshotService.readSnapshot(anything(), anything())).thenCall((pid: string, nid?: string) => { + readSnapshotCallCount++; + readArgs.push({ projectId: pid, notebookId: nid }); + if (nid === nb1Id) { + return Promise.resolve(nb1Outputs); + } + if (nid === nb2Id) { + return Promise.resolve(nb2Outputs); + } + + return Promise.resolve(undefined); + }); + + when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); + + const uriNb1 = testFileUri('multi-scope-nb1.deepnote'); + const uriNb2 = testFileUri('multi-scope-nb2.deepnote'); + const notebook1 = createMockNotebook({ + uri: uriNb1, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: nb1Id + }, + cells: [ + { + metadata: { id: 'block-nb1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb1")', languageId: 'python' } + } + ] + }); + const notebook2 = createMockNotebook({ + uri: uriNb2, + metadata: { + deepnoteProjectId: 'project-1', + deepnoteNotebookId: nb2Id + }, + cells: [ + { + metadata: { id: 'block-nb2', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("nb2")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); + + const snapshotUri = testFileUri('snapshots', `proj_project-1_${nb1Id}_latest.snapshot.deepnote`); + + snapshotApplyEditCount = 0; + snapshotOnDidChange.fire(snapshotUri); + + await waitFor(() => snapshotApplyEditCount >= 2); + + assert.deepStrictEqual(readArgs, [{ projectId: 'project-1', notebookId: nb1Id }]); + assert.strictEqual( + snapshotApplyEditCount, + 2, + 'only the matching open notebook should receive snapshot edits' + ); + }); + suite('snapshot and deserialization interaction', () => { let interactionCaptures: SnapshotInteractionCapture[]; let snapshotApplyEditStub: sinon.SinonStub; From 966bf963200e172af13cc9a0e8adcf6a726be535 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 17 Apr 2026 06:01:01 +0000 Subject: [PATCH 38/47] feat(deepnote): Implement notebook ID support in project management - Enhanced `DeepnoteNotebookManager` to manage projects using a (projectId, notebookId) pair, allowing for distinct storage of projects across multiple notebooks. - Updated methods `getOriginalProject`, `storeOriginalProject`, and `updateOriginalProject` to accept an optional notebook ID, improving project data handling for split files. - Refactored related classes and tests to ensure compatibility with the new notebook ID structure, enhancing overall functionality and reliability. - Improved snapshot handling in `SnapshotService` and `DeepnoteSerializer` to utilize notebook IDs for accurate project retrieval and serialization. --- .../deepnote/deepnoteServerStarter.node.ts | 42 +++++------ .../deepnote/deepnoteExplorerView.ts | 5 ++ .../deepnoteExplorerView.unit.test.ts | 24 ++++-- .../deepnote/deepnoteFileChangeWatcher.ts | 14 ++-- .../deepnoteFileChangeWatcher.unit.test.ts | 30 ++++---- .../deepnoteKernelAutoSelector.node.ts | 3 +- .../deepnote/deepnoteNotebookManager.ts | 74 ++++++++++++++----- .../deepnoteNotebookManager.unit.test.ts | 73 +++++++++++++----- src/notebooks/deepnote/deepnoteSerializer.ts | 28 ++++--- .../deepnote/deepnoteSerializer.unit.test.ts | 34 ++++----- src/notebooks/deepnote/deepnoteTreeItem.ts | 5 +- .../deepnote/snapshots/snapshotService.ts | 18 ++--- src/notebooks/types.ts | 6 +- 13 files changed, 222 insertions(+), 134 deletions(-) diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 0ea58022cd..8c1d51717c 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -343,11 +343,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension this.serverOutputByFile.delete(fileKey); - const disposables = this.disposablesByFile.get(fileKey); - if (disposables) { - disposables.forEach((d) => d.dispose()); - this.disposablesByFile.delete(fileKey); - } + this.disposeOutputListeners(fileKey); } /** @@ -402,16 +398,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension */ private monitorServerOutput(fileKey: string, serverInfo: DeepnoteServerInfo): void { const proc = serverInfo.process; - const existing = this.disposablesByFile.get(fileKey); - if (existing) { - for (const d of existing) { - try { - d.dispose(); - } catch (ex) { - logger.warn(`Error disposing listener for ${fileKey}`, ex); - } - } - } + this.disposeOutputListeners(fileKey); const disposables: IDisposable[] = []; this.disposablesByFile.set(fileKey, disposables); @@ -456,6 +443,22 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } + /** + * Dispose all output listeners registered for a given file key. + * Per-listener try/catch ensures one failure doesn't prevent others from being disposed. + */ + private disposeOutputListeners(fileKey: string): void { + const disposables = this.disposablesByFile.get(fileKey) ?? []; + for (const d of disposables) { + try { + d.dispose(); + } catch (ex) { + logger.warn(`Error disposing output listener for ${fileKey}`, ex); + } + } + this.disposablesByFile.delete(fileKey); + } + public async dispose(): Promise { logger.info('Disposing DeepnoteServerStarter - stopping all servers...'); @@ -495,12 +498,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension await this.deleteLockFile(pid); } - for (const [fileKey, disposables] of this.disposablesByFile.entries()) { - try { - disposables.forEach((d) => d.dispose()); - } catch (ex) { - logger.error(`Error disposing resources for ${fileKey}`, ex); - } + // Snapshot the keys first: disposeOutputListeners mutates disposablesByFile via .delete(). + for (const fileKey of Array.from(this.disposablesByFile.keys())) { + this.disposeOutputListeners(fileKey); } this.disposablesByFile.clear(); diff --git a/src/notebooks/deepnote/deepnoteExplorerView.ts b/src/notebooks/deepnote/deepnoteExplorerView.ts index a8fc8173a5..eac7c6fb3d 100644 --- a/src/notebooks/deepnote/deepnoteExplorerView.ts +++ b/src/notebooks/deepnote/deepnoteExplorerView.ts @@ -29,6 +29,11 @@ export class DeepnoteExplorerView { this.treeDataProvider = new DeepnoteTreeDataProvider(logger); } + public dispose(): void { + this.treeView?.dispose(); + this.treeDataProvider.dispose(); + } + public activate(): void { this.treeView = window.createTreeView('deepnoteExplorer', { treeDataProvider: this.treeDataProvider, diff --git a/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts b/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts index a1f12f9604..f3afc275f3 100644 --- a/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts @@ -53,6 +53,10 @@ suite('DeepnoteExplorerView', () => { explorerView = new DeepnoteExplorerView(mockExtensionContext, mockLogger); }); + teardown(() => { + explorerView.dispose(); + }); + suite('constructor', () => { test('should create instance with extension context', () => { assert.isDefined(explorerView); @@ -189,13 +193,18 @@ suite('DeepnoteExplorerView', () => { const view1 = new DeepnoteExplorerView(context1, logger1); const view2 = new DeepnoteExplorerView(context2, logger2); - // Verify each view has its own context - assert.strictEqual((view1 as any).extensionContext, context1); - assert.strictEqual((view2 as any).extensionContext, context2); - assert.notStrictEqual((view1 as any).extensionContext, (view2 as any).extensionContext); - - // Verify views are independent instances - assert.notStrictEqual(view1, view2); + try { + // Verify each view has its own context + assert.strictEqual((view1 as any).extensionContext, context1); + assert.strictEqual((view2 as any).extensionContext, context2); + assert.notStrictEqual((view1 as any).extensionContext, (view2 as any).extensionContext); + + // Verify views are independent instances + assert.notStrictEqual(view1, view2); + } finally { + view1.dispose(); + view2.dispose(); + } }); test('should maintain component references', () => { @@ -232,6 +241,7 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { }); teardown(() => { + explorerView.dispose(); sandbox.restore(); uuidStubs.forEach((stub) => stub.restore()); uuidStubs = []; diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index e883987d01..38a9c0f046 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -435,8 +435,13 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic return; } + const metadataNotebookId = + typeof notebook.metadata?.deepnoteNotebookId === 'string' + ? notebook.metadata.deepnoteNotebookId + : undefined; + // Look up original project blocks for fallback block ID resolution - const originalProject = this.notebookManager.getOriginalProject(projectId); + const originalProject = this.notebookManager.getOriginalProject(projectId, notebookId ?? metadataNotebookId); const notebookBlocksMap = new Map(); if (originalProject) { for (const nb of originalProject.project.notebooks) { @@ -445,12 +450,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } const liveCells = notebook.getCells(); - const docNotebookId = - notebook.metadata?.deepnoteNotebookId !== undefined && - typeof notebook.metadata.deepnoteNotebookId === 'string' - ? notebook.metadata.deepnoteNotebookId - : undefined; - const originalBlocks = docNotebookId ? notebookBlocksMap.get(docNotebookId) : undefined; + const originalBlocks = metadataNotebookId ? notebookBlocksMap.get(metadataNotebookId) : undefined; // Collect cells that need output updates const cellUpdates: Array<{ diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 2c31132d04..0afce2b407 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -162,8 +162,8 @@ suite('DeepnoteFileChangeWatcher', () => { mockDisposables = []; mockedNotebookManager = mock(); - when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(validProject); - when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); + when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(validProject); + when(mockedNotebookManager.updateOriginalProject(anything(), anything(), anything())).thenReturn(); mockNotebookManager = instance(mockedNotebookManager); // Set up FileSystemWatcher mock @@ -563,7 +563,7 @@ project: this.timeout(15_000); const uri = testFileUri('self-write-leak.deepnote'); - when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); + when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(multiNotebookProject); // Initial state: editor content matches disk — use the real converter const converter = new DeepnoteDataConverter(); @@ -1086,7 +1086,7 @@ project: test('should apply snapshot outputs using original blocks when metadata is lost', async () => { // Create a mock notebook manager that returns an original project const mockedManager = mock(); - when(mockedManager.getOriginalProject('project-1')).thenReturn({ + when(mockedManager.getOriginalProject('project-1', anything())).thenReturn({ version: '1.0', metadata: { createdAt: '2025-01-01T00:00:00Z' }, project: { @@ -1353,7 +1353,7 @@ project: when(nfSnapshotService.onFileWritten(anything())).thenReturn({ dispose: () => {} } as Disposable); const nfManager = mock(); - when(nfManager.getOriginalProject(anything())).thenReturn(undefined); + when(nfManager.getOriginalProject(anything(), anything())).thenReturn(undefined); const nfWatcher = new DeepnoteFileChangeWatcher( noFallbackDisposables, @@ -1555,8 +1555,8 @@ project: } as any); const mockedManagerEx = mock(); - when(mockedManagerEx.getOriginalProject(anything())).thenReturn(validProject); - when(mockedManagerEx.updateOriginalProject(anything(), anything())).thenReturn(); + when(mockedManagerEx.getOriginalProject(anything(), anything())).thenReturn(validProject); + when(mockedManagerEx.updateOriginalProject(anything(), anything(), anything())).thenReturn(); const exMdWatcher = new DeepnoteFileChangeWatcher( exMdDisposables, @@ -1703,7 +1703,7 @@ project: return Promise.resolve(undefined); }); - when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); + when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(multiNotebookProject); const uriNb1 = testFileUri('multi-scope-nb1.deepnote'); const uriNb2 = testFileUri('multi-scope-nb2.deepnote'); @@ -1763,8 +1763,8 @@ project: interactionCaptures = []; reset(mockedNotebookManager); - when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(validProject); - when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); + when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(validProject); + when(mockedNotebookManager.updateOriginalProject(anything(), anything(), anything())).thenReturn(); resetCalls(mockedNotebookManager); snapshotApplyEditStub = sinon.stub(snapshotWatcher, 'applyNotebookEdits').callsFake(async function ( @@ -1809,7 +1809,7 @@ project: }); test('snapshot change with multi-notebook project applies only matching block outputs per notebook', async () => { - when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); + when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(multiNotebookProject); const multiOutputs = new Map([ [ @@ -2110,7 +2110,7 @@ project: // Multi-notebook test removed — multi-notebook support has been replaced by auto-splitting into separate files test.skip('multi-notebook: snapshot outputs then external YAML update keeps per-notebook sources', async function () { this.timeout(12_000); - when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); + when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(multiNotebookProject); const multiOutputs = new Map([ [ @@ -2238,7 +2238,7 @@ project: }); test('snapshot outputs for sibling notebook blocks do not leak into a single open notebook', async () => { - when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); + when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(multiNotebookProject); const multiOutputs = new Map([ [ @@ -2307,8 +2307,8 @@ project: setup(() => { reset(mockedNotebookManager); - when(mockedNotebookManager.getOriginalProject(anything())).thenReturn(multiNotebookProject); - when(mockedNotebookManager.updateOriginalProject(anything(), anything())).thenReturn(); + when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(multiNotebookProject); + when(mockedNotebookManager.updateOriginalProject(anything(), anything(), anything())).thenReturn(); resetCalls(mockedNotebookManager); workspaceSetCaptures = []; sinon.stub(watcher, 'applyNotebookEdits' as any).callsFake(async (...args: unknown[]) => { diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 63e7129200..d1330ae14f 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -670,8 +670,9 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Prepare init notebook execution const projectId = notebook.metadata?.deepnoteProjectId; + const notebookIdForProject = notebook.metadata?.deepnoteNotebookId; const project = projectId - ? (this.notebookManager.getOriginalProject(projectId) as DeepnoteFile | undefined) + ? (this.notebookManager.getOriginalProject(projectId, notebookIdForProject) as DeepnoteFile | undefined) : undefined; if (project) { diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.ts b/src/notebooks/deepnote/deepnoteNotebookManager.ts index 26dad36a01..d515fd8bf2 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.ts @@ -6,19 +6,41 @@ import type { DeepnoteProject } from '../../platform/deepnote/deepnoteTypes'; /** * Centralized manager for tracking Deepnote project state. * Manages per-project data caching and init notebook tracking. + * + * Cache keys are (projectId, notebookId) pairs so that split files sharing + * a projectId but owning distinct notebooks don't clobber each other. */ @injectable() export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { - private readonly originalProjects = new Map(); + private readonly originalProjects = new Map>(); private readonly projectsWithInitNotebookRun = new Set(); /** - * Retrieves the original project data for a given project ID. + * Retrieves the original project data for a given project ID and optional notebook ID. + * When `notebookId` is provided and matches a cached entry, that exact entry is returned. + * Otherwise (or when no exact match exists), returns any cached entry for the projectId + * to support project-wide reads (e.g. integrations list, project name) where any + * cached notebook's project snapshot is equivalent. * @param projectId Project identifier + * @param notebookId Optional notebook identifier * @returns Original project data or undefined if not found */ - getOriginalProject(projectId: string): DeepnoteProject | undefined { - return this.originalProjects.get(projectId); + getOriginalProject(projectId: string, notebookId?: string): DeepnoteProject | undefined { + const notebookMap = this.originalProjects.get(projectId); + + if (!notebookMap) { + return undefined; + } + + if (notebookId !== undefined) { + const exactMatch = notebookMap.get(notebookId); + + if (exactMatch) { + return exactMatch; + } + } + + return notebookMap.values().next().value; } /** @@ -39,45 +61,63 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { } /** - * Stores the original project data. + * Stores the original project data for the given (projectId, notebookId) pair. * This is used during deserialization to cache project data. * @param projectId Project identifier + * @param notebookId Notebook identifier within the project * @param project Original project data to store */ - storeOriginalProject(projectId: string, project: DeepnoteProject): void { + storeOriginalProject(projectId: string, notebookId: string, project: DeepnoteProject): void { const clonedProject = structuredClone(project); - this.originalProjects.set(projectId, clonedProject); + let notebookMap = this.originalProjects.get(projectId); + + if (!notebookMap) { + notebookMap = new Map(); + this.originalProjects.set(projectId, notebookMap); + } + + notebookMap.set(notebookId, clonedProject); } /** - * Updates the stored project data. + * Updates the stored project data for the given (projectId, notebookId) pair. * Used during serialization where we need to cache the updated project state. * @param projectId Project identifier + * @param notebookId Notebook identifier within the project * @param project Updated project data to store */ - updateOriginalProject(projectId: string, project: DeepnoteProject): void { + updateOriginalProject(projectId: string, notebookId: string, project: DeepnoteProject): void { const clonedProject = structuredClone(project); - this.originalProjects.set(projectId, clonedProject); + let notebookMap = this.originalProjects.get(projectId); + + if (!notebookMap) { + notebookMap = new Map(); + this.originalProjects.set(projectId, notebookMap); + } + + notebookMap.set(notebookId, clonedProject); } /** * Updates the integrations list in the project data. - * This modifies the stored project to reflect changes in configured integrations. + * This modifies every cached entry for the project to reflect changes in configured integrations. * * @param projectId - Project identifier * @param integrations - Array of integration metadata to store in the project - * @returns `true` if the project was found and updated successfully, `false` if the project does not exist + * @returns `true` if at least one cached entry was found and updated, `false` if no entries exist */ updateProjectIntegrations(projectId: string, integrations: ProjectIntegration[]): boolean { - const project = this.originalProjects.get(projectId); + const notebookMap = this.originalProjects.get(projectId); - if (!project) { + if (!notebookMap || notebookMap.size === 0) { return false; } - const updatedProject = structuredClone(project); - updatedProject.project.integrations = integrations; - this.originalProjects.set(projectId, updatedProject); + for (const [notebookId, project] of notebookMap.entries()) { + const updatedProject = structuredClone(project); + updatedProject.project.integrations = integrations; + notebookMap.set(notebookId, updatedProject); + } return true; } diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts index 7d3b226843..2ebd2c7c8a 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts @@ -33,7 +33,7 @@ suite('DeepnoteNotebookManager', () => { }); test('should return original project after storing', () => { - manager.storeOriginalProject('project-123', mockProject); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); const result = manager.getOriginalProject('project-123'); @@ -43,7 +43,7 @@ suite('DeepnoteNotebookManager', () => { suite('storeOriginalProject', () => { test('should store project data', () => { - manager.storeOriginalProject('project-123', mockProject); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); const storedProject = manager.getOriginalProject('project-123'); @@ -59,13 +59,30 @@ suite('DeepnoteNotebookManager', () => { } }; - manager.storeOriginalProject('project-123', mockProject); - manager.storeOriginalProject('project-123', updatedProject); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); + manager.storeOriginalProject('project-123', 'notebook-1', updatedProject); const storedProject = manager.getOriginalProject('project-123'); assert.deepStrictEqual(storedProject, updatedProject); }); + + test('stores split files with the same projectId under distinct notebook IDs', () => { + const projectA: DeepnoteProject = { + ...mockProject, + project: { ...mockProject.project, name: 'File A' } + }; + const projectB: DeepnoteProject = { + ...mockProject, + project: { ...mockProject.project, name: 'File B' } + }; + + manager.storeOriginalProject('project-123', 'notebook-a', projectA); + manager.storeOriginalProject('project-123', 'notebook-b', projectB); + + assert.strictEqual(manager.getOriginalProject('project-123', 'notebook-a')?.project.name, 'File A'); + assert.strictEqual(manager.getOriginalProject('project-123', 'notebook-b')?.project.name, 'File B'); + }); }); suite('updateOriginalProject', () => { @@ -78,8 +95,8 @@ suite('DeepnoteNotebookManager', () => { } }; - manager.storeOriginalProject('project-123', mockProject); - manager.updateOriginalProject('project-123', updatedProject); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); + manager.updateOriginalProject('project-123', 'notebook-1', updatedProject); const storedProject = manager.getOriginalProject('project-123'); @@ -95,8 +112,8 @@ suite('DeepnoteNotebookManager', () => { } }; - manager.storeOriginalProject('project-123', mockProject); - manager.updateOriginalProject('project-123', updatedProject); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); + manager.updateOriginalProject('project-123', 'notebook-1', updatedProject); updatedProject.project.name = 'After Mutation'; @@ -115,9 +132,9 @@ suite('DeepnoteNotebookManager', () => { project: { ...mockProject.project, name: 'Second Update' } }; - manager.storeOriginalProject('project-123', mockProject); - manager.updateOriginalProject('project-123', firstUpdate); - manager.updateOriginalProject('project-123', secondUpdate); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); + manager.updateOriginalProject('project-123', 'notebook-1', firstUpdate); + manager.updateOriginalProject('project-123', 'notebook-1', secondUpdate); assert.strictEqual(manager.getOriginalProject('project-123')?.project.name, 'Second Update'); }); @@ -128,7 +145,7 @@ suite('DeepnoteNotebookManager', () => { project: { ...mockProject.project, name: 'No Notebook Id Yet' } }; - manager.updateOriginalProject('project-123', projectOnly); + manager.updateOriginalProject('project-123', 'notebook-1', projectOnly); assert.deepStrictEqual(manager.getOriginalProject('project-123'), projectOnly); }); @@ -136,7 +153,7 @@ suite('DeepnoteNotebookManager', () => { suite('updateProjectIntegrations', () => { test('should update integrations list for existing project and return true', () => { - manager.storeOriginalProject('project-123', mockProject); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); const integrations: ProjectIntegration[] = [ { id: 'int-1', name: 'PostgreSQL', type: 'pgsql' }, @@ -160,7 +177,7 @@ suite('DeepnoteNotebookManager', () => { } }; - manager.storeOriginalProject('project-123', projectWithIntegrations); + manager.storeOriginalProject('project-123', 'notebook-1', projectWithIntegrations); const newIntegrations: ProjectIntegration[] = [ { id: 'new-int-1', name: 'New Integration 1', type: 'pgsql' }, @@ -184,7 +201,7 @@ suite('DeepnoteNotebookManager', () => { } }; - manager.storeOriginalProject('project-123', projectWithIntegrations); + manager.storeOriginalProject('project-123', 'notebook-1', projectWithIntegrations); const result = manager.updateProjectIntegrations('project-123', []); @@ -206,7 +223,7 @@ suite('DeepnoteNotebookManager', () => { }); test('should preserve other project properties and return true', () => { - manager.storeOriginalProject('project-123', mockProject); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); const integrations: ProjectIntegration[] = [{ id: 'int-1', name: 'PostgreSQL', type: 'pgsql' }]; @@ -221,9 +238,25 @@ suite('DeepnoteNotebookManager', () => { assert.deepStrictEqual(updatedProject?.metadata, mockProject.metadata); }); + test('updateProjectIntegrations updates every cached entry for a projectId', () => { + manager.storeOriginalProject('project-123', 'notebook-a', mockProject); + manager.storeOriginalProject('project-123', 'notebook-b', mockProject); + + const integrations: ProjectIntegration[] = [{ id: 'int-1', name: 'PG', type: 'pgsql' }]; + assert.strictEqual(manager.updateProjectIntegrations('project-123', integrations), true); + + assert.deepStrictEqual( + manager.getOriginalProject('project-123', 'notebook-a')?.project.integrations, + integrations + ); + assert.deepStrictEqual( + manager.getOriginalProject('project-123', 'notebook-b')?.project.integrations, + integrations + ); + }); + test('should update integrations when project was stored via updateOriginalProject and return true', () => { - // Use updateOriginalProject which doesn't set currentNotebookId - manager.updateOriginalProject('project-123', mockProject); + manager.updateOriginalProject('project-123', 'notebook-1', mockProject); const integrations: ProjectIntegration[] = [ { id: 'int-1', name: 'PostgreSQL', type: 'pgsql' }, @@ -266,8 +299,8 @@ suite('DeepnoteNotebookManager', () => { suite('integration scenarios', () => { test('should handle complete workflow for multiple projects', () => { - manager.storeOriginalProject('project-1', mockProject); - manager.storeOriginalProject('project-2', mockProject); + manager.storeOriginalProject('project-1', 'notebook-1', mockProject); + manager.storeOriginalProject('project-2', 'notebook-1', mockProject); assert.deepStrictEqual(manager.getOriginalProject('project-1'), mockProject); assert.deepStrictEqual(manager.getOriginalProject('project-2'), mockProject); diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 83d89ec145..fe3207a2e1 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -150,7 +150,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { ); } - this.notebookManager.storeOriginalProject(deepnoteFile.project.id, deepnoteFile); + this.notebookManager.storeOriginalProject(deepnoteFile.project.id, selectedNotebook.id, deepnoteFile); logger.debug(`DeepnoteSerializer: Stored project ${projectId} in notebook manager`); return { @@ -205,10 +205,20 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { logger.debug(`SerializeNotebook: Project ID: ${projectId}`); + const notebookId = data.metadata?.deepnoteNotebookId; + + if (!notebookId) { + throw new Error('Cannot determine which notebook to save'); + } + + logger.debug(`SerializeNotebook: Notebook ID: ${notebookId}`); + // Clone the project before modifying to prevent state corruption // This is critical for multi-notebook projects where the stored project // is shared between notebook serialization calls - const storedProject = this.notebookManager.getOriginalProject(projectId) as DeepnoteFile | undefined; + const storedProject = this.notebookManager.getOriginalProject(projectId, notebookId) as + | DeepnoteFile + | undefined; if (!storedProject) { throw new Error('Original Deepnote project not found. Cannot save changes.'); @@ -218,14 +228,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { logger.debug('SerializeNotebook: Got and cloned original project'); - const notebookId = data.metadata?.deepnoteNotebookId; - - if (!notebookId) { - throw new Error('Cannot determine which notebook to save'); - } - - logger.debug(`SerializeNotebook: Notebook ID: ${notebookId}`); - const notebook = originalProject.project.notebooks.find((nb: { id: string }) => nb.id === notebookId); if (!notebook) { @@ -306,11 +308,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } // Store the updated project back so subsequent saves start from correct state. - // Use updateOriginalProject (not storeOriginalProject) to avoid overwriting - // currentNotebookId — when multiple notebooks share the same file, changing - // currentNotebookId here would cause VS Code's follow-up deserialize calls - // for other open notebooks to resolve to the wrong notebook. - this.notebookManager.updateOriginalProject(projectId, originalProject); + this.notebookManager.updateOriginalProject(projectId, notebookId, originalProject); logger.debug('SerializeNotebook: Serializing to YAML'); diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index a5028c27a6..c865dcbdf5 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -180,7 +180,7 @@ project: test('should serialize notebook when original project exists', async () => { // First store the original project - manager.storeOriginalProject('project-123', mockProject); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); const mockNotebookData = { cells: [ @@ -208,7 +208,7 @@ project: }); test('should throw error when metadata notebook ID is missing', async () => { - manager.storeOriginalProject('project-123', mockProject); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); const mockNotebookData = { cells: [ @@ -328,7 +328,7 @@ project: } }; - manager.storeOriginalProject('project-circular', projectWithCircularRef); + manager.storeOriginalProject('project-circular', 'notebook-1', projectWithCircularRef); const notebookData = { cells: [ @@ -396,7 +396,7 @@ project: }; // Store the project - manager.storeOriginalProject('project-id-test', projectData); + manager.storeOriginalProject('project-id-test', 'notebook-1', projectData); // Create cells with the EXACT metadata structure that deserializeNotebook produces // This simulates what VS Code should preserve from deserialization @@ -482,7 +482,7 @@ project: } }; - manager.storeOriginalProject('project-recover-ids', projectData); + manager.storeOriginalProject('project-recover-ids', 'notebook-1', projectData); // Cells WITHOUT id metadata (simulating what VS Code might provide if it strips metadata) // But content matches the original block @@ -549,7 +549,7 @@ project: } }; - manager.storeOriginalProject('project-new-content', projectData); + manager.storeOriginalProject('project-new-content', 'notebook-1', projectData); // Cell with different content than any original block const notebookData = { @@ -1102,7 +1102,7 @@ project: } }; - manager.storeOriginalProject('project-snapshot-hash', projectData); + manager.storeOriginalProject('project-snapshot-hash', 'notebook-1', projectData); const notebookData = { cells: [ @@ -1177,13 +1177,13 @@ project: }; // Serialize twice - manager.storeOriginalProject('project-deterministic', structuredClone(projectData)); + manager.storeOriginalProject('project-deterministic', 'notebook-1', structuredClone(projectData)); const result1 = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed1 = parseYaml(new TextDecoder().decode(result1)) as DeepnoteFile & { metadata: { snapshotHash?: string }; }; - manager.storeOriginalProject('project-deterministic', structuredClone(projectData)); + manager.storeOriginalProject('project-deterministic', 'notebook-1', structuredClone(projectData)); const result2 = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed2 = parseYaml(new TextDecoder().decode(result2)) as DeepnoteFile & { metadata: { snapshotHash?: string }; @@ -1278,7 +1278,7 @@ project: // Serialize 5 times and collect all hashes for (let i = 0; i < 5; i++) { - manager.storeOriginalProject('project-multi-serialize', structuredClone(projectData)); + manager.storeOriginalProject('project-multi-serialize', 'notebook-1', structuredClone(projectData)); const result = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed = parseYaml(new TextDecoder().decode(result)) as DeepnoteFile & { metadata: { snapshotHash?: string }; @@ -1327,7 +1327,7 @@ project: } }; - manager.storeOriginalProject('project-content-change', projectData1); + manager.storeOriginalProject('project-content-change', 'notebook-1', projectData1); const notebookData1 = { cells: [ @@ -1405,7 +1405,7 @@ project: } }; - manager.storeOriginalProject('project-version-change', projectData1); + manager.storeOriginalProject('project-version-change', 'notebook-1', projectData1); const notebookData = { cells: [ @@ -1429,7 +1429,7 @@ project: // Change version const projectData2: DeepnoteFile = { ...structuredClone(projectData1), version: '2.0' }; - manager.storeOriginalProject('project-version-change', projectData2); + manager.storeOriginalProject('project-version-change', 'notebook-1', projectData2); const result2 = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed2 = parseYaml(new TextDecoder().decode(result2)) as DeepnoteFile & { @@ -1471,7 +1471,7 @@ project: } }; - manager.storeOriginalProject('project-integrations-change', projectData1); + manager.storeOriginalProject('project-integrations-change', 'notebook-1', projectData1); const notebookData = { cells: [ @@ -1496,7 +1496,7 @@ project: // Add integrations const projectData2 = structuredClone(projectData1); projectData2.project.integrations = [{ id: 'int-1', name: 'PostgreSQL', type: 'postgres' }]; - manager.storeOriginalProject('project-integrations-change', projectData2); + manager.storeOriginalProject('project-integrations-change', 'notebook-1', projectData2); const result2 = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed2 = parseYaml(new TextDecoder().decode(result2)) as DeepnoteFile & { @@ -1538,7 +1538,7 @@ project: } }; - manager.storeOriginalProject('project-env-hash', projectData1); + manager.storeOriginalProject('project-env-hash', 'notebook-1', projectData1); const notebookData = { cells: [ @@ -1563,7 +1563,7 @@ project: // Add environment hash const projectData2 = structuredClone(projectData1); projectData2.environment = { hash: 'env-hash-123' }; - manager.storeOriginalProject('project-env-hash', projectData2); + manager.storeOriginalProject('project-env-hash', 'notebook-1', projectData2); const result2 = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed2 = parseYaml(new TextDecoder().decode(result2)) as DeepnoteFile & { diff --git a/src/notebooks/deepnote/deepnoteTreeItem.ts b/src/notebooks/deepnote/deepnoteTreeItem.ts index 15577142da..5679334a64 100644 --- a/src/notebooks/deepnote/deepnoteTreeItem.ts +++ b/src/notebooks/deepnote/deepnoteTreeItem.ts @@ -1,6 +1,7 @@ import { TreeItem, TreeItemCollapsibleState, ThemeIcon } from 'vscode'; import type { DeepnoteProject, DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; +import { basename } from '../../platform/vscode-path/path'; /** * Represents different types of items in the Deepnote tree view @@ -75,7 +76,7 @@ export class DeepnoteTreeItem extends TreeItem { // getLabel() inline if (this.type === DeepnoteTreeItemType.ProjectFile) { const project = this.data as DeepnoteProject; - const fileName = this.context.filePath.split('/').pop() ?? ''; + const fileName = basename(this.context.filePath); this.label = fileName || project.project.name || 'Untitled Project'; } else { const notebook = this.data as DeepnoteNotebook; @@ -138,7 +139,7 @@ export class DeepnoteTreeItem extends TreeItem { if (this.type === DeepnoteTreeItemType.ProjectFile) { const project = this.data as DeepnoteProject; - const fileName = this.context.filePath.split('/').pop() ?? ''; + const fileName = basename(this.context.filePath); this.label = fileName || project.project.name || 'Untitled Project'; this.tooltip = `Deepnote Project: ${project.project.name}\nFile: ${this.context.filePath}`; diff --git a/src/notebooks/deepnote/snapshots/snapshotService.ts b/src/notebooks/deepnote/snapshots/snapshotService.ts index 4ed8095628..c20f5948f3 100644 --- a/src/notebooks/deepnote/snapshots/snapshotService.ts +++ b/src/notebooks/deepnote/snapshots/snapshotService.ts @@ -754,26 +754,26 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync return; } - const originalProject = this.notebookManager?.getOriginalProject(projectId); + const notebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; - if (!originalProject) { - logger.warn(`[Snapshot] No original project found for ${projectId}`); + if (!notebookId) { + logger.warn(`[Snapshot] No notebook ID in notebook metadata`); return; } - const projectUri = this.findProjectUriFromId(projectId); + const originalProject = this.notebookManager?.getOriginalProject(projectId, notebookId); - if (!projectUri) { - logger.warn(`[Snapshot] Could not find project URI for ${projectId}`); + if (!originalProject) { + logger.warn(`[Snapshot] No original project found for ${projectId}`); return; } - const notebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; + const projectUri = this.findProjectUriFromId(projectId); - if (!notebookId) { - logger.warn(`[Snapshot] No notebook ID in notebook metadata`); + if (!projectUri) { + logger.warn(`[Snapshot] Could not find project URI for ${projectId}`); return; } diff --git a/src/notebooks/types.ts b/src/notebooks/types.ts index de3c3e664b..a8823edb95 100644 --- a/src/notebooks/types.ts +++ b/src/notebooks/types.ts @@ -37,11 +37,11 @@ export interface ProjectIntegration { export const IDeepnoteNotebookManager = Symbol('IDeepnoteNotebookManager'); export interface IDeepnoteNotebookManager { - getOriginalProject(projectId: string): DeepnoteProject | undefined; + getOriginalProject(projectId: string, notebookId?: string): DeepnoteProject | undefined; hasInitNotebookBeenRun(projectId: string): boolean; markInitNotebookAsRun(projectId: string): void; - storeOriginalProject(projectId: string, project: DeepnoteProject): void; - updateOriginalProject(projectId: string, project: DeepnoteProject): void; + storeOriginalProject(projectId: string, notebookId: string, project: DeepnoteProject): void; + updateOriginalProject(projectId: string, notebookId: string, project: DeepnoteProject): void; /** * Updates the integrations list in the project data. From 367db44599a08a9a4083bf5c26016ee6aee33ca4 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 17 Apr 2026 07:29:53 +0000 Subject: [PATCH 39/47] Fix test --- .../deepnote/deepnoteFileChangeWatcher.unit.test.ts | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 0afce2b407..c122bbabc7 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -508,7 +508,8 @@ project: }); test('should not suppress real changes after auto-save', async function () { - this.timeout(10_000); + // First reload (debounce + I/O) + second waitFor can exceed 10s on slow CI. + this.timeout(20_000); const uri = testFileUri('test.deepnote'); // First change: notebook has no cells, YAML has one cell -> different -> reload @@ -527,13 +528,16 @@ project: setupMockFs(validYaml); onDidChangeFile.fire(uri); - await waitFor(() => saveCount >= 1); + await waitFor(() => saveCount >= 1 && applyEditCount >= 1); // The first reload sets 2 self-write markers (writeFile + save). // Consume them both with simulated fs events. onDidChangeFile.fire(uri); onDidChangeFile.fire(uri); + // Let debounce/async work from the first reload settle (matches sibling regression test). + await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); + // Second real external change: use different YAML content const changedYaml = ` version: '1.0.0' @@ -554,7 +558,7 @@ project: `; setupMockFs(changedYaml); onDidChangeFile.fire(uri); - await waitFor(() => applyEditCount >= 2, waitForTimeoutMs); + await waitFor(() => applyEditCount >= 2, 15_000); assert.isAtLeast(applyEditCount, 2, 'applyEdit should be called for both external changes'); }); From b73d313c9fe807e69cef574d75e75090f8468f25 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 17 Apr 2026 08:59:23 +0000 Subject: [PATCH 40/47] feat(deepnote): Refactor notebook file handling and introduce new factory methods - Added `deepnoteNotebookFileFactory` to encapsulate logic for creating single notebook files and generating sibling file URIs. - Updated `DeepnoteAutoSplitter` and `DeepnoteExplorerView` to utilize the new factory methods, improving code organization and readability. - Enhanced snapshot hash computation and slugification logic for notebook names, ensuring consistent file naming and metadata handling. - Introduced unit tests for the new factory methods to validate functionality and maintain code quality. --- .../deepnote/deepnoteAutoSplitter.ts | 84 +---- .../deepnote/deepnoteExplorerView.ts | 121 ++++--- .../deepnoteExplorerView.unit.test.ts | 321 ++++++++++++++++-- .../deepnote/deepnoteNotebookFileFactory.ts | 118 +++++++ .../deepnoteNotebookFileFactory.unit.test.ts | 305 +++++++++++++++++ 5 files changed, 801 insertions(+), 148 deletions(-) create mode 100644 src/notebooks/deepnote/deepnoteNotebookFileFactory.ts create mode 100644 src/notebooks/deepnote/deepnoteNotebookFileFactory.unit.test.ts diff --git a/src/notebooks/deepnote/deepnoteAutoSplitter.ts b/src/notebooks/deepnote/deepnoteAutoSplitter.ts index 31b88b6cd0..caeaf655ed 100644 --- a/src/notebooks/deepnote/deepnoteAutoSplitter.ts +++ b/src/notebooks/deepnote/deepnoteAutoSplitter.ts @@ -1,8 +1,13 @@ import { serializeDeepnoteFile, type DeepnoteFile } from '@deepnote/blocks'; import { l10n, RelativePattern, Uri, window, workspace } from 'vscode'; -import { computeHash } from '../../platform/common/crypto'; import { logger } from '../../platform/logging'; +import { + buildSingleNotebookFile, + computeSnapshotHash, + getFileStem, + slugifyNotebookNameOrFallback +} from './deepnoteNotebookFileFactory'; import { slugifyProjectName } from './snapshots/snapshotFiles'; /** @@ -39,34 +44,16 @@ export class DeepnoteAutoSplitter { ); const parentDir = Uri.joinPath(fileUri, '..'); - const originalStem = this.getFileStem(fileUri); + const originalStem = getFileStem(fileUri); const newFiles: Uri[] = []; // Create a new file for each extra notebook for (const notebook of extraNotebooks) { - const notebookSlug = this.slugifyNotebookName(notebook.name); + const notebookSlug = slugifyNotebookNameOrFallback(notebook.name); const newFileName = `${originalStem}_${notebookSlug}.deepnote`; const newFileUri = Uri.joinPath(parentDir, newFileName); - const notebooks = initNotebook ? [structuredClone(initNotebook), notebook] : [notebook]; - - const newProject: DeepnoteFile = { - metadata: deepnoteFile.metadata - ? structuredClone(deepnoteFile.metadata) - : { createdAt: new Date().toISOString() }, - project: { - ...deepnoteFile.project, - notebooks - }, - version: deepnoteFile.version - }; - - if (initNotebook && initNotebookId) { - newProject.project.initNotebookId = initNotebookId; - } - - // Recompute snapshotHash for the new file - (newProject.metadata as Record).snapshotHash = await this.computeSnapshotHash(newProject); + const newProject = await buildSingleNotebookFile(deepnoteFile, notebook); const yaml = serializeDeepnoteFile(newProject); await workspace.fs.writeFile(newFileUri, new TextEncoder().encode(yaml)); @@ -81,9 +68,7 @@ export class DeepnoteAutoSplitter { deepnoteFile.project.notebooks = originalNotebooks; if (deepnoteFile.metadata) { - (deepnoteFile.metadata as Record).snapshotHash = await this.computeSnapshotHash( - deepnoteFile - ); + (deepnoteFile.metadata as Record).snapshotHash = await computeSnapshotHash(deepnoteFile); } const updatedYaml = serializeDeepnoteFile(deepnoteFile); @@ -198,7 +183,7 @@ export class DeepnoteAutoSplitter { // Recompute hash if (notebookData.metadata) { - (notebookData.metadata as Record).snapshotHash = await this.computeSnapshotHash( + (notebookData.metadata as Record).snapshotHash = await computeSnapshotHash( notebookData ); } @@ -241,51 +226,4 @@ export class DeepnoteAutoSplitter { return withoutSuffix.slice(projectIdIndex + 1 + projectId.length + 1); } - - /** - * Computes snapshotHash using the same algorithm as DeepnoteNotebookSerializer. - */ - private async computeSnapshotHash(project: DeepnoteFile): Promise { - const contentHashes: string[] = []; - - for (const notebook of project.project.notebooks) { - for (const block of notebook.blocks ?? []) { - if (block.contentHash) { - contentHashes.push(block.contentHash); - } - } - } - - contentHashes.sort(); - - const hashInput = { - contentHashes, - environmentHash: project.environment?.hash ?? null, - integrations: (project.project.integrations ?? []) - .map((i: { id: string; name: string; type: string }) => ({ id: i.id, name: i.name, type: i.type })) - .sort((a: { id: string }, b: { id: string }) => a.id.localeCompare(b.id)), - version: project.version - }; - - const hashData = JSON.stringify(hashInput); - const hash = await computeHash(hashData, 'SHA-256'); - - return `sha256:${hash}`; - } - - private getFileStem(uri: Uri): string { - const basename = uri.path.split('/').pop() ?? ''; - const dotIndex = basename.indexOf('.'); - - return dotIndex > 0 ? basename.slice(0, dotIndex) : basename; - } - - private slugifyNotebookName(name: string): string { - try { - return slugifyProjectName(name); - } catch { - // Fallback for names that produce empty slugs - return 'notebook'; - } - } } diff --git a/src/notebooks/deepnote/deepnoteExplorerView.ts b/src/notebooks/deepnote/deepnoteExplorerView.ts index eac7c6fb3d..0aef5fb395 100644 --- a/src/notebooks/deepnote/deepnoteExplorerView.ts +++ b/src/notebooks/deepnote/deepnoteExplorerView.ts @@ -1,5 +1,5 @@ import { injectable, inject } from 'inversify'; -import { commands, window, workspace, type TreeView, Uri, l10n } from 'vscode'; +import { commands, RelativePattern, window, workspace, type TreeView, Uri, l10n } from 'vscode'; import { serializeDeepnoteFile, type DeepnoteBlock, type DeepnoteFile } from '@deepnote/blocks'; import { convertDeepnoteToJupyterNotebooks, convertIpynbFilesToDeepnoteFile } from '@deepnote/convert'; @@ -10,7 +10,9 @@ import { type DeepnoteTreeItem, DeepnoteTreeItemType, type DeepnoteTreeItemConte import { uuidUtils } from '../../platform/common/uuid'; import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; import { Commands } from '../../platform/common/constants'; +import { buildSiblingNotebookFileUri, buildSingleNotebookFile } from './deepnoteNotebookFileFactory'; import { readDeepnoteProjectFile } from './deepnoteProjectUtils'; +import { SNAPSHOT_FILE_SUFFIX } from './snapshots/snapshotFiles'; import { ILogger } from '../../platform/logging/types'; /** @@ -47,25 +49,27 @@ export class DeepnoteExplorerView { } /** - * Shared helper that creates and adds a new notebook to a project - * @param fileUri The URI of the project file + * Shared helper that creates a new notebook in a new sibling `.deepnote` file. + * The new file shares the same project id/name/version/metadata as the source file + * and contains only the newly-created notebook (plus the source's init notebook if any). + * @param fileUri The URI of the source project file (not modified) * @returns Object with notebook ID and name if successful, or null if aborted/failed */ public async createAndAddNotebookToProject(fileUri: Uri): Promise<{ id: string; name: string } | null> { // Read the Deepnote project file - const projectData = await readDeepnoteProjectFile(fileUri); + const sourceData = await readDeepnoteProjectFile(fileUri); - if (!projectData?.project) { + if (!sourceData?.project) { await window.showErrorMessage(l10n.t('Invalid Deepnote file format')); return null; } + // Aggregate notebook names across all sibling files sharing the same project ID + const existingNames = await this.collectNotebookNamesForProject(sourceData.project.id); + // Generate suggested name and prompt user - const suggestedName = this.generateSuggestedNotebookName(projectData); - const notebookName = await this.promptForNotebookName( - suggestedName, - new Set(projectData.project.notebooks?.map((nb: DeepnoteNotebook) => nb.name) ?? []) - ); + const suggestedName = this.generateSuggestedNotebookName(existingNames); + const notebookName = await this.promptForNotebookName(suggestedName, existingNames); if (!notebookName) { return null; @@ -74,14 +78,32 @@ export class DeepnoteExplorerView { // Create new notebook with initial block const newNotebook = this.createNotebookWithFirstBlock(notebookName); - // Add new notebook to the project (initialize array if needed) - if (!projectData.project.notebooks) { - projectData.project.notebooks = []; - } - projectData.project.notebooks.push(newNotebook); + // Build a single-notebook sibling file (preserves source project metadata and init notebook) + const newProject = await buildSingleNotebookFile(sourceData, newNotebook); + const newFileUri = await buildSiblingNotebookFileUri(fileUri, notebookName, async (u) => { + try { + await workspace.fs.stat(u); + + return true; + } catch { + return false; + } + }); + + const yaml = serializeDeepnoteFile(newProject); - // Save and open the new notebook - await this.saveProjectAndOpenNotebook(fileUri, projectData); + await workspace.fs.writeFile(newFileUri, new TextEncoder().encode(yaml)); + + // Refresh the tree view + this.treeDataProvider.refresh(); + + // Open the newly-created file + const document = await workspace.openNotebookDocument(newFileUri); + + await window.showNotebookDocument(document, { + preserveFocus: false, + preview: false + }); return { id: newNotebook.id, name: notebookName }; } @@ -414,16 +436,15 @@ export class DeepnoteExplorerView { } /** - * Generates a suggested unique notebook name based on existing notebooks - * @param projectData The project data containing existing notebooks + * Generates a suggested unique notebook name based on the set of existing notebook names + * across the project group (i.e., all sibling files sharing the same project ID). + * @param existingNames The set of already-taken notebook names * @returns A unique suggested notebook name */ - private generateSuggestedNotebookName(projectData: DeepnoteFile): string { - const notebookCount = projectData.project.notebooks?.length || 0; - const existingNames = new Set(projectData.project.notebooks?.map((nb: DeepnoteNotebook) => nb.name) || []); - - let nextNumber = notebookCount + 1; + private generateSuggestedNotebookName(existingNames: Set): string { + let nextNumber = existingNames.size + 1; let suggestedName = `Notebook ${nextNumber}`; + while (existingNames.has(suggestedName)) { nextNumber++; suggestedName = `Notebook ${nextNumber}`; @@ -485,32 +506,42 @@ export class DeepnoteExplorerView { } /** - * Saves the project data to file and opens the specified notebook - * @param fileUri The URI of the project file - * @param projectData The project data to save - * @param notebookId The notebook ID to open + * Aggregates notebook names across all `.deepnote` files in the workspace whose + * `project.id` matches the given project ID. Used for project-wide uniqueness validation. + * @param projectId The project ID to match against + * @returns Set of notebook names taken across the project group */ - private async saveProjectAndOpenNotebook(fileUri: Uri, projectData: DeepnoteFile): Promise { - // Update metadata timestamp - if (!projectData.metadata) { - projectData.metadata = { createdAt: new Date().toISOString() }; + private async collectNotebookNamesForProject(projectId: string): Promise> { + const names = new Set(); + const workspaceFolders = workspace.workspaceFolders; + + if (!workspaceFolders || workspaceFolders.length === 0) { + return names; } - projectData.metadata.modifiedAt = new Date().toISOString(); - // Write the updated YAML - const updatedYaml = serializeDeepnoteFile(projectData); - const encoder = new TextEncoder(); - await workspace.fs.writeFile(fileUri, encoder.encode(updatedYaml)); + for (const folder of workspaceFolders) { + const pattern = new RelativePattern(folder, '**/*.deepnote'); + const files = await workspace.findFiles(pattern); + const projectFiles = files.filter((file) => !file.path.endsWith(SNAPSHOT_FILE_SUFFIX)); - // Refresh the tree view - await this.treeDataProvider.refreshNotebook(projectData.project.id); + for (const file of projectFiles) { + try { + const project = await readDeepnoteProjectFile(file); - // Open the notebook - const document = await workspace.openNotebookDocument(fileUri); - await window.showNotebookDocument(document, { - preserveFocus: false, - preview: false - }); + if (project?.project?.id !== projectId) { + continue; + } + + for (const notebook of project.project.notebooks ?? []) { + names.add(notebook.name); + } + } catch { + // Skip files that fail to parse; uniqueness is best-effort across siblings + } + } + } + + return names; } public refreshTree(): void { diff --git a/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts b/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts index f3afc275f3..915078d714 100644 --- a/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts @@ -779,7 +779,7 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { }); suite('createAndAddNotebookToProject', () => { - test('should create and add a new notebook to an existing project', async () => { + test('should create and add a new notebook as a new sibling .deepnote file', async () => { const projectId = 'test-project-id'; const existingNotebookId = 'existing-notebook-id'; const newNotebookId = 'new-notebook-id'; @@ -787,6 +787,7 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { const blockId = 'test-block-id'; const fileUri = Uri.file('/workspace/test-project.deepnote'); const notebookName = 'New Notebook'; + const expectedNewFilePath = '/workspace/test-project_new-notebook.deepnote'; // Mock existing project data const existingProjectData: DeepnoteFile = { @@ -811,13 +812,14 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { const yamlContent = serializeDeepnoteFile(existingProjectData); - // Mock file system + // Mock file system: track writes per URI path and stat rejects (no collision) const mockFS = mock(); when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(yamlContent))); + when(mockFS.stat(anything())).thenReject(new Error('File not found')); - let capturedWriteContent: Uint8Array | undefined; - when(mockFS.writeFile(anything(), anything())).thenCall((_uri: Uri, content: Uint8Array) => { - capturedWriteContent = content; + const writes = new Map(); + when(mockFS.writeFile(anything(), anything())).thenCall((uri: Uri, content: Uint8Array) => { + writes.set(uri.path, content); return Promise.resolve(); }); when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); @@ -825,20 +827,24 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { // Mock user input when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(notebookName)); - // Mock UUID generation by mocking crypto.randomUUID + // Mock UUID generation const uuidStub = createUuidMock([newNotebookId, blockGroupId, blockId]); uuidStubs.push(uuidStub); - // Mock notebook opening + // Capture openNotebookDocument URI + let openedUri: Uri | undefined; const mockNotebook = { notebookType: 'deepnote' }; - when(mockedVSCodeNamespaces.workspace.openNotebookDocument(anything())).thenReturn( - Promise.resolve(mockNotebook as any) - ); + when(mockedVSCodeNamespaces.workspace.openNotebookDocument(anything())).thenCall((u: Uri) => { + openedUri = u; + return Promise.resolve(mockNotebook as any); + }); when(mockedVSCodeNamespaces.window.showNotebookDocument(anything(), anything())).thenReturn( Promise.resolve(undefined as any) ); - // Execute the method + // Spy on treeDataProvider.refresh to verify it is called after the new file is written + const refreshSpy = sandbox.spy((explorerView as any).treeDataProvider, 'refresh'); + const result = await explorerView.createAndAddNotebookToProject(fileUri); // Verify result @@ -846,25 +852,278 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { expect(result?.id).to.equal(newNotebookId); expect(result?.name).to.equal(notebookName); - // Verify file was written - expect(capturedWriteContent).to.exist; + // Verify original file was NOT written to + expect(writes.has(fileUri.path)).to.be.false; - // Verify YAML content - const updatedYamlContent = Buffer.from(capturedWriteContent!).toString('utf8'); - const updatedProjectData = deserializeDeepnoteFile(updatedYamlContent) as any; + // Verify new sibling file at expected path was written + expect(writes.has(expectedNewFilePath)).to.be.true; + const newFileContent = writes.get(expectedNewFilePath)!; + const newFileYaml = Buffer.from(newFileContent).toString('utf8'); + const newProjectData = deserializeDeepnoteFile(newFileYaml); - expect(updatedProjectData.project.notebooks).to.have.lengthOf(2); - expect(updatedProjectData.project.notebooks[1].id).to.equal(newNotebookId); - expect(updatedProjectData.project.notebooks[1].name).to.equal(notebookName); - expect(updatedProjectData.project.notebooks[1].blocks).to.have.lengthOf(1); - expect(updatedProjectData.project.notebooks[1].executionMode).to.equal('block'); + // The new file should contain exactly 1 notebook (the new one) + expect(newProjectData.project.notebooks).to.have.lengthOf(1); + expect(newProjectData.project.notebooks[0].id).to.equal(newNotebookId); + expect(newProjectData.project.notebooks[0].name).to.equal(notebookName); + expect(newProjectData.project.notebooks[0].blocks).to.have.lengthOf(1); + expect(newProjectData.project.notebooks[0].executionMode).to.equal('block'); + + // The new file should share the source project.id + expect(newProjectData.project.id).to.equal(projectId); + + // openNotebookDocument should have been called with the NEW file URI + expect(openedUri).to.exist; + expect(openedUri!.path).to.equal(expectedNewFilePath); + + // Tree view should have been refreshed after the new file was written + expect(refreshSpy.called).to.be.true; + }); + + test('should clone init notebook into the new sibling file', async () => { + const projectId = 'test-project-id'; + const initNotebookId = 'init-notebook-id'; + const newNotebookId = 'new-notebook-id'; + const blockGroupId = 'test-blockgroup-id'; + const blockId = 'test-block-id'; + const fileUri = Uri.file('/workspace/test-project.deepnote'); + const notebookName = 'New Notebook'; + const expectedNewFilePath = '/workspace/test-project_new-notebook.deepnote'; + + const existingProjectData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2024-01-01T00:00:00.000Z', + modifiedAt: '2024-01-01T00:00:00.000Z' + }, + project: { + id: projectId, + name: 'Test Project', + initNotebookId, + notebooks: [ + { + id: initNotebookId, + name: 'Init', + blocks: [ + { + blockGroup: 'init-bg', + content: 'print("init")', + id: 'init-block', + metadata: {}, + sortingKey: '0', + type: 'code', + version: 1 + } + ], + executionMode: 'block' + }, + { + id: 'other-nb', + name: 'Other Notebook', + blocks: [], + executionMode: 'block' + } + ] + } + }; + + const yamlContent = serializeDeepnoteFile(existingProjectData); + + const mockFS = mock(); + when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(yamlContent))); + when(mockFS.stat(anything())).thenReject(new Error('File not found')); + + const writes = new Map(); + when(mockFS.writeFile(anything(), anything())).thenCall((uri: Uri, content: Uint8Array) => { + writes.set(uri.path, content); + return Promise.resolve(); + }); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(notebookName)); + + const uuidStub = createUuidMock([newNotebookId, blockGroupId, blockId]); + uuidStubs.push(uuidStub); + + when(mockedVSCodeNamespaces.workspace.openNotebookDocument(anything())).thenReturn( + Promise.resolve({ notebookType: 'deepnote' } as any) + ); + when(mockedVSCodeNamespaces.window.showNotebookDocument(anything(), anything())).thenReturn( + Promise.resolve(undefined as any) + ); + + await explorerView.createAndAddNotebookToProject(fileUri); + + // Original file is not written + expect(writes.has(fileUri.path)).to.be.false; + + // New file exists with both init + new notebook + const newFileYaml = Buffer.from(writes.get(expectedNewFilePath)!).toString('utf8'); + const newProjectData = deserializeDeepnoteFile(newFileYaml); + + expect(newProjectData.project.notebooks).to.have.lengthOf(2); + // Init notebook should be present and have preserved id + content + const initInNew = newProjectData.project.notebooks.find((nb) => nb.id === initNotebookId); + expect(initInNew).to.exist; + expect(initInNew!.name).to.equal('Init'); + expect(initInNew!.blocks).to.have.lengthOf(1); + expect(initInNew!.blocks[0].content).to.equal('print("init")'); + + // New user notebook is present + const newNotebook = newProjectData.project.notebooks.find((nb) => nb.id === newNotebookId); + expect(newNotebook).to.exist; + expect(newNotebook!.name).to.equal(notebookName); + + // initNotebookId preserved on new file + expect(newProjectData.project.initNotebookId).to.equal(initNotebookId); + }); + + test('should append numeric suffix when sibling filename collides', async () => { + const projectId = 'test-project-id'; + const newNotebookId = 'new-notebook-id'; + const blockGroupId = 'test-blockgroup-id'; + const blockId = 'test-block-id'; + const fileUri = Uri.file('/workspace/test-project.deepnote'); + const notebookName = 'New Notebook'; + const collidingPath = '/workspace/test-project_new-notebook.deepnote'; + const expectedPath = '/workspace/test-project_new-notebook_2.deepnote'; + + const existingProjectData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2024-01-01T00:00:00.000Z', + modifiedAt: '2024-01-01T00:00:00.000Z' + }, + project: { + id: projectId, + name: 'Test Project', + notebooks: [{ id: 'existing', name: 'Existing', blocks: [], executionMode: 'block' }] + } + }; + + const yamlContent = serializeDeepnoteFile(existingProjectData); + + const mockFS = mock(); + when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(yamlContent))); + + // First candidate exists, second does not + when(mockFS.stat(anything())).thenCall((uri: Uri) => { + if (uri.path === collidingPath) { + return Promise.resolve({} as any); + } + return Promise.reject(new Error('File not found')); + }); + + const writes = new Map(); + when(mockFS.writeFile(anything(), anything())).thenCall((uri: Uri, content: Uint8Array) => { + writes.set(uri.path, content); + return Promise.resolve(); + }); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(notebookName)); + + const uuidStub = createUuidMock([newNotebookId, blockGroupId, blockId]); + uuidStubs.push(uuidStub); + + let openedUri: Uri | undefined; + when(mockedVSCodeNamespaces.workspace.openNotebookDocument(anything())).thenCall((u: Uri) => { + openedUri = u; + return Promise.resolve({ notebookType: 'deepnote' } as any); + }); + when(mockedVSCodeNamespaces.window.showNotebookDocument(anything(), anything())).thenReturn( + Promise.resolve(undefined as any) + ); + + await explorerView.createAndAddNotebookToProject(fileUri); + + // Ensure no write to the original or the colliding path + expect(writes.has(fileUri.path)).to.be.false; + expect(writes.has(collidingPath)).to.be.false; + + // The new file should have been written to the `_2` path + expect(writes.has(expectedPath)).to.be.true; + expect(openedUri?.path).to.equal(expectedPath); + }); + + test('should validate name uniqueness across sibling project files', async () => { + const projectId = 'test-project-id'; + const fileUri = Uri.file('/workspace/test-project.deepnote'); + const siblingUri = Uri.file('/workspace/test-project_other.deepnote'); + + const sourceData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2024-01-01T00:00:00.000Z', + modifiedAt: '2024-01-01T00:00:00.000Z' + }, + project: { + id: projectId, + name: 'Test Project', + notebooks: [{ id: 'nb-source', name: 'Source Notebook', blocks: [], executionMode: 'block' }] + } + }; + + const siblingData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2024-01-01T00:00:00.000Z', + modifiedAt: '2024-01-01T00:00:00.000Z' + }, + project: { + id: projectId, + name: 'Test Project', + notebooks: [{ id: 'nb-shared', name: 'Shared Name', blocks: [], executionMode: 'block' }] + } + }; + + const sourceYaml = serializeDeepnoteFile(sourceData); + const siblingYaml = serializeDeepnoteFile(siblingData); + + // Activate workspaceFolders so collectNotebookNamesForProject runs findFiles + const workspaceFolder = { uri: Uri.file('/workspace'), name: 'workspace', index: 0 }; + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder as any]); + + // Return both files from findFiles + when(mockedVSCodeNamespaces.workspace.findFiles(anything())).thenReturn( + Promise.resolve([fileUri, siblingUri]) + ); + + const mockFS = mock(); + when(mockFS.readFile(anything())).thenCall((uri: Uri) => { + if (uri.path === siblingUri.path) { + return Promise.resolve(Buffer.from(siblingYaml)); + } + return Promise.resolve(Buffer.from(sourceYaml)); + }); + when(mockFS.stat(anything())).thenReject(new Error('File not found')); + when(mockFS.writeFile(anything(), anything())).thenReturn(Promise.resolve()); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + + // Capture validateInput from showInputBox options; have the user cancel so no further work runs + let capturedValidateInput: ((value: string) => string | null | undefined) | undefined; + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenCall((options: any) => { + capturedValidateInput = options?.validateInput; + return Promise.resolve(undefined); + }); + + await explorerView.createAndAddNotebookToProject(fileUri); + + expect(capturedValidateInput).to.exist; + + // 'Shared Name' is taken by sibling -> rejection (non-null string) + const result = capturedValidateInput!('Shared Name'); + expect(result).to.be.a('string'); + expect(result).to.not.be.null; + + // A unique name should pass validation (null return) + const okResult = capturedValidateInput!('Totally Unique Name'); + expect(okResult).to.be.null; }); test('should return null if user cancels notebook name input', async () => { const projectId = 'test-project-id'; const fileUri = Uri.file('/workspace/test-project.deepnote'); - // Mock existing project data const existingProjectData: DeepnoteFile = { version: '1.0.0', metadata: { @@ -880,19 +1139,17 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { const yamlContent = serializeDeepnoteFile(existingProjectData); - // Mock file system const mockFS = mock(); when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(yamlContent))); + when(mockFS.stat(anything())).thenReject(new Error('File not found')); when(mockFS.writeFile(anything(), anything())).thenReturn(Promise.resolve()); when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - // Mock user cancelling input when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); - // Execute the method const result = await explorerView.createAndAddNotebookToProject(fileUri); - // Verify result is null and file was not written + // Null result and no file writes occurred (neither original nor sibling) expect(result).to.be.null; verify(mockFS.writeFile(anything(), anything())).never(); }); @@ -901,7 +1158,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { const projectId = 'test-project-id'; const fileUri = Uri.file('/workspace/test-project.deepnote'); - // Mock existing project data with multiple notebooks const existingProjectData: DeepnoteFile = { version: '1.0.0', metadata: { @@ -920,9 +1176,15 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { const yamlContent = serializeDeepnoteFile(existingProjectData); - // Mock file system + // Tolerate workspace.findFiles being called - return only the source URI so only + // its own notebook names contribute to the suggested name logic + const workspaceFolder = { uri: Uri.file('/workspace'), name: 'workspace', index: 0 }; + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder as any]); + when(mockedVSCodeNamespaces.workspace.findFiles(anything())).thenReturn(Promise.resolve([fileUri])); + const mockFS = mock(); when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(yamlContent))); + when(mockFS.stat(anything())).thenReject(new Error('File not found')); when(mockFS.writeFile(anything(), anything())).thenReturn(Promise.resolve()); when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); @@ -942,10 +1204,9 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { Promise.resolve(undefined as any) ); - // Execute the method await explorerView.createAndAddNotebookToProject(fileUri); - // Verify suggested name is 'Notebook 3' (next in sequence) + // With two existing notebooks, suggestion is `Notebook ${size + 1}` = 'Notebook 3' expect(capturedInputBoxOptions).to.exist; expect(capturedInputBoxOptions.value).to.equal('Notebook 3'); }); diff --git a/src/notebooks/deepnote/deepnoteNotebookFileFactory.ts b/src/notebooks/deepnote/deepnoteNotebookFileFactory.ts new file mode 100644 index 0000000000..60fa71cdc4 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteNotebookFileFactory.ts @@ -0,0 +1,118 @@ +import { type DeepnoteFile } from '@deepnote/blocks'; +import { Uri } from 'vscode'; + +import { computeHash } from '../../platform/common/crypto'; +import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; +import { slugifyProjectName } from './snapshots/snapshotFiles'; + +/** + * Builds a new DeepnoteFile containing a single user notebook (plus optional init notebook), + * sharing the source's project id, name, version, and metadata. + */ +export async function buildSingleNotebookFile(source: DeepnoteFile, notebook: DeepnoteNotebook): Promise { + const initNotebookId = source.project.initNotebookId; + const initNotebook = initNotebookId ? source.project.notebooks.find((nb) => nb.id === initNotebookId) : undefined; + + const metadata = source.metadata ? structuredClone(source.metadata) : { createdAt: new Date().toISOString() }; + + metadata.modifiedAt = new Date().toISOString(); + + const notebooks = initNotebook ? [structuredClone(initNotebook), notebook] : [notebook]; + + const newProject: DeepnoteFile = { + metadata, + project: { + ...source.project, + notebooks + }, + version: source.version + }; + + if (initNotebook && initNotebookId) { + newProject.project.initNotebookId = initNotebookId; + } else { + delete newProject.project.initNotebookId; + } + + (newProject.metadata as Record).snapshotHash = await computeSnapshotHash(newProject); + + return newProject; +} + +/** + * Builds a sibling file URI for a notebook based on the original file's stem and a slugified notebook name. + * If the resulting path already exists, appends `_2`, `_3`, ... until a unique name is found. + */ +export async function buildSiblingNotebookFileUri( + originalUri: Uri, + notebookName: string, + exists: (uri: Uri) => Promise +): Promise { + const parentDir = Uri.joinPath(originalUri, '..'); + const originalStem = getFileStem(originalUri); + const slug = slugifyNotebookNameOrFallback(notebookName); + const baseName = `${originalStem}_${slug}`; + + let candidate = Uri.joinPath(parentDir, `${baseName}.deepnote`); + let suffix = 2; + + while (await exists(candidate)) { + candidate = Uri.joinPath(parentDir, `${baseName}_${suffix}.deepnote`); + suffix++; + } + + return candidate; +} + +/** + * Computes snapshotHash using the same algorithm as DeepnoteNotebookSerializer. + */ +export async function computeSnapshotHash(project: DeepnoteFile): Promise { + const contentHashes: string[] = []; + + for (const notebook of project.project.notebooks) { + for (const block of notebook.blocks ?? []) { + if (block.contentHash) { + contentHashes.push(block.contentHash); + } + } + } + + contentHashes.sort(); + + const hashInput = { + contentHashes, + environmentHash: project.environment?.hash ?? null, + integrations: (project.project.integrations ?? []) + .map((i: { id: string; name: string; type: string }) => ({ id: i.id, name: i.name, type: i.type })) + .sort((a: { id: string }, b: { id: string }) => a.id.localeCompare(b.id)), + version: project.version + }; + + const hashData = JSON.stringify(hashInput); + const hash = await computeHash(hashData, 'SHA-256'); + + return `sha256:${hash}`; +} + +/** + * Extracts the file stem (portion before the first dot) from a URI's basename. + */ +export function getFileStem(uri: Uri): string { + const basename = uri.path.split('/').pop() ?? ''; + const dotIndex = basename.indexOf('.'); + + return dotIndex > 0 ? basename.slice(0, dotIndex) : basename; +} + +/** + * Slugifies a notebook name, falling back to 'notebook' if the name cannot be slugified. + */ +export function slugifyNotebookNameOrFallback(name: string): string { + try { + return slugifyProjectName(name); + } catch { + // Fallback for names that produce empty slugs + return 'notebook'; + } +} diff --git a/src/notebooks/deepnote/deepnoteNotebookFileFactory.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookFileFactory.unit.test.ts new file mode 100644 index 0000000000..dfda1049f5 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteNotebookFileFactory.unit.test.ts @@ -0,0 +1,305 @@ +import { type DeepnoteFile } from '@deepnote/blocks'; +import { assert, expect } from 'chai'; +import { Uri } from 'vscode'; + +import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; +import { + buildSiblingNotebookFileUri, + buildSingleNotebookFile, + computeSnapshotHash, + getFileStem, + slugifyNotebookNameOrFallback +} from './deepnoteNotebookFileFactory'; + +function createNotebook(id: string, name: string, blockId = `block-${id}`, contentHash?: string): DeepnoteNotebook { + return { + blocks: [ + { + blockGroup: `bg-${id}`, + content: '', + contentHash, + id: blockId, + metadata: {}, + sortingKey: '0', + type: 'code', + version: 1 + } + ], + executionMode: 'block', + id, + name + }; +} + +function createSourceFile(overrides?: Partial): DeepnoteFile { + return { + metadata: { + createdAt: '2024-01-01T00:00:00.000Z', + modifiedAt: '2024-01-01T00:00:00.000Z' + }, + project: { + id: 'project-1', + name: 'Source Project', + notebooks: [createNotebook('nb-1', 'First Notebook', 'block-nb-1', 'hash-a')], + ...overrides + }, + version: '1.0.0' + }; +} + +suite('deepnoteNotebookFileFactory', () => { + suite('buildSingleNotebookFile', () => { + test('should preserve project id/name/version and set metadata.modifiedAt while preserving createdAt', async () => { + const source = createSourceFile(); + const notebook = createNotebook('nb-new', 'My New Notebook', 'block-nb-new', 'hash-b'); + + const result = await buildSingleNotebookFile(source, notebook); + + assert.strictEqual(result.project.id, 'project-1'); + assert.strictEqual(result.project.name, 'Source Project'); + assert.strictEqual(result.version, '1.0.0'); + assert.strictEqual(result.metadata?.createdAt, '2024-01-01T00:00:00.000Z'); + assert.isString(result.metadata?.modifiedAt); + // modifiedAt should be freshly set (not the frozen source value) + assert.notStrictEqual(result.metadata?.modifiedAt, '2024-01-01T00:00:00.000Z'); + }); + + test('should compute snapshotHash metadata', async () => { + const source = createSourceFile(); + const notebook = createNotebook('nb-new', 'My New Notebook', 'block-nb-new', 'hash-b'); + + const result = await buildSingleNotebookFile(source, notebook); + const snapshotHash = (result.metadata as Record).snapshotHash; + + assert.isString(snapshotHash); + expect(snapshotHash as string).to.match(/^sha256:/); + }); + + test('should include init notebook clone when initNotebookId is set', async () => { + const initNotebook = createNotebook('init-nb', 'Init Notebook', 'block-init', 'hash-init'); + const source = createSourceFile({ + initNotebookId: 'init-nb', + notebooks: [initNotebook, createNotebook('nb-1', 'First Notebook', 'block-nb-1', 'hash-a')] + }); + const newNotebook = createNotebook('nb-new', 'My New Notebook', 'block-nb-new', 'hash-b'); + + const result = await buildSingleNotebookFile(source, newNotebook); + + assert.strictEqual(result.project.notebooks.length, 2); + assert.strictEqual(result.project.notebooks[0].id, 'init-nb'); + assert.strictEqual(result.project.notebooks[0].name, 'Init Notebook'); + assert.strictEqual(result.project.notebooks[1].id, 'nb-new'); + assert.strictEqual(result.project.initNotebookId, 'init-nb'); + + // Init notebook should be a clone (different reference) but same content + assert.notStrictEqual(result.project.notebooks[0], initNotebook); + }); + + test('should not include init notebook when initNotebookId is not set', async () => { + const source = createSourceFile(); + const newNotebook = createNotebook('nb-new', 'My New Notebook', 'block-nb-new', 'hash-b'); + + const result = await buildSingleNotebookFile(source, newNotebook); + + assert.strictEqual(result.project.notebooks.length, 1); + assert.strictEqual(result.project.notebooks[0].id, 'nb-new'); + assert.isUndefined(result.project.initNotebookId); + }); + + test('should not include init notebook when initNotebookId is set but init notebook is missing', async () => { + const source = createSourceFile({ + initNotebookId: 'missing-init-id' + }); + const newNotebook = createNotebook('nb-new', 'My New Notebook', 'block-nb-new', 'hash-b'); + + const result = await buildSingleNotebookFile(source, newNotebook); + + assert.strictEqual(result.project.notebooks.length, 1); + assert.strictEqual(result.project.notebooks[0].id, 'nb-new'); + assert.isUndefined(result.project.initNotebookId); + }); + }); + + suite('buildSiblingNotebookFileUri', () => { + test('should return ${stem}_${slug}.deepnote when the path does not exist', async () => { + const sourceUri = Uri.file('/workspace/test-project.deepnote'); + const exists = async (_u: Uri) => false; + + const result = await buildSiblingNotebookFileUri(sourceUri, 'My Notebook', exists); + + assert.strictEqual(result.path, '/workspace/test-project_my-notebook.deepnote'); + }); + + test('should append _2 on first collision', async () => { + const sourceUri = Uri.file('/workspace/test-project.deepnote'); + const exists = async (u: Uri) => u.path === '/workspace/test-project_my-notebook.deepnote'; + + const result = await buildSiblingNotebookFileUri(sourceUri, 'My Notebook', exists); + + assert.strictEqual(result.path, '/workspace/test-project_my-notebook_2.deepnote'); + }); + + test('should keep incrementing suffix on repeated collisions', async () => { + const sourceUri = Uri.file('/workspace/test-project.deepnote'); + const taken = new Set([ + '/workspace/test-project_my-notebook.deepnote', + '/workspace/test-project_my-notebook_2.deepnote', + '/workspace/test-project_my-notebook_3.deepnote' + ]); + const exists = async (u: Uri) => taken.has(u.path); + + const result = await buildSiblingNotebookFileUri(sourceUri, 'My Notebook', exists); + + assert.strictEqual(result.path, '/workspace/test-project_my-notebook_4.deepnote'); + }); + + test("should use 'notebook' fallback slug for names that slugify to empty", async () => { + const sourceUri = Uri.file('/workspace/test-project.deepnote'); + const exists = async (_u: Uri) => false; + + const result = await buildSiblingNotebookFileUri(sourceUri, '!!!', exists); + + assert.strictEqual(result.path, '/workspace/test-project_notebook.deepnote'); + }); + }); + + suite('slugifyNotebookNameOrFallback', () => { + test("should return 'my-notebook' for 'My Notebook'", () => { + assert.strictEqual(slugifyNotebookNameOrFallback('My Notebook'), 'my-notebook'); + }); + + test("should return 'notebook' for '!!!' (unslugifiable name)", () => { + assert.strictEqual(slugifyNotebookNameOrFallback('!!!'), 'notebook'); + }); + + test("should return 'notebook' for empty name", () => { + assert.strictEqual(slugifyNotebookNameOrFallback(''), 'notebook'); + }); + }); + + suite('getFileStem', () => { + test("should return 'test-project' for /a/b/test-project.deepnote", () => { + const uri = Uri.file('/a/b/test-project.deepnote'); + + assert.strictEqual(getFileStem(uri), 'test-project'); + }); + + test("should return 'file' for /a/b/file (no dot)", () => { + const uri = Uri.file('/a/b/file'); + + assert.strictEqual(getFileStem(uri), 'file'); + }); + + test('should stop at the first dot for multi-dot filenames', () => { + const uri = Uri.file('/a/b/my.snapshot.deepnote'); + + assert.strictEqual(getFileStem(uri), 'my'); + }); + }); + + suite('computeSnapshotHash', () => { + test('should be deterministic - same input twice returns same hash', async () => { + const file = createSourceFile(); + + const hashA = await computeSnapshotHash(file); + const hashB = await computeSnapshotHash(file); + + assert.strictEqual(hashA, hashB); + expect(hashA).to.match(/^sha256:/); + }); + + test("should change when a block's contentHash changes", async () => { + const fileA = createSourceFile(); + const fileB = createSourceFile({ + notebooks: [createNotebook('nb-1', 'First Notebook', 'block-nb-1', 'hash-CHANGED')] + }); + + const hashA = await computeSnapshotHash(fileA); + const hashB = await computeSnapshotHash(fileB); + + assert.notStrictEqual(hashA, hashB); + }); + + test('should be insensitive to block order (hashes are sorted internally)', async () => { + const fileA: DeepnoteFile = { + metadata: { createdAt: '2024-01-01T00:00:00.000Z' }, + project: { + id: 'p', + name: 'p', + notebooks: [ + { + blocks: [ + { + blockGroup: 'bg', + content: '', + contentHash: 'h-a', + id: 'b1', + metadata: {}, + sortingKey: '0', + type: 'code', + version: 1 + }, + { + blockGroup: 'bg', + content: '', + contentHash: 'h-b', + id: 'b2', + metadata: {}, + sortingKey: '1', + type: 'code', + version: 1 + } + ], + executionMode: 'block', + id: 'nb', + name: 'nb' + } + ] + }, + version: '1.0.0' + }; + const fileB: DeepnoteFile = { + metadata: { createdAt: '2024-01-01T00:00:00.000Z' }, + project: { + id: 'p', + name: 'p', + notebooks: [ + { + blocks: [ + { + blockGroup: 'bg', + content: '', + contentHash: 'h-b', + id: 'b2', + metadata: {}, + sortingKey: '1', + type: 'code', + version: 1 + }, + { + blockGroup: 'bg', + content: '', + contentHash: 'h-a', + id: 'b1', + metadata: {}, + sortingKey: '0', + type: 'code', + version: 1 + } + ], + executionMode: 'block', + id: 'nb', + name: 'nb' + } + ] + }, + version: '1.0.0' + }; + + const hashA = await computeSnapshotHash(fileA); + const hashB = await computeSnapshotHash(fileB); + + assert.strictEqual(hashA, hashB); + }); + }); +}); From cea5ecd92dc9e7ecccc4f7ad41f5686c1662518c Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 17 Apr 2026 09:06:16 +0000 Subject: [PATCH 41/47] Rename test --- src/notebooks/deepnote/deepnoteNotebookFileFactory.unit.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/notebooks/deepnote/deepnoteNotebookFileFactory.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookFileFactory.unit.test.ts index dfda1049f5..407235e71e 100644 --- a/src/notebooks/deepnote/deepnoteNotebookFileFactory.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookFileFactory.unit.test.ts @@ -168,7 +168,7 @@ suite('deepnoteNotebookFileFactory', () => { assert.strictEqual(slugifyNotebookNameOrFallback('My Notebook'), 'my-notebook'); }); - test("should return 'notebook' for '!!!' (unslugifiable name)", () => { + test("should return 'notebook' for '!!!' (name that produces no slug)", () => { assert.strictEqual(slugifyNotebookNameOrFallback('!!!'), 'notebook'); }); From a47d889647e8ddebc3a34a581734b7eef71c32df Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 20 Apr 2026 07:58:30 +0000 Subject: [PATCH 42/47] fix(deepnote): Update context handling for project files and notebooks - Refactored conditions in `deepnoteExplorerView` to correctly identify and handle project groups and notebook files. - Enhanced the `DeepnoteTreeItem` to adopt the label of a single non-init notebook when applicable, improving user experience in the explorer view. - Adjusted the `DeepnoteTreeDataProvider` to ensure proper collapsible states for project groups based on the number of notebooks. - Updated unit tests to validate the new behavior and ensure accurate representation of project and notebook states. --- package.json | 16 +- .../deepnote/deepnoteExplorerView.ts | 330 ++++--- .../deepnoteExplorerView.unit.test.ts | 866 +++++++++++++++--- .../deepnote/deepnoteTreeDataProvider.ts | 30 +- .../deepnoteTreeDataProvider.unit.test.ts | 197 +++- src/notebooks/deepnote/deepnoteTreeItem.ts | 41 +- .../deepnote/deepnoteTreeItem.unit.test.ts | 366 +++++++- 7 files changed, 1513 insertions(+), 333 deletions(-) diff --git a/package.json b/package.json index c9f71a0d14..05af26b156 100644 --- a/package.json +++ b/package.json @@ -1600,42 +1600,42 @@ }, { "command": "deepnote.addNotebookToProject", - "when": "view == deepnoteExplorer && viewItem == projectFile", + "when": "view == deepnoteExplorer && viewItem == projectGroup", "group": "1_add@1" }, { "command": "deepnote.renameProject", - "when": "view == deepnoteExplorer && viewItem == projectFile", + "when": "view == deepnoteExplorer && viewItem == projectGroup", "group": "2_manage@1" }, { "command": "deepnote.exportProject", - "when": "view == deepnoteExplorer && viewItem == projectFile", + "when": "view == deepnoteExplorer && viewItem == projectGroup", "group": "2_manage@2" }, { "command": "deepnote.deleteProject", - "when": "view == deepnoteExplorer && viewItem == projectFile", + "when": "view == deepnoteExplorer && viewItem == projectGroup", "group": "3_delete@1" }, { "command": "deepnote.renameNotebook", - "when": "view == deepnoteExplorer && viewItem == notebook", + "when": "view == deepnoteExplorer && (viewItem == notebookFile || viewItem == notebook)", "group": "1_edit@1" }, { "command": "deepnote.duplicateNotebook", - "when": "view == deepnoteExplorer && viewItem == notebook", + "when": "view == deepnoteExplorer && (viewItem == notebookFile || viewItem == notebook)", "group": "1_edit@2" }, { "command": "deepnote.exportNotebook", - "when": "view == deepnoteExplorer && viewItem == notebook", + "when": "view == deepnoteExplorer && (viewItem == notebookFile || viewItem == notebook)", "group": "2_export@1" }, { "command": "deepnote.deleteNotebook", - "when": "view == deepnoteExplorer && viewItem == notebook", + "when": "view == deepnoteExplorer && (viewItem == notebookFile || viewItem == notebook)", "group": "3_delete@1" } ] diff --git a/src/notebooks/deepnote/deepnoteExplorerView.ts b/src/notebooks/deepnote/deepnoteExplorerView.ts index 0aef5fb395..a2fb193002 100644 --- a/src/notebooks/deepnote/deepnoteExplorerView.ts +++ b/src/notebooks/deepnote/deepnoteExplorerView.ts @@ -6,9 +6,16 @@ import { convertDeepnoteToJupyterNotebooks, convertIpynbFilesToDeepnoteFile } fr import { IExtensionContext } from '../../platform/common/types'; import { DeepnoteTreeDataProvider } from './deepnoteTreeDataProvider'; -import { type DeepnoteTreeItem, DeepnoteTreeItemType, type DeepnoteTreeItemContext } from './deepnoteTreeItem'; +import { + type DeepnoteTreeItem, + DeepnoteTreeItemType, + type DeepnoteTreeItemContext, + getSingleNonInitNotebook, + NOTEBOOK_FILE_CONTEXT_VALUE, + type ProjectGroupData +} from './deepnoteTreeItem'; import { uuidUtils } from '../../platform/common/uuid'; -import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; +import type { DeepnoteNotebook, DeepnoteProject } from '../../platform/deepnote/deepnoteTypes'; import { Commands } from '../../platform/common/constants'; import { buildSiblingNotebookFileUri, buildSingleNotebookFile } from './deepnoteNotebookFileFactory'; import { readDeepnoteProjectFile } from './deepnoteProjectUtils'; @@ -20,6 +27,7 @@ import { ILogger } from '../../platform/logging/types'; */ @injectable() export class DeepnoteExplorerView { + private readonly logger: ILogger; private readonly treeDataProvider: DeepnoteTreeDataProvider; private treeView: TreeView; @@ -28,6 +36,7 @@ export class DeepnoteExplorerView { @inject(IExtensionContext) private readonly extensionContext: IExtensionContext, @inject(ILogger) logger: ILogger ) { + this.logger = logger; this.treeDataProvider = new DeepnoteTreeDataProvider(logger); } @@ -109,33 +118,27 @@ export class DeepnoteExplorerView { } public async renameNotebook(treeItem: DeepnoteTreeItem): Promise { - if (treeItem.type !== DeepnoteTreeItemType.Notebook) { + const target = this.resolveNotebookTarget(treeItem); + + if (!target) { return; } try { - const fileUri = Uri.file(treeItem.context.filePath); + const { fileUri, notebookId } = target; const projectData = await readDeepnoteProjectFile(fileUri); if (!projectData?.project?.notebooks) { await window.showErrorMessage(l10n.t('Invalid Deepnote file format')); return; } - const targetNotebook = projectData.project.notebooks.find( - (nb: DeepnoteNotebook) => nb.id === treeItem.context.notebookId - ); + const targetNotebook = projectData.project.notebooks.find((nb: DeepnoteNotebook) => nb.id === notebookId); if (!targetNotebook) { await window.showErrorMessage(l10n.t('Notebook not found')); return; } - const itemNotebook = treeItem.data as DeepnoteNotebook; - const currentName = itemNotebook.name; - - if (targetNotebook.id !== itemNotebook.id) { - await window.showErrorMessage(l10n.t('Selected notebook is not the target notebook')); - return; - } + const currentName = targetNotebook.name; const existingNames = new Set( projectData.project.notebooks @@ -169,25 +172,16 @@ export class DeepnoteExplorerView { } public async deleteNotebook(treeItem: DeepnoteTreeItem): Promise { - if (treeItem.type !== DeepnoteTreeItemType.Notebook) { - return; - } - - const notebook = treeItem.data as DeepnoteNotebook; - const notebookName = notebook.name; + const target = this.resolveNotebookTarget(treeItem); - const confirmation = await window.showWarningMessage( - l10n.t('Are you sure you want to delete notebook "{0}"?', notebookName), - { modal: true }, - l10n.t('Delete') - ); - - if (confirmation !== l10n.t('Delete')) { + if (!target) { return; } + const { fileUri, notebookId } = target; + const isSingleNotebookFile = treeItem.contextValue === NOTEBOOK_FILE_CONTEXT_VALUE; + try { - const fileUri = Uri.file(treeItem.context.filePath); const projectData = await readDeepnoteProjectFile(fileUri); if (!projectData?.project?.notebooks) { @@ -195,8 +189,35 @@ export class DeepnoteExplorerView { return; } + const targetNotebook = projectData.project.notebooks.find((nb: DeepnoteNotebook) => nb.id === notebookId); + + if (!targetNotebook) { + await window.showErrorMessage(l10n.t('Notebook not found')); + return; + } + + const notebookName = targetNotebook.name; + + const confirmation = await window.showWarningMessage( + l10n.t('Are you sure you want to delete notebook "{0}"?', notebookName), + { modal: true }, + l10n.t('Delete') + ); + + if (confirmation !== l10n.t('Delete')) { + return; + } + + // Single-notebook file: removing the sole notebook would leave an empty file; delete the file instead + if (isSingleNotebookFile) { + await workspace.fs.delete(fileUri); + this.treeDataProvider.refresh(); + await window.showInformationMessage(l10n.t('Notebook deleted: {0}', notebookName)); + return; + } + projectData.project.notebooks = projectData.project.notebooks.filter( - (nb: DeepnoteNotebook) => nb.id !== treeItem.context.notebookId + (nb: DeepnoteNotebook) => nb.id !== notebookId ); if (!projectData.metadata) { @@ -217,15 +238,16 @@ export class DeepnoteExplorerView { } public async duplicateNotebook(treeItem: DeepnoteTreeItem): Promise { - if (treeItem.type !== DeepnoteTreeItemType.Notebook) { + const target = this.resolveNotebookTarget(treeItem); + + if (!target) { return; } - const notebook = treeItem.data as DeepnoteNotebook; - const originalName = notebook.name; + const { fileUri, notebookId } = target; + const isSingleNotebookFile = treeItem.contextValue === NOTEBOOK_FILE_CONTEXT_VALUE; try { - const fileUri = Uri.file(treeItem.context.filePath); const projectData = await readDeepnoteProjectFile(fileUri); if (!projectData?.project?.notebooks) { @@ -233,17 +255,21 @@ export class DeepnoteExplorerView { return; } - const targetNotebook = projectData.project.notebooks.find( - (nb: DeepnoteNotebook) => nb.id === treeItem.context.notebookId - ); + const targetNotebook = projectData.project.notebooks.find((nb: DeepnoteNotebook) => nb.id === notebookId); if (!targetNotebook) { await window.showErrorMessage(l10n.t('Notebook not found')); return; } - // Generate new name - const existingNames = new Set(projectData.project.notebooks.map((nb: DeepnoteNotebook) => nb.name)); + const originalName = targetNotebook.name; + + // For sibling-file duplicates we need to avoid name collisions across the whole project group, + // not just within a single file — otherwise two siblings could end up with the same notebook name. + const existingNames = isSingleNotebookFile + ? await this.collectNotebookNamesForProject(projectData.project.id) + : new Set(projectData.project.notebooks.map((nb: DeepnoteNotebook) => nb.name)); + let copyNumber = 1; let newName = `${originalName} (Copy)`; while (existingNames.has(newName)) { @@ -251,26 +277,37 @@ export class DeepnoteExplorerView { newName = `${originalName} (Copy ${copyNumber})`; } - // Deep clone the notebook and generate new IDs - const newNotebook: DeepnoteNotebook = { - ...targetNotebook, - id: uuidUtils.generateUuid(), - name: newName, - blocks: targetNotebook.blocks.map((block: DeepnoteBlock) => { - // Use structuredClone for deep cloning if available, otherwise fall back to JSON - const clonedBlock = - typeof structuredClone !== 'undefined' - ? structuredClone(block) - : JSON.parse(JSON.stringify(block)); - - // Update cloned block with new IDs and reset execution state - clonedBlock.id = uuidUtils.generateUuid(); - clonedBlock.blockGroup = uuidUtils.generateUuid(); - clonedBlock.executionCount = undefined; - - return clonedBlock; - }) - }; + const newNotebook = this.cloneNotebookWithFreshIds(targetNotebook, newName); + + if (isSingleNotebookFile) { + // Build a sibling `.deepnote` so each single-notebook file carries exactly one notebook + const newProject = await buildSingleNotebookFile(projectData, newNotebook); + const newFileUri = await buildSiblingNotebookFileUri(fileUri, newName, async (u) => { + try { + await workspace.fs.stat(u); + + return true; + } catch { + return false; + } + }); + + const yaml = serializeDeepnoteFile(newProject); + + await workspace.fs.writeFile(newFileUri, new TextEncoder().encode(yaml)); + + this.treeDataProvider.refresh(); + + const document = await workspace.openNotebookDocument(newFileUri); + + await window.showNotebookDocument(document, { + preserveFocus: false, + preview: false + }); + + await window.showInformationMessage(l10n.t('Notebook duplicated: {0}', newName)); + return; + } projectData.project.notebooks.push(newNotebook); @@ -300,12 +337,12 @@ export class DeepnoteExplorerView { } public async renameProject(treeItem: DeepnoteTreeItem): Promise { - if (treeItem.type !== DeepnoteTreeItemType.ProjectFile) { + if (treeItem.type !== DeepnoteTreeItemType.ProjectGroup) { return; } - const project = treeItem.data as DeepnoteFile; - const currentName = project.project.name; + const groupData = treeItem.data as ProjectGroupData; + const currentName = groupData.projectName; const newName = await window.showInputBox({ prompt: l10n.t('Enter new project name'), @@ -323,26 +360,29 @@ export class DeepnoteExplorerView { } try { - const fileUri = Uri.file(treeItem.context.filePath); - const projectData = await readDeepnoteProjectFile(fileUri); + // Mirror the new name across every sibling file in the group so siblings don't diverge + for (const file of groupData.files) { + const fileUri = Uri.file(file.filePath); + const projectData = await readDeepnoteProjectFile(fileUri); + + if (!projectData?.project) { + this.logger.error(`Failed to parse Deepnote file during rename: ${file.filePath}`); + continue; + } - if (!projectData?.project) { - await window.showErrorMessage(l10n.t('Invalid Deepnote file format')); - return; - } + projectData.project.name = newName; - projectData.project.name = newName; + if (!projectData.metadata) { + projectData.metadata = { createdAt: new Date().toISOString() }; + } + projectData.metadata.modifiedAt = new Date().toISOString(); - if (!projectData.metadata) { - projectData.metadata = { createdAt: new Date().toISOString() }; + const updatedYaml = serializeDeepnoteFile(projectData); + const encoder = new TextEncoder(); + await workspace.fs.writeFile(fileUri, encoder.encode(updatedYaml)); } - projectData.metadata.modifiedAt = new Date().toISOString(); - const updatedYaml = serializeDeepnoteFile(projectData); - const encoder = new TextEncoder(); - await workspace.fs.writeFile(fileUri, encoder.encode(updatedYaml)); - - await this.treeDataProvider.refreshProject(treeItem.context.filePath); + this.treeDataProvider.refresh(); await window.showInformationMessage(l10n.t('Project renamed to: {0}', newName)); } catch (error) { const errorMessage = error instanceof Error ? error.message : 'Unknown error'; @@ -435,6 +475,66 @@ export class DeepnoteExplorerView { ); } + /** + * Deep-clones a notebook with a new id, freshly-generated block ids/blockGroups, and + * reset `executionCount`. Used by duplicate flows so cached state keyed on + * `(projectId, notebookId)` or on block identity can't collide across copies. + */ + private cloneNotebookWithFreshIds(source: DeepnoteNotebook, newName: string): DeepnoteNotebook { + return { + ...source, + id: uuidUtils.generateUuid(), + name: newName, + blocks: source.blocks.map((block: DeepnoteBlock) => { + const clonedBlock = + typeof structuredClone !== 'undefined' ? structuredClone(block) : JSON.parse(JSON.stringify(block)); + + clonedBlock.id = uuidUtils.generateUuid(); + clonedBlock.blockGroup = uuidUtils.generateUuid(); + clonedBlock.executionCount = undefined; + + return clonedBlock; + }) + }; + } + + /** + * Resolves a tree item to the `(fileUri, notebookId)` pair that notebook-level handlers operate on. + * Accepts both legacy inner `Notebook` items and single-notebook `ProjectFile` items + * (contextValue `notebookFile`). Returns undefined for any other tree item kind. + */ + private resolveNotebookTarget(treeItem: DeepnoteTreeItem): { fileUri: Uri; notebookId: string } | undefined { + if (treeItem.type === DeepnoteTreeItemType.Notebook) { + if (!treeItem.context.notebookId) { + return undefined; + } + + return { + fileUri: Uri.file(treeItem.context.filePath), + notebookId: treeItem.context.notebookId + }; + } + + if ( + treeItem.type === DeepnoteTreeItemType.ProjectFile && + treeItem.contextValue === NOTEBOOK_FILE_CONTEXT_VALUE + ) { + const project = treeItem.data as DeepnoteProject; + const singleNotebook = getSingleNonInitNotebook(project); + + if (!singleNotebook) { + return undefined; + } + + return { + fileUri: Uri.file(treeItem.context.filePath), + notebookId: singleNotebook.id + }; + } + + return undefined; + } + /** * Generates a suggested unique notebook name based on the set of existing notebook names * across the project group (i.e., all sibling files sharing the same project ID). @@ -936,12 +1036,12 @@ export class DeepnoteExplorerView { } private async deleteProject(treeItem: DeepnoteTreeItem): Promise { - if (treeItem.type !== DeepnoteTreeItemType.ProjectFile) { + if (treeItem.type !== DeepnoteTreeItemType.ProjectGroup) { return; } - const project = treeItem.data as DeepnoteFile; - const projectName = project.project.name; + const groupData = treeItem.data as ProjectGroupData; + const projectName = groupData.projectName; const confirmation = await window.showWarningMessage( l10n.t('Are you sure you want to delete project "{0}"?', projectName), @@ -954,8 +1054,14 @@ export class DeepnoteExplorerView { } try { - const fileUri = Uri.file(treeItem.context.filePath); - await workspace.fs.delete(fileUri); + for (const file of groupData.files) { + try { + await workspace.fs.delete(Uri.file(file.filePath)); + } catch (error) { + this.logger.error(`Failed to delete ${file.filePath}`, error); + } + } + this.treeDataProvider.refresh(); await window.showInformationMessage(l10n.t('Project deleted: {0}', projectName)); } catch (error) { @@ -965,14 +1071,20 @@ export class DeepnoteExplorerView { } private async addNotebookToProject(treeItem: DeepnoteTreeItem): Promise { - if (treeItem.type !== DeepnoteTreeItemType.ProjectFile) { + if (treeItem.type !== DeepnoteTreeItemType.ProjectGroup) { + return; + } + + const groupData = treeItem.data as ProjectGroupData; + + if (groupData.files.length === 0) { return; } try { - const fileUri = Uri.file(treeItem.context.filePath); + // Use the first file in the group as the template source (project id/name/metadata carry over) + const fileUri = Uri.file(groupData.files[0].filePath); - // Use shared helper to create and add notebook const result = await this.createAndAddNotebookToProject(fileUri); if (result) { @@ -985,14 +1097,15 @@ export class DeepnoteExplorerView { } /** - * Exports all notebooks from a Deepnote project to Jupyter format - * @param treeItem The tree item representing a project + * Exports all notebooks from a Deepnote project group (across every sibling file) to Jupyter format. */ private async exportProject(treeItem: DeepnoteTreeItem): Promise { - if (treeItem.type !== DeepnoteTreeItemType.ProjectFile) { + if (treeItem.type !== DeepnoteTreeItemType.ProjectGroup) { return; } + const groupData = treeItem.data as ProjectGroupData; + try { const format = await window.showQuickPick([{ label: 'Jupyter Notebook (.ipynb)', value: 'jupyter' }], { placeHolder: l10n.t('Select export format') @@ -1002,15 +1115,6 @@ export class DeepnoteExplorerView { return; } - const fileUri = Uri.file(treeItem.context.filePath); - const projectData = await readDeepnoteProjectFile(fileUri); - - if (!projectData?.project) { - await window.showErrorMessage(l10n.t('Invalid Deepnote file format')); - - return; - } - const outputFolder = await window.showOpenDialog({ canSelectFiles: false, canSelectFolders: true, @@ -1023,7 +1127,26 @@ export class DeepnoteExplorerView { return; } - const jupyterNotebooks = convertDeepnoteToJupyterNotebooks(projectData); + type JupyterNotebookEntry = { filename: string; notebook: unknown }; + const jupyterNotebooks: JupyterNotebookEntry[] = []; + + for (const file of groupData.files) { + const fileUri = Uri.file(file.filePath); + const projectData = await readDeepnoteProjectFile(fileUri); + + if (!projectData?.project) { + this.logger.error(`Failed to parse Deepnote file during export: ${file.filePath}`); + continue; + } + + const perFile = convertDeepnoteToJupyterNotebooks(projectData); + jupyterNotebooks.push(...perFile); + } + + if (jupyterNotebooks.length === 0) { + await window.showErrorMessage(l10n.t('Invalid Deepnote file format')); + return; + } // Check for existing files before writing const existingFiles: string[] = []; @@ -1071,14 +1194,18 @@ export class DeepnoteExplorerView { } /** - * Exports a single notebook from a Deepnote project to Jupyter format - * @param treeItem The tree item representing a notebook + * Exports a single notebook (either a single-notebook file or a legacy inner notebook) + * from a Deepnote project to Jupyter format. */ private async exportNotebook(treeItem: DeepnoteTreeItem): Promise { - if (treeItem.type !== DeepnoteTreeItemType.Notebook) { + const target = this.resolveNotebookTarget(treeItem); + + if (!target) { return; } + const { fileUri, notebookId } = target; + try { const format = await window.showQuickPick([{ label: 'Jupyter Notebook (.ipynb)', value: 'jupyter' }], { placeHolder: l10n.t('Select export format') @@ -1088,7 +1215,6 @@ export class DeepnoteExplorerView { return; } - const fileUri = Uri.file(treeItem.context.filePath); const projectData = await readDeepnoteProjectFile(fileUri); if (!projectData?.project) { @@ -1109,7 +1235,7 @@ export class DeepnoteExplorerView { return; } - const targetNotebook = projectData.project.notebooks.find((nb) => nb.id === treeItem.context.notebookId); + const targetNotebook = projectData.project.notebooks.find((nb) => nb.id === notebookId); if (!targetNotebook) { await window.showErrorMessage(l10n.t('Notebook not found')); diff --git a/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts b/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts index 915078d714..695ea2b146 100644 --- a/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts @@ -7,7 +7,13 @@ import { stringify as yamlStringify } from 'yaml'; import { DeepnoteExplorerView } from './deepnoteExplorerView'; -import { DeepnoteTreeItem, DeepnoteTreeItemType, type DeepnoteTreeItemContext } from './deepnoteTreeItem'; +import { + DeepnoteTreeItem, + DeepnoteTreeItemType, + NOTEBOOK_FILE_CONTEXT_VALUE, + type DeepnoteTreeItemContext, + type ProjectGroupData +} from './deepnoteTreeItem'; import type { IExtensionContext } from '../../platform/common/types'; import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; @@ -1525,8 +1531,23 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { const notebookName = 'Notebook to Delete'; const fileUri = Uri.file('/workspace/test-project.deepnote'); + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { createdAt: '2024-01-01T00:00:00.000Z', modifiedAt: '2024-01-01T00:00:00.000Z' }, + project: { + id: projectId, + name: 'Test Project', + notebooks: [ + { id: 'other-nb', name: 'Other', blocks: [], executionMode: 'block' }, + { id: notebookId, name: notebookName, blocks: [], executionMode: 'block' } + ] + } + }; + const mockFS = mock(); - when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(''))); + when(mockFS.readFile(anything())).thenReturn( + Promise.resolve(Buffer.from(serializeDeepnoteFile(projectData))) + ); when(mockFS.writeFile(anything(), anything())).thenReturn(Promise.resolve()); when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); @@ -1535,7 +1556,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { Promise.resolve(undefined) ); - // Create mock tree item const mockTreeItem: Partial = { type: DeepnoteTreeItemType.Notebook, context: { @@ -1551,12 +1571,11 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { } as DeepnoteNotebook }; - // Execute the method await explorerView.deleteNotebook(mockTreeItem as DeepnoteTreeItem); - // Verify file operations were not called (user cancelled) - verify(mockFS.readFile(anything())).never(); + // File is read to look up the notebook name for the prompt, but never written after cancellation verify(mockFS.writeFile(anything(), anything())).never(); + verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).once(); }); }); @@ -1897,14 +1916,14 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { }); suite('renameProject', () => { - test('should successfully rename a project with valid input', async () => { + test('should update project.name in every file in the group', async () => { const oldProjectName = 'Old Project Name'; const newProjectName = 'New Project Name'; const projectId = 'test-project-id'; - const fileUri = Uri.file('/workspace/test-project.deepnote'); + const fileA = '/workspace/test-project.deepnote'; + const fileB = '/workspace/test-project_sibling.deepnote'; - // Mock existing project data - const existingProjectData: DeepnoteFile = { + const projectDataA: DeepnoteFile = { version: '1.0.0', metadata: { createdAt: '2024-01-01T00:00:00.000Z', @@ -1913,71 +1932,76 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { project: { id: projectId, name: oldProjectName, - notebooks: [ - { - id: 'notebook-1', - name: 'Notebook 1', - blocks: [], - executionMode: 'block' - } - ] + notebooks: [{ id: 'notebook-1', name: 'Notebook 1', blocks: [], executionMode: 'block' }] } }; - const yamlContent = serializeDeepnoteFile(existingProjectData); + const projectDataB: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2024-01-01T00:00:00.000Z', + modifiedAt: '2024-01-01T00:00:00.000Z' + }, + project: { + id: projectId, + name: oldProjectName, + notebooks: [{ id: 'notebook-2', name: 'Notebook 2', blocks: [], executionMode: 'block' }] + } + }; - // Mock file system const mockFS = mock(); - when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(yamlContent))); - let capturedWriteContent: Uint8Array | undefined; - when(mockFS.writeFile(anything(), anything())).thenCall((_uri: Uri, content: Uint8Array) => { - capturedWriteContent = content; + when(mockFS.readFile(anything())).thenCall((uri: Uri) => { + if (uri.path === fileB) { + return Promise.resolve(Buffer.from(serializeDeepnoteFile(projectDataB))); + } + + return Promise.resolve(Buffer.from(serializeDeepnoteFile(projectDataA))); + }); + + const writes = new Map(); + when(mockFS.writeFile(anything(), anything())).thenCall((uri: Uri, content: Uint8Array) => { + writes.set(uri.path, content); return Promise.resolve(); }); when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - // Mock user input for new name when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(newProjectName)); when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenReturn( Promise.resolve(undefined) ); - // Create mock tree item + const groupData = { + projectId, + projectName: oldProjectName, + files: [ + { filePath: fileA, project: projectDataA }, + { filePath: fileB, project: projectDataB } + ] + }; + const mockTreeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: fileUri.fsPath, - projectId: projectId - }, - data: existingProjectData as unknown as DeepnoteFile + type: DeepnoteTreeItemType.ProjectGroup, + context: { filePath: fileA, projectId }, + data: groupData as any }; - // Execute the method await explorerView.renameProject(mockTreeItem as DeepnoteTreeItem); - // Verify file was written - expect(capturedWriteContent).to.exist; - - // Verify YAML content - const updatedYamlContent = Buffer.from(capturedWriteContent!).toString('utf8'); - const updatedProjectData = deserializeDeepnoteFile(updatedYamlContent); - - // Verify project was renamed - expect(updatedProjectData.project.name).to.equal(newProjectName); + // Both files should have been written with the new name + expect(writes.size).to.equal(2); + expect(writes.has(fileA)).to.be.true; + expect(writes.has(fileB)).to.be.true; - // Verify notebooks were not affected - expect(updatedProjectData.project.notebooks).to.have.lengthOf(1); - expect(updatedProjectData.project.notebooks[0].name).to.equal('Notebook 1'); - - // Verify metadata was updated - expect(updatedProjectData.metadata.modifiedAt).to.exist; + for (const content of writes.values()) { + const updated = deserializeDeepnoteFile(Buffer.from(content).toString('utf8')); + expect(updated.project.name).to.equal(newProjectName); + } - // Verify success message was shown verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); }); - test('should return early if tree item type is not ProjectFile', async () => { + test('should return early if tree item type is not ProjectGroup', async () => { const mockTreeItem: Partial = { type: DeepnoteTreeItemType.Notebook, context: { @@ -2000,12 +2024,28 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { verify(mockFS.readFile(anything())).never(); }); + test('should return early if the tree item is a ProjectFile (not a group)', async () => { + const mockTreeItem: Partial = { + type: DeepnoteTreeItemType.ProjectFile, + context: { + filePath: '/workspace/test-project.deepnote', + projectId: 'test-project-id' + } + }; + + const mockFS = mock(); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + + await explorerView.renameProject(mockTreeItem as DeepnoteTreeItem); + + verify(mockedVSCodeNamespaces.window.showInputBox(anything())).never(); + verify(mockFS.readFile(anything())).never(); + }); + test('should return early if user cancels input or provides same name', async () => { const projectId = 'test-project-id'; const currentName = 'Current Project Name'; - const fileUri = Uri.file('/workspace/test-project.deepnote'); - // Mock existing project data const existingProjectData: DeepnoteFile = { version: '1.0.0', metadata: { @@ -2019,30 +2059,28 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { } }; - const yamlContent = serializeDeepnoteFile(existingProjectData); - const mockFS = mock(); - when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(yamlContent))); + when(mockFS.readFile(anything())).thenReturn( + Promise.resolve(Buffer.from(serializeDeepnoteFile(existingProjectData))) + ); when(mockFS.writeFile(anything(), anything())).thenReturn(Promise.resolve()); when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); // Test 1: User cancels input (returns undefined) when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); - // Create mock tree item const mockTreeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: fileUri.fsPath, - projectId: projectId - }, - data: existingProjectData as unknown as DeepnoteFile + type: DeepnoteTreeItemType.ProjectGroup, + context: { filePath: '/workspace/test-project.deepnote', projectId }, + data: { + projectId, + projectName: currentName, + files: [{ filePath: '/workspace/test-project.deepnote', project: existingProjectData }] + } as any }; - // Execute the method await explorerView.renameProject(mockTreeItem as DeepnoteTreeItem); - // Verify file was not written (early return) verify(mockFS.writeFile(anything(), anything())).never(); // Test 2: User provides same name @@ -2050,34 +2088,55 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { await explorerView.renameProject(mockTreeItem as DeepnoteTreeItem); - // Verify file was still not written (early return) verify(mockFS.writeFile(anything(), anything())).never(); }); }); suite('exportProject', () => { + function buildGroupTreeItem( + projectId: string, + filesInGroup: Array<{ filePath: string; project: DeepnoteFile }> + ): Partial { + return { + type: DeepnoteTreeItemType.ProjectGroup, + context: { + filePath: filesInGroup[0]?.filePath ?? '/test/project.deepnote', + projectId + }, + data: { + projectId, + projectName: 'Test Project', + files: filesInGroup + } as any + }; + } + test('should return early if user cancels format selection', async () => { resetVSCodeMocks(); const mockFS = mock(); when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - // User cancels format selection when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( Promise.resolve(undefined) ); - const treeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/test/project.deepnote', - projectId: 'project-id' + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { createdAt: '2024-01-01T00:00:00.000Z', modifiedAt: '2024-01-01T00:00:00.000Z' }, + project: { + id: 'project-id', + name: 'Test Project', + notebooks: [{ id: 'nb-1', name: 'Notebook 1', blocks: [], executionMode: 'block' }] } }; + const treeItem = buildGroupTreeItem('project-id', [ + { filePath: '/test/project.deepnote', project: projectData } + ]); + await (explorerView as any).exportProject(treeItem); - // Verify no file operations occurred verify(mockFS.writeFile(anything(), anything())).never(); verify(mockFS.readFile(anything())).never(); }); @@ -2087,10 +2146,7 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { const projectData: DeepnoteFile = { version: '1.0.0', - metadata: { - createdAt: '2024-01-01T00:00:00.000Z', - modifiedAt: '2024-01-01T00:00:00.000Z' - }, + metadata: { createdAt: '2024-01-01T00:00:00.000Z', modifiedAt: '2024-01-01T00:00:00.000Z' }, project: { id: 'project-id', name: 'Test Project', @@ -2104,30 +2160,23 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { ); when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - // User selects format but cancels folder selection when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( Promise.resolve({ label: 'Jupyter Notebook (.ipynb)', value: 'jupyter' }) as any ); when(mockedVSCodeNamespaces.window.showOpenDialog(anything())).thenReturn(Promise.resolve(undefined)); - const treeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/test/project.deepnote', - projectId: 'project-id' - } - }; + const treeItem = buildGroupTreeItem('project-id', [ + { filePath: '/test/project.deepnote', project: projectData } + ]); await (explorerView as any).exportProject(treeItem); - // Verify no file was written verify(mockFS.writeFile(anything(), anything())).never(); }); - test('should show error for invalid Deepnote file format', async () => { + test('should show error when every file in the group fails to parse', async () => { resetVSCodeMocks(); - // Invalid project data (no project property) const invalidData = { version: '1.0.0', metadata: { createdAt: '2024-01-01T00:00:00.000Z' } @@ -2140,20 +2189,81 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( Promise.resolve({ label: 'Jupyter Notebook (.ipynb)', value: 'jupyter' }) as any ); + when(mockedVSCodeNamespaces.window.showOpenDialog(anything())).thenReturn( + Promise.resolve([Uri.file('/output/folder')]) + ); when(mockedVSCodeNamespaces.window.showErrorMessage(anything())).thenReturn(Promise.resolve(undefined)); - const treeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/test/project.deepnote', - projectId: 'project-id' + const treeItem = buildGroupTreeItem('project-id', [ + { filePath: '/test/project.deepnote', project: invalidData as any } + ]); + + await (explorerView as any).exportProject(treeItem); + + verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); + }); + + test('should concatenate Jupyter outputs across every file in the group', async () => { + resetVSCodeMocks(); + + const projectA: DeepnoteFile = { + version: '1.0.0', + metadata: { createdAt: '2024-01-01T00:00:00.000Z', modifiedAt: '2024-01-01T00:00:00.000Z' }, + project: { + id: 'project-id', + name: 'Test Project', + notebooks: [{ id: 'nb-a', name: 'Notebook A', blocks: [], executionMode: 'block' }] + } + }; + + const projectB: DeepnoteFile = { + version: '1.0.0', + metadata: { createdAt: '2024-01-01T00:00:00.000Z', modifiedAt: '2024-01-01T00:00:00.000Z' }, + project: { + id: 'project-id', + name: 'Test Project', + notebooks: [{ id: 'nb-b', name: 'Notebook B', blocks: [], executionMode: 'block' }] } }; + const mockFS = mock(); + + when(mockFS.readFile(anything())).thenCall((uri: Uri) => { + if (uri.path === '/test/project_b.deepnote') { + return Promise.resolve(Buffer.from(serializeDeepnoteFile(projectB))); + } + + return Promise.resolve(Buffer.from(serializeDeepnoteFile(projectA))); + }); + when(mockFS.stat(anything())).thenReject(new Error('File not found')); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( + Promise.resolve({ label: 'Jupyter Notebook (.ipynb)', value: 'jupyter' }) as any + ); + when(mockedVSCodeNamespaces.window.showOpenDialog(anything())).thenReturn( + Promise.resolve([Uri.file('/output/folder')]) + ); + when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenReturn( + Promise.resolve(undefined) + ); + + let writeCount = 0; + when(mockFS.writeFile(anything(), anything())).thenCall(() => { + writeCount++; + return Promise.resolve(); + }); + + const treeItem = buildGroupTreeItem('project-id', [ + { filePath: '/test/project.deepnote', project: projectA }, + { filePath: '/test/project_b.deepnote', project: projectB } + ]); + await (explorerView as any).exportProject(treeItem); - // Verify error message was shown - verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); + // One notebook per file -> two writes + assert.strictEqual(writeCount, 2); + verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); }); test('should export all notebooks when triggered from project', async () => { @@ -2198,17 +2308,12 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { return Promise.resolve(); }); - const treeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/test/project.deepnote', - projectId: 'project-id' - } - }; + const treeItem = buildGroupTreeItem('project-id', [ + { filePath: '/test/project.deepnote', project: projectData } + ]); await (explorerView as any).exportProject(treeItem); - // Verify both notebooks were exported assert.strictEqual(writeCount, 2); verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); }); @@ -2268,17 +2373,12 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { return Promise.resolve(); }); - const treeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/test/project.deepnote', - projectId: 'project-id' - } - }; + const treeItem = buildGroupTreeItem('project-id', [ + { filePath: '/test/project.deepnote', project: projectData } + ]); await (explorerView as any).exportProject(treeItem); - // Verify the exported content is valid Jupyter notebook JSON assert.isDefined(capturedContent); const notebook = JSON.parse(Buffer.from(capturedContent!).toString('utf8')); @@ -2327,17 +2427,12 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { return Promise.resolve(); }); - const treeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/test/project.deepnote', - projectId: 'project-id' - } - }; + const treeItem = buildGroupTreeItem('project-id', [ + { filePath: '/test/project.deepnote', project: projectData } + ]); await (explorerView as any).exportProject(treeItem); - // Verify the output path is correctly constructed assert.isDefined(capturedUri); assert.isTrue(capturedUri!.fsPath.startsWith('/output/folder')); assert.isTrue(capturedUri!.fsPath.endsWith('.ipynb')); @@ -2374,20 +2469,14 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { ); when(mockedVSCodeNamespaces.window.showErrorMessage(anything())).thenReturn(Promise.resolve(undefined)); - // Simulate write error when(mockFS.writeFile(anything(), anything())).thenReject(new Error('Permission denied')); - const treeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/test/project.deepnote', - projectId: 'project-id' - } - }; + const treeItem = buildGroupTreeItem('project-id', [ + { filePath: '/test/project.deepnote', project: projectData } + ]); await (explorerView as any).exportProject(treeItem); - // Verify error message was shown verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); }); @@ -2414,7 +2503,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { when(mockFS.readFile(anything())).thenReturn( Promise.resolve(Buffer.from(serializeDeepnoteFile(projectData))) ); - // Files exist - stat returns successfully when(mockFS.stat(anything())).thenReturn(Promise.resolve({} as any)); when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); @@ -2424,24 +2512,18 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { when(mockedVSCodeNamespaces.window.showOpenDialog(anything())).thenReturn( Promise.resolve([Uri.file('/output/folder')]) ); - // User cancels overwrite when(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).thenReturn( Promise.resolve(undefined) ); - const treeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/test/project.deepnote', - projectId: 'project-id' - } - }; + const treeItem = buildGroupTreeItem('project-id', [ + { filePath: '/test/project.deepnote', project: projectData } + ]); await (explorerView as any).exportProject(treeItem); - // Verify warning message was shown about files existing + // Only one warning even though two output files exist (one prompt for the whole batch) verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).once(); - // Verify no files were written verify(mockFS.writeFile(anything(), anything())).never(); }); @@ -2465,7 +2547,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { when(mockFS.readFile(anything())).thenReturn( Promise.resolve(Buffer.from(serializeDeepnoteFile(projectData))) ); - // File exists - stat returns successfully when(mockFS.stat(anything())).thenReturn(Promise.resolve({} as any)); when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); @@ -2475,7 +2556,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { when(mockedVSCodeNamespaces.window.showOpenDialog(anything())).thenReturn( Promise.resolve([Uri.file('/output/folder')]) ); - // User confirms overwrite when(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).thenReturn( Promise.resolve('Overwrite') as any ); @@ -2489,18 +2569,31 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { return Promise.resolve(); }); + const treeItem = buildGroupTreeItem('project-id', [ + { filePath: '/test/project.deepnote', project: projectData } + ]); + + await (explorerView as any).exportProject(treeItem); + + assert.strictEqual(writeCount, 1); + }); + + test('should be a no-op when tree item is not a ProjectGroup', async () => { + resetVSCodeMocks(); + + const mockFS = mock(); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + const treeItem: Partial = { type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/test/project.deepnote', - projectId: 'project-id' - } + context: { filePath: '/test/project.deepnote', projectId: 'project-id' } }; await (explorerView as any).exportProject(treeItem); - // Verify file was written after user confirmed overwrite - assert.strictEqual(writeCount, 1); + verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).never(); + verify(mockFS.readFile(anything())).never(); + verify(mockFS.writeFile(anything(), anything())).never(); }); }); @@ -2882,3 +2975,484 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { }); }); }); + +suite('DeepnoteExplorerView - Project group handlers', () => { + let explorerView: DeepnoteExplorerView; + let sandbox: sinon.SinonSandbox; + let uuidStubs: sinon.SinonStub[] = []; + + setup(() => { + sandbox = sinon.createSandbox(); + resetVSCodeMocks(); + uuidStubs = []; + + const mockContext = { subscriptions: [] } as unknown as IExtensionContext; + const mockLogger = createMockLogger(); + + explorerView = new DeepnoteExplorerView(mockContext, mockLogger); + }); + + teardown(() => { + explorerView.dispose(); + sandbox.restore(); + uuidStubs.forEach((stub) => stub.restore()); + uuidStubs = []; + resetVSCodeMocks(); + }); + + suite('deleteProject', () => { + test('should delete every file in the group after a single confirmation', async () => { + const projectId = 'project-id'; + const fileA = '/workspace/a.deepnote'; + const fileB = '/workspace/b.deepnote'; + const fileC = '/workspace/c.deepnote'; + + const mockFS = mock(); + const deletedPaths: string[] = []; + + when(mockFS.delete(anything())).thenCall((uri: Uri) => { + deletedPaths.push(uri.path); + return Promise.resolve(); + }); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + + let warningCount = 0; + when(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).thenCall(() => { + warningCount++; + return Promise.resolve('Delete') as any; + }); + when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenReturn( + Promise.resolve(undefined) + ); + + const treeItem: Partial = { + type: DeepnoteTreeItemType.ProjectGroup, + context: { filePath: fileA, projectId }, + data: { + projectId, + projectName: 'My Project', + files: [ + { filePath: fileA, project: {} as any }, + { filePath: fileB, project: {} as any }, + { filePath: fileC, project: {} as any } + ] + } as ProjectGroupData as any + }; + + await (explorerView as any).deleteProject(treeItem); + + assert.strictEqual(warningCount, 1, 'should show exactly one confirmation'); + assert.deepStrictEqual(deletedPaths.sort(), [fileA, fileB, fileC].sort()); + verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); + }); + + test('should abort if user declines confirmation', async () => { + const mockFS = mock(); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + + when(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).thenReturn( + Promise.resolve(undefined) + ); + + const treeItem: Partial = { + type: DeepnoteTreeItemType.ProjectGroup, + context: { filePath: '/a', projectId: 'p' }, + data: { + projectId: 'p', + projectName: 'P', + files: [{ filePath: '/workspace/a.deepnote', project: {} as any }] + } as ProjectGroupData as any + }; + + await (explorerView as any).deleteProject(treeItem); + + verify(mockFS.delete(anything())).never(); + }); + + test('should be a no-op when called with a non-group tree item', async () => { + const mockFS = mock(); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + + const treeItem: Partial = { + type: DeepnoteTreeItemType.ProjectFile, + context: { filePath: '/workspace/a.deepnote', projectId: 'p' } + }; + + await (explorerView as any).deleteProject(treeItem); + + verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); + verify(mockFS.delete(anything())).never(); + }); + }); + + suite('addNotebookToProject', () => { + test('uses groupData.files[0] as the source file', async () => { + const projectId = 'project-id'; + const firstFile = '/workspace/a.deepnote'; + const secondFile = '/workspace/b.deepnote'; + + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { createdAt: '2024-01-01T00:00:00.000Z', modifiedAt: '2024-01-01T00:00:00.000Z' }, + project: { + id: projectId, + name: 'Test Project', + notebooks: [{ id: 'nb-1', name: 'Existing', blocks: [], executionMode: 'block' }] + } + }; + + const createStub = sandbox + .stub(explorerView, 'createAndAddNotebookToProject') + .resolves({ id: 'new-nb', name: 'New Notebook' }); + + when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenReturn( + Promise.resolve(undefined) + ); + + const treeItem: Partial = { + type: DeepnoteTreeItemType.ProjectGroup, + context: { filePath: firstFile, projectId }, + data: { + projectId, + projectName: 'Test Project', + files: [ + { filePath: firstFile, project: projectData }, + { filePath: secondFile, project: projectData } + ] + } as ProjectGroupData as any + }; + + await (explorerView as any).addNotebookToProject(treeItem); + + assert.isTrue(createStub.calledOnce); + const callArg = createStub.firstCall.args[0] as Uri; + + assert.strictEqual(callArg.path, firstFile); + }); + + test('is a no-op for non-group tree items', async () => { + const createStub = sandbox.stub(explorerView, 'createAndAddNotebookToProject'); + + const treeItem: Partial = { + type: DeepnoteTreeItemType.ProjectFile, + context: { filePath: '/workspace/a.deepnote', projectId: 'p' } + }; + + await (explorerView as any).addNotebookToProject(treeItem); + + assert.isFalse(createStub.called); + }); + + test('is a no-op when the group has zero files', async () => { + const createStub = sandbox.stub(explorerView, 'createAndAddNotebookToProject'); + + const treeItem: Partial = { + type: DeepnoteTreeItemType.ProjectGroup, + context: { filePath: '', projectId: 'p' }, + data: { projectId: 'p', projectName: 'P', files: [] } as ProjectGroupData as any + }; + + await (explorerView as any).addNotebookToProject(treeItem); + + assert.isFalse(createStub.called); + }); + }); +}); + +suite('DeepnoteExplorerView - Notebook target resolution', () => { + let explorerView: DeepnoteExplorerView; + + setup(() => { + resetVSCodeMocks(); + + const mockContext = { subscriptions: [] } as unknown as IExtensionContext; + const mockLogger = createMockLogger(); + + explorerView = new DeepnoteExplorerView(mockContext, mockLogger); + }); + + teardown(() => { + explorerView.dispose(); + resetVSCodeMocks(); + }); + + test('resolveNotebookTarget returns (fileUri, notebookId) for a Notebook tree item', () => { + const treeItem: Partial = { + type: DeepnoteTreeItemType.Notebook, + contextValue: 'notebook', + context: { + filePath: '/workspace/project.deepnote', + projectId: 'p', + notebookId: 'nb-inner' + } + }; + + const result = (explorerView as any).resolveNotebookTarget(treeItem); + + assert.isDefined(result); + assert.strictEqual(result!.notebookId, 'nb-inner'); + assert.strictEqual(result!.fileUri.path, '/workspace/project.deepnote'); + }); + + test('resolveNotebookTarget returns the sole non-init notebook id for a single-notebook ProjectFile', () => { + const project: DeepnoteFile = { + version: '1.0.0', + metadata: { createdAt: '2024-01-01T00:00:00.000Z', modifiedAt: '2024-01-01T00:00:00.000Z' }, + project: { + id: 'p', + name: 'P', + notebooks: [{ id: 'sole-nb', name: 'Sole', blocks: [], executionMode: 'block' }] + } + }; + + const treeItem: Partial = { + type: DeepnoteTreeItemType.ProjectFile, + contextValue: NOTEBOOK_FILE_CONTEXT_VALUE, + context: { filePath: '/workspace/solo.deepnote', projectId: 'p' }, + data: project as any + }; + + const result = (explorerView as any).resolveNotebookTarget(treeItem); + + assert.isDefined(result); + assert.strictEqual(result!.notebookId, 'sole-nb'); + assert.strictEqual(result!.fileUri.path, '/workspace/solo.deepnote'); + }); + + test('resolveNotebookTarget returns undefined for a legacy multi-notebook ProjectFile', () => { + const project: DeepnoteFile = { + version: '1.0.0', + metadata: { createdAt: '2024-01-01T00:00:00.000Z', modifiedAt: '2024-01-01T00:00:00.000Z' }, + project: { + id: 'p', + name: 'P', + notebooks: [ + { id: 'nb-a', name: 'A', blocks: [], executionMode: 'block' }, + { id: 'nb-b', name: 'B', blocks: [], executionMode: 'block' } + ] + } + }; + + const treeItem: Partial = { + type: DeepnoteTreeItemType.ProjectFile, + contextValue: 'projectFile', + context: { filePath: '/workspace/multi.deepnote', projectId: 'p' }, + data: project as any + }; + + const result = (explorerView as any).resolveNotebookTarget(treeItem); + + assert.isUndefined(result); + }); + + test('resolveNotebookTarget returns undefined for a ProjectGroup', () => { + const treeItem: Partial = { + type: DeepnoteTreeItemType.ProjectGroup, + contextValue: 'projectGroup', + context: { filePath: '/workspace/a.deepnote', projectId: 'p' } + }; + + const result = (explorerView as any).resolveNotebookTarget(treeItem); + + assert.isUndefined(result); + }); + + test('resolveNotebookTarget returns undefined for a Notebook tree item without notebookId', () => { + const treeItem: Partial = { + type: DeepnoteTreeItemType.Notebook, + contextValue: 'notebook', + context: { filePath: '/workspace/a.deepnote', projectId: 'p' } + }; + + const result = (explorerView as any).resolveNotebookTarget(treeItem); + + assert.isUndefined(result); + }); +}); + +suite('DeepnoteExplorerView - Single-notebook file handlers', () => { + let explorerView: DeepnoteExplorerView; + let sandbox: sinon.SinonSandbox; + let uuidStubs: sinon.SinonStub[] = []; + + setup(() => { + sandbox = sinon.createSandbox(); + resetVSCodeMocks(); + uuidStubs = []; + + const mockContext = { subscriptions: [] } as unknown as IExtensionContext; + const mockLogger = createMockLogger(); + + explorerView = new DeepnoteExplorerView(mockContext, mockLogger); + }); + + teardown(() => { + explorerView.dispose(); + sandbox.restore(); + uuidStubs.forEach((stub) => stub.restore()); + uuidStubs = []; + resetVSCodeMocks(); + }); + + suite('duplicateNotebook (single-notebook file)', () => { + test('creates a sibling .deepnote file with freshly regenerated ids', async () => { + const projectId = 'project-id'; + const originalNotebookId = 'original-nb'; + const clonedNotebookId = 'cloned-nb-uuid'; + const clonedBlockId = 'cloned-block-uuid'; + const clonedBlockGroup = 'cloned-block-group-uuid'; + const fileUri = Uri.file('/workspace/solo.deepnote'); + const expectedNewPath = '/workspace/solo_solo-copy.deepnote'; + + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { createdAt: '2024-01-01T00:00:00.000Z', modifiedAt: '2024-01-01T00:00:00.000Z' }, + project: { + id: projectId, + name: 'Solo Project', + notebooks: [ + { + id: originalNotebookId, + name: 'Solo', + blocks: [ + { + id: 'original-block', + blockGroup: 'original-group', + type: 'code', + content: 'print("hello")', + sortingKey: '0', + metadata: {}, + executionCount: 7 + } + ], + executionMode: 'block' + } + ] + } + }; + + const yaml = serializeDeepnoteFile(projectData); + + const mockFS = mock(); + when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(yaml))); + when(mockFS.stat(anything())).thenReject(new Error('File not found')); + + const writes = new Map(); + when(mockFS.writeFile(anything(), anything())).thenCall((uri: Uri, content: Uint8Array) => { + writes.set(uri.path, content); + return Promise.resolve(); + }); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + + const workspaceFolder = { uri: Uri.file('/workspace'), name: 'workspace', index: 0 }; + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder as any]); + when(mockedVSCodeNamespaces.workspace.findFiles(anything())).thenReturn(Promise.resolve([fileUri])); + + uuidStubs.push(createUuidMock([clonedNotebookId, clonedBlockId, clonedBlockGroup])); + + when(mockedVSCodeNamespaces.workspace.openNotebookDocument(anything())).thenReturn( + Promise.resolve({ notebookType: 'deepnote' } as any) + ); + when(mockedVSCodeNamespaces.window.showNotebookDocument(anything(), anything())).thenReturn( + Promise.resolve(undefined as any) + ); + when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenReturn( + Promise.resolve(undefined) + ); + + const treeItem: Partial = { + type: DeepnoteTreeItemType.ProjectFile, + contextValue: NOTEBOOK_FILE_CONTEXT_VALUE, + context: { filePath: fileUri.fsPath, projectId }, + data: projectData as any + }; + + await explorerView.duplicateNotebook(treeItem as DeepnoteTreeItem); + + // Source file is not touched + expect(writes.has(fileUri.path)).to.be.false; + + // The new sibling file exists + const newEntry = Array.from(writes.entries()).find(([p]) => p !== fileUri.path); + + expect(newEntry).to.exist; + const [newPath, content] = newEntry!; + + expect(newPath).to.equal(expectedNewPath); + + const newDoc = deserializeDeepnoteFile(Buffer.from(content).toString('utf8')); + + // Single notebook inside — must have the NEW id, not the original id + expect(newDoc.project.notebooks).to.have.lengthOf(1); + const duplicated = newDoc.project.notebooks[0]; + + expect(duplicated.id).to.equal(clonedNotebookId); + expect(duplicated.id).to.not.equal(originalNotebookId); + expect(duplicated.name).to.equal('Solo (Copy)'); + expect(duplicated.blocks).to.have.lengthOf(1); + expect(duplicated.blocks[0].id).to.equal(clonedBlockId); + expect(duplicated.blocks[0].id).to.not.equal('original-block'); + expect(duplicated.blocks[0].blockGroup).to.equal(clonedBlockGroup); + expect(duplicated.blocks[0].blockGroup).to.not.equal('original-group'); + expect((duplicated.blocks[0] as ExecutableBlock).executionCount).to.be.undefined; + + // Project id is preserved (same project group) + expect(newDoc.project.id).to.equal(projectId); + }); + }); + + suite('deleteNotebook (single-notebook file)', () => { + test('deletes the sibling file instead of rewriting it', async () => { + const projectId = 'project-id'; + const notebookId = 'sole-nb'; + const fileUri = Uri.file('/workspace/solo.deepnote'); + + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { createdAt: '2024-01-01T00:00:00.000Z', modifiedAt: '2024-01-01T00:00:00.000Z' }, + project: { + id: projectId, + name: 'Solo', + notebooks: [{ id: notebookId, name: 'Sole Notebook', blocks: [], executionMode: 'block' }] + } + }; + + const mockFS = mock(); + when(mockFS.readFile(anything())).thenReturn( + Promise.resolve(Buffer.from(serializeDeepnoteFile(projectData))) + ); + + const deletedPaths: string[] = []; + when(mockFS.delete(anything())).thenCall((uri: Uri) => { + deletedPaths.push(uri.path); + return Promise.resolve(); + }); + + let writeCalled = false; + when(mockFS.writeFile(anything(), anything())).thenCall(() => { + writeCalled = true; + return Promise.resolve(); + }); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + + when(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).thenReturn( + Promise.resolve('Delete') as any + ); + when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenReturn( + Promise.resolve(undefined) + ); + + const treeItem: Partial = { + type: DeepnoteTreeItemType.ProjectFile, + contextValue: NOTEBOOK_FILE_CONTEXT_VALUE, + context: { filePath: fileUri.fsPath, projectId }, + data: projectData as any + }; + + await explorerView.deleteNotebook(treeItem as DeepnoteTreeItem); + + assert.deepStrictEqual(deletedPaths, [fileUri.path]); + assert.isFalse(writeCalled); + }); + }); +}); diff --git a/src/notebooks/deepnote/deepnoteTreeDataProvider.ts b/src/notebooks/deepnote/deepnoteTreeDataProvider.ts index 96c4bc2d1e..5587ccc759 100644 --- a/src/notebooks/deepnote/deepnoteTreeDataProvider.ts +++ b/src/notebooks/deepnote/deepnoteTreeDataProvider.ts @@ -204,30 +204,6 @@ export class DeepnoteTreeDataProvider implements TreeDataProvider nb.id !== initNotebookId) ?? []; - - if (nonInitNotebooks.length <= 1) { - const context: DeepnoteTreeItemContext = { - filePath: file.filePath, - projectId - }; - - const treeItem = new DeepnoteTreeItem( - DeepnoteTreeItemType.ProjectFile, - context, - file.project, - TreeItemCollapsibleState.None - ); - groups.push(treeItem); - continue; - } - } - - // Multiple files or multi-notebook file: create a group const groupData: ProjectGroupData = { projectId, projectName, @@ -239,11 +215,15 @@ export class DeepnoteTreeDataProvider implements TreeDataProvider { let provider: DeepnoteTreeDataProvider; @@ -338,11 +346,11 @@ suite('DeepnoteTreeDataProvider', () => { }); }); - test('should update visual fields when project data changes', async () => { + test('should update visual fields when project data changes from single to multi notebook', async () => { // Access private caches const treeItemCache = (provider as any).treeItemCache as Map; - // Create initial project with 1 notebook + // Create initial project with 1 notebook (single-notebook file -> label = notebook name) const filePath = '/workspace/test-project.deepnote'; const cacheKey = `project:${filePath}`; const initialProject: DeepnoteProject = { @@ -378,11 +386,12 @@ suite('DeepnoteTreeDataProvider', () => { ); treeItemCache.set(cacheKey, mockTreeItem); - // Verify initial state - assert.strictEqual(mockTreeItem.label, 'test-project.deepnote'); + // With a single non-init notebook the label adopts the notebook's name. + assert.strictEqual(mockTreeItem.label, 'Notebook 1'); + assert.strictEqual(mockTreeItem.contextValue, NOTEBOOK_FILE_CONTEXT_VALUE); assert.strictEqual(mockTreeItem.description, '0 cells'); - // Update the project data (simulating rename and adding notebooks) + // Upgrade to a multi-notebook project (simulating a rename + add) const updatedProject: DeepnoteProject = { ...initialProject, project: { @@ -402,25 +411,14 @@ suite('DeepnoteTreeDataProvider', () => { }; mockTreeItem.data = updatedProject; - // Call updateVisualFields if it exists (it may not work properly in test environment due to proxy limitations) - if (typeof mockTreeItem.updateVisualFields === 'function') { - mockTreeItem.updateVisualFields(); - } else { - // Manually update visual fields for testing purposes - const fileName = mockTreeItem.context.filePath.split('/').pop() ?? ''; - mockTreeItem.label = fileName || updatedProject.project.name || 'Untitled Project'; - mockTreeItem.tooltip = `Deepnote Project: ${updatedProject.project.name}\nFile: ${mockTreeItem.context.filePath}`; - const blockCount = - updatedProject.project.notebooks?.reduce( - (sum: number, nb: any) => sum + (nb.blocks?.length ?? 0), - 0 - ) ?? 0; - mockTreeItem.description = `${blockCount} cell${blockCount !== 1 ? 's' : ''}`; - } + // VS Code's mock TreeItem loses subclass prototype methods under the Proxy-based + // class mock, so invoke updateVisualFields via the class prototype directly. + DeepnoteTreeItem.prototype.updateVisualFields.call(mockTreeItem); - // Verify visual fields were updated - assert.strictEqual(mockTreeItem.label, 'test-project.deepnote', 'Label should reflect filename'); - assert.strictEqual(mockTreeItem.description, '0 cells', 'Description should reflect cell count'); + // Now the row is multi-notebook: label reverts to basename and contextValue reverts to projectFile + assert.strictEqual(mockTreeItem.label, 'test-project.deepnote'); + assert.strictEqual(mockTreeItem.contextValue, 'projectFile'); + assert.strictEqual(mockTreeItem.description, '0 cells'); assert.include( mockTreeItem.tooltip as string, 'Renamed Project', @@ -650,3 +648,152 @@ suite('DeepnoteTreeDataProvider', () => { }); }); }); + +suite('DeepnoteTreeDataProvider - getProjectGroups (no flattening)', () => { + let provider: DeepnoteTreeDataProvider; + + setup(() => { + resetVSCodeMocks(); + provider = new DeepnoteTreeDataProvider(); + }); + + teardown(() => { + provider.dispose(); + resetVSCodeMocks(); + }); + + async function invokeGetProjectGroups( + files: Array<{ uri: Uri; project: DeepnoteProject }> + ): Promise { + const workspaceFolder = { uri: Uri.file('/workspace'), name: 'workspace', index: 0 }; + + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder as any]); + when(mockedVSCodeNamespaces.workspace.findFiles(anything())).thenReturn( + Promise.resolve(files.map((f) => f.uri)) + ); + + const mockFS = mock(); + + when(mockFS.readFile(anything())).thenCall((uri: Uri) => { + const entry = files.find((f) => f.uri.path === uri.path); + + if (!entry) { + return Promise.reject(new Error('File not found')); + } + + return Promise.resolve(Buffer.from(serializeDeepnoteFile(entry.project))); + }); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + + return (provider as any).getProjectGroups(); + } + + test('single-file project with one non-init notebook emits a ProjectGroup (Expanded)', async () => { + const project: DeepnoteProject = { + version: '1.0.0', + metadata: { + createdAt: '2024-01-01T00:00:00.000Z', + modifiedAt: '2024-01-01T00:00:00.000Z' + }, + project: { + id: 'proj-1', + name: 'Single File', + notebooks: [ + { + id: 'nb-1', + name: 'The Notebook', + blocks: [], + executionMode: 'block' + } + ] + } + }; + + const groups = await invokeGetProjectGroups([{ uri: Uri.file('/workspace/single.deepnote'), project }]); + + assert.strictEqual(groups.length, 1); + assert.strictEqual(groups[0].type, DeepnoteTreeItemType.ProjectGroup); + assert.strictEqual(groups[0].contextValue, 'projectGroup'); + assert.strictEqual(groups[0].collapsibleState, TreeItemCollapsibleState.Expanded); + + const data = groups[0].data as ProjectGroupData; + + assert.strictEqual(data.projectId, 'proj-1'); + assert.strictEqual(data.projectName, 'Single File'); + assert.strictEqual(data.files.length, 1); + assert.strictEqual(data.files[0].filePath, '/workspace/single.deepnote'); + }); + + test('single-file project with multiple notebooks emits a ProjectGroup', async () => { + const project: DeepnoteProject = { + version: '1.0.0', + metadata: { + createdAt: '2024-01-01T00:00:00.000Z', + modifiedAt: '2024-01-01T00:00:00.000Z' + }, + project: { + id: 'proj-2', + name: 'Multi Notebook', + notebooks: [ + { id: 'nb-a', name: 'A', blocks: [], executionMode: 'block' }, + { id: 'nb-b', name: 'B', blocks: [], executionMode: 'block' } + ] + } + }; + + const groups = await invokeGetProjectGroups([{ uri: Uri.file('/workspace/multi.deepnote'), project }]); + + assert.strictEqual(groups.length, 1); + assert.strictEqual(groups[0].type, DeepnoteTreeItemType.ProjectGroup); + // Single file -> expanded by default even when the file carries multiple notebooks + assert.strictEqual(groups[0].collapsibleState, TreeItemCollapsibleState.Expanded); + + const data = groups[0].data as ProjectGroupData; + + assert.strictEqual(data.files.length, 1); + }); + + test('multi-file project emits a single ProjectGroup (Collapsed) that aggregates every file', async () => { + const fileA: DeepnoteProject = { + version: '1.0.0', + metadata: { + createdAt: '2024-01-01T00:00:00.000Z', + modifiedAt: '2024-01-01T00:00:00.000Z' + }, + project: { + id: 'proj-shared', + name: 'Shared Project', + notebooks: [{ id: 'nb-a', name: 'A', blocks: [], executionMode: 'block' }] + } + }; + + const fileB: DeepnoteProject = { + version: '1.0.0', + metadata: { + createdAt: '2024-01-01T00:00:00.000Z', + modifiedAt: '2024-01-01T00:00:00.000Z' + }, + project: { + id: 'proj-shared', + name: 'Shared Project', + notebooks: [{ id: 'nb-b', name: 'B', blocks: [], executionMode: 'block' }] + } + }; + + const groups = await invokeGetProjectGroups([ + { uri: Uri.file('/workspace/a.deepnote'), project: fileA }, + { uri: Uri.file('/workspace/b.deepnote'), project: fileB } + ]); + + assert.strictEqual(groups.length, 1); + assert.strictEqual(groups[0].type, DeepnoteTreeItemType.ProjectGroup); + assert.strictEqual(groups[0].collapsibleState, TreeItemCollapsibleState.Collapsed); + + const data = groups[0].data as ProjectGroupData; + + assert.strictEqual(data.files.length, 2); + const paths = data.files.map((f) => f.filePath).sort(); + + assert.deepStrictEqual(paths, ['/workspace/a.deepnote', '/workspace/b.deepnote']); + }); +}); diff --git a/src/notebooks/deepnote/deepnoteTreeItem.ts b/src/notebooks/deepnote/deepnoteTreeItem.ts index 5679334a64..9211098793 100644 --- a/src/notebooks/deepnote/deepnoteTreeItem.ts +++ b/src/notebooks/deepnote/deepnoteTreeItem.ts @@ -31,6 +31,12 @@ export interface ProjectGroupData { readonly files: Array<{ filePath: string; project: DeepnoteProject }>; } +/** + * contextValue assigned to a `.deepnote` file row that represents a single notebook + * (i.e., a `ProjectFile` that carries notebook-level actions in the tree menu). + */ +export const NOTEBOOK_FILE_CONTEXT_VALUE = 'notebookFile'; + /** * Tree item representing a Deepnote project group, file, or notebook in the explorer view */ @@ -76,8 +82,15 @@ export class DeepnoteTreeItem extends TreeItem { // getLabel() inline if (this.type === DeepnoteTreeItemType.ProjectFile) { const project = this.data as DeepnoteProject; - const fileName = basename(this.context.filePath); - this.label = fileName || project.project.name || 'Untitled Project'; + const singleNonInitNotebook = getSingleNonInitNotebook(project); + + if (singleNonInitNotebook) { + this.label = singleNonInitNotebook.name || project.project.name || 'Untitled Notebook'; + this.contextValue = NOTEBOOK_FILE_CONTEXT_VALUE; + } else { + const fileName = basename(this.context.filePath); + this.label = fileName || project.project.name || 'Untitled Project'; + } } else { const notebook = this.data as DeepnoteNotebook; this.label = notebook.name || 'Untitled Notebook'; @@ -139,8 +152,16 @@ export class DeepnoteTreeItem extends TreeItem { if (this.type === DeepnoteTreeItemType.ProjectFile) { const project = this.data as DeepnoteProject; - const fileName = basename(this.context.filePath); - this.label = fileName || project.project.name || 'Untitled Project'; + const singleNonInitNotebook = getSingleNonInitNotebook(project); + + if (singleNonInitNotebook) { + this.label = singleNonInitNotebook.name || project.project.name || 'Untitled Notebook'; + this.contextValue = NOTEBOOK_FILE_CONTEXT_VALUE; + } else { + const fileName = basename(this.context.filePath); + this.label = fileName || project.project.name || 'Untitled Project'; + this.contextValue = this.type; + } this.tooltip = `Deepnote Project: ${project.project.name}\nFile: ${this.context.filePath}`; const initNotebookId = project.project.initNotebookId; @@ -159,3 +180,15 @@ export class DeepnoteTreeItem extends TreeItem { } } } + +/** + * Returns the sole non-init notebook on a project, or undefined if the project has zero + * or multiple non-init notebooks. Used to decide whether a `ProjectFile` row should act as + * a notebook (label + notebook actions) versus a legacy multi-notebook container. + */ +export function getSingleNonInitNotebook(project: DeepnoteProject): DeepnoteNotebook | undefined { + const initNotebookId = project.project.initNotebookId; + const nonInitNotebooks = project.project.notebooks?.filter((nb) => nb.id !== initNotebookId) ?? []; + + return nonInitNotebooks.length === 1 ? nonInitNotebooks[0] : undefined; +} diff --git a/src/notebooks/deepnote/deepnoteTreeItem.unit.test.ts b/src/notebooks/deepnote/deepnoteTreeItem.unit.test.ts index cbc2a99bf3..e09042a0eb 100644 --- a/src/notebooks/deepnote/deepnoteTreeItem.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteTreeItem.unit.test.ts @@ -1,11 +1,19 @@ import { assert } from 'chai'; import { TreeItemCollapsibleState, ThemeIcon } from 'vscode'; -import { DeepnoteTreeItem, DeepnoteTreeItemType, DeepnoteTreeItemContext } from './deepnoteTreeItem'; +import { + DeepnoteTreeItem, + DeepnoteTreeItemType, + DeepnoteTreeItemContext, + NOTEBOOK_FILE_CONTEXT_VALUE, + getSingleNonInitNotebook +} from './deepnoteTreeItem'; import type { DeepnoteProject, DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; suite('DeepnoteTreeItem', () => { - const mockProject: DeepnoteProject = { + // A project with a single non-init notebook. Under the refactored rules this means + // the `ProjectFile` row should adopt the notebook's label and `notebookFile` contextValue. + const singleNotebookProject: DeepnoteProject = { metadata: { createdAt: '2023-01-01T00:00:00Z', modifiedAt: '2023-01-02T00:00:00Z' @@ -27,6 +35,37 @@ suite('DeepnoteTreeItem', () => { version: '1.0.0' }; + // A legacy multi-notebook project. The `ProjectFile` row should use the file basename + // as its label and keep the default `projectFile` contextValue. + const multiNotebookProject: DeepnoteProject = { + metadata: { + createdAt: '2023-01-01T00:00:00Z', + modifiedAt: '2023-01-02T00:00:00Z' + }, + project: { + id: 'project-123', + name: 'Test Project', + notebooks: [ + { + id: 'notebook-1', + name: 'First Notebook', + blocks: [], + executionMode: 'block', + isModule: false + }, + { + id: 'notebook-2', + name: 'Second Notebook', + blocks: [], + executionMode: 'block', + isModule: false + } + ], + settings: {} + }, + version: '1.0.0' + }; + const mockNotebook: DeepnoteNotebook = { id: 'notebook-456', name: 'Analysis Notebook', @@ -45,7 +84,7 @@ suite('DeepnoteTreeItem', () => { }; suite('constructor', () => { - test('should create project file item with basic properties', () => { + test('should create multi-notebook project file item with basic properties', () => { const context: DeepnoteTreeItemContext = { filePath: '/test/project.deepnote', projectId: 'project-123' @@ -54,7 +93,7 @@ suite('DeepnoteTreeItem', () => { const item = new DeepnoteTreeItem( DeepnoteTreeItemType.ProjectFile, context, - mockProject, + multiNotebookProject, TreeItemCollapsibleState.Collapsed ); @@ -95,7 +134,7 @@ suite('DeepnoteTreeItem', () => { const item = new DeepnoteTreeItem( DeepnoteTreeItemType.ProjectFile, context, - mockProject, + multiNotebookProject, TreeItemCollapsibleState.Expanded ); @@ -103,7 +142,136 @@ suite('DeepnoteTreeItem', () => { }); }); - suite('ProjectFile type', () => { + suite('ProjectFile type (single non-init notebook)', () => { + test('should use the sole notebook name as label and notebookFile contextValue', () => { + const context: DeepnoteTreeItemContext = { + filePath: '/workspace/my-project.deepnote', + projectId: 'project-456' + }; + + const item = new DeepnoteTreeItem( + DeepnoteTreeItemType.ProjectFile, + context, + singleNotebookProject, + TreeItemCollapsibleState.None + ); + + assert.strictEqual(item.label, 'First Notebook'); + assert.strictEqual(item.contextValue, NOTEBOOK_FILE_CONTEXT_VALUE); + assert.strictEqual(item.contextValue, 'notebookFile'); + assert.strictEqual(item.tooltip, 'Deepnote Project: Test Project\nFile: /workspace/my-project.deepnote'); + }); + + test('should fall back to project name when notebook name is empty', () => { + const projectWithEmptyNotebookName: DeepnoteProject = { + ...singleNotebookProject, + project: { + ...singleNotebookProject.project, + notebooks: [ + { + id: 'notebook-1', + name: '', + blocks: [], + executionMode: 'block', + isModule: false + } + ] + } + }; + + const context: DeepnoteTreeItemContext = { + filePath: '/workspace/my-project.deepnote', + projectId: 'project-456' + }; + + const item = new DeepnoteTreeItem( + DeepnoteTreeItemType.ProjectFile, + context, + projectWithEmptyNotebookName, + TreeItemCollapsibleState.None + ); + + assert.strictEqual(item.label, 'Test Project'); + assert.strictEqual(item.contextValue, NOTEBOOK_FILE_CONTEXT_VALUE); + }); + + test('should fall back to Untitled Notebook when both notebook and project names are empty', () => { + const projectWithNoNames: DeepnoteProject = { + ...singleNotebookProject, + project: { + ...singleNotebookProject.project, + name: '', + notebooks: [ + { + id: 'notebook-1', + name: '', + blocks: [], + executionMode: 'block', + isModule: false + } + ] + } + }; + + const context: DeepnoteTreeItemContext = { + filePath: '/workspace/my-project.deepnote', + projectId: 'project-456' + }; + + const item = new DeepnoteTreeItem( + DeepnoteTreeItemType.ProjectFile, + context, + projectWithNoNames, + TreeItemCollapsibleState.None + ); + + assert.strictEqual(item.label, 'Untitled Notebook'); + assert.strictEqual(item.contextValue, NOTEBOOK_FILE_CONTEXT_VALUE); + }); + + test('should treat a file with one non-init notebook as single-notebook even when an init notebook exists', () => { + const projectWithInit: DeepnoteProject = { + ...singleNotebookProject, + project: { + ...singleNotebookProject.project, + initNotebookId: 'init-notebook', + notebooks: [ + { + id: 'init-notebook', + name: 'Init', + blocks: [], + executionMode: 'block', + isModule: false + }, + { + id: 'notebook-1', + name: 'Only Real Notebook', + blocks: [], + executionMode: 'block', + isModule: false + } + ] + } + }; + + const context: DeepnoteTreeItemContext = { + filePath: '/workspace/my-project.deepnote', + projectId: 'project-456' + }; + + const item = new DeepnoteTreeItem( + DeepnoteTreeItemType.ProjectFile, + context, + projectWithInit, + TreeItemCollapsibleState.None + ); + + assert.strictEqual(item.label, 'Only Real Notebook'); + assert.strictEqual(item.contextValue, NOTEBOOK_FILE_CONTEXT_VALUE); + }); + }); + + suite('ProjectFile type (legacy multi-notebook)', () => { test('should have correct properties for project file', () => { const context: DeepnoteTreeItemContext = { filePath: '/workspace/my-project.deepnote', @@ -113,7 +281,7 @@ suite('DeepnoteTreeItem', () => { const item = new DeepnoteTreeItem( DeepnoteTreeItemType.ProjectFile, context, - mockProject, + multiNotebookProject, TreeItemCollapsibleState.Collapsed ); @@ -133,11 +301,11 @@ suite('DeepnoteTreeItem', () => { assert.strictEqual(item.command!.command, 'deepnote.openNotebook'); }); - test('should handle project with multiple notebooks', () => { + test('should handle project with three notebooks', () => { const projectWithMultipleNotebooks: DeepnoteProject = { - ...mockProject, + ...multiNotebookProject, project: { - ...mockProject.project, + ...multiNotebookProject.project, notebooks: [ { id: 'notebook-1', @@ -177,13 +345,15 @@ suite('DeepnoteTreeItem', () => { ); assert.strictEqual(item.description, '0 cells'); + assert.strictEqual(item.label, 'project.deepnote'); + assert.strictEqual(item.contextValue, 'projectFile'); }); - test('should handle project with no notebooks', () => { + test('should handle project with no notebooks (label falls back to basename)', () => { const projectWithNoNotebooks = { - ...mockProject, + ...multiNotebookProject, project: { - ...mockProject.project, + ...multiNotebookProject.project, notebooks: [] } }; @@ -201,13 +371,15 @@ suite('DeepnoteTreeItem', () => { ); assert.strictEqual(item.description, '0 cells'); + assert.strictEqual(item.label, 'project.deepnote'); + assert.strictEqual(item.contextValue, 'projectFile'); }); test('should handle unnamed project', () => { const unnamedProject = { - ...mockProject, + ...multiNotebookProject, project: { - ...mockProject.project, + ...multiNotebookProject.project, name: undefined } }; @@ -395,13 +567,22 @@ suite('DeepnoteTreeItem', () => { projectId: 'project-1' }; + // Multi-notebook project file -> contextValue 'projectFile' const projectItem = new DeepnoteTreeItem( DeepnoteTreeItemType.ProjectFile, baseContext, - mockProject, + multiNotebookProject, TreeItemCollapsibleState.Collapsed ); + // Single-notebook project file -> contextValue 'notebookFile' + const notebookFileItem = new DeepnoteTreeItem( + DeepnoteTreeItemType.ProjectFile, + baseContext, + singleNotebookProject, + TreeItemCollapsibleState.None + ); + const notebookItem = new DeepnoteTreeItem( DeepnoteTreeItemType.Notebook, { ...baseContext, notebookId: 'notebook-1' }, @@ -410,6 +591,7 @@ suite('DeepnoteTreeItem', () => { ); assert.strictEqual(projectItem.contextValue, 'projectFile'); + assert.strictEqual(notebookFileItem.contextValue, NOTEBOOK_FILE_CONTEXT_VALUE); assert.strictEqual(notebookItem.contextValue, 'notebook'); }); }); @@ -424,7 +606,7 @@ suite('DeepnoteTreeItem', () => { const item = new DeepnoteTreeItem( DeepnoteTreeItemType.ProjectFile, context, - mockProject, + multiNotebookProject, TreeItemCollapsibleState.Collapsed ); @@ -466,7 +648,7 @@ suite('DeepnoteTreeItem', () => { const item = new DeepnoteTreeItem( DeepnoteTreeItemType.ProjectFile, context, - mockProject, + multiNotebookProject, TreeItemCollapsibleState.Collapsed ); @@ -501,9 +683,9 @@ suite('DeepnoteTreeItem', () => { }; const projectWithName = { - ...mockProject, + ...multiNotebookProject, project: { - ...mockProject.project, + ...multiNotebookProject.project, name: 'My Amazing Project' } }; @@ -637,9 +819,147 @@ suite('DeepnoteTreeItem', () => { }); }); + suite('updateVisualFields', () => { + // VS Code's mock TreeItem is a Proxy-based class (see build/mocha-esm-loader.js); + // when DeepnoteTreeItem extends it, `super(...)` creates a plain TreeItem-mock + // instance rather than a DeepnoteTreeItem, so instance methods on the subclass + // prototype are NOT reachable via the usual `item.updateVisualFields()` call. + // We invoke the method by reading it off the class prototype and calling it on + // the instance — same behaviour, just side-steps the proxy's prototype chain loss. + function callUpdateVisualFields(item: DeepnoteTreeItem): void { + const method = DeepnoteTreeItem.prototype.updateVisualFields; + + assert.isFunction(method, 'updateVisualFields should be defined on the prototype'); + method.call(item); + } + + test('should keep notebook label and notebookFile contextValue for single-notebook ProjectFile', () => { + const context: DeepnoteTreeItemContext = { + filePath: '/workspace/my-project.deepnote', + projectId: 'project-456' + }; + + const item = new DeepnoteTreeItem( + DeepnoteTreeItemType.ProjectFile, + context, + singleNotebookProject, + TreeItemCollapsibleState.None + ); + + // Pre-conditions set in constructor + assert.strictEqual(item.label, 'First Notebook'); + assert.strictEqual(item.contextValue, NOTEBOOK_FILE_CONTEXT_VALUE); + + callUpdateVisualFields(item); + + assert.strictEqual(item.label, 'First Notebook'); + assert.strictEqual(item.contextValue, NOTEBOOK_FILE_CONTEXT_VALUE); + }); + + test('should keep file basename label and projectFile contextValue for multi-notebook ProjectFile', () => { + const context: DeepnoteTreeItemContext = { + filePath: '/workspace/my-project.deepnote', + projectId: 'project-456' + }; + + const item = new DeepnoteTreeItem( + DeepnoteTreeItemType.ProjectFile, + context, + multiNotebookProject, + TreeItemCollapsibleState.Collapsed + ); + + // Pre-conditions set in constructor + assert.strictEqual(item.label, 'my-project.deepnote'); + assert.strictEqual(item.contextValue, 'projectFile'); + + callUpdateVisualFields(item); + + assert.strictEqual(item.label, 'my-project.deepnote'); + assert.strictEqual(item.contextValue, 'projectFile'); + }); + + test('should flip contextValue back to projectFile when data is mutated from single to multi notebook', () => { + const context: DeepnoteTreeItemContext = { + filePath: '/workspace/my-project.deepnote', + projectId: 'project-456' + }; + + const item = new DeepnoteTreeItem( + DeepnoteTreeItemType.ProjectFile, + context, + singleNotebookProject, + TreeItemCollapsibleState.None + ); + + assert.strictEqual(item.contextValue, NOTEBOOK_FILE_CONTEXT_VALUE); + + // Mutate data to a multi-notebook project + item.data = multiNotebookProject; + callUpdateVisualFields(item); + + assert.strictEqual(item.contextValue, 'projectFile'); + assert.strictEqual(item.label, 'my-project.deepnote'); + }); + }); + + suite('getSingleNonInitNotebook', () => { + test('returns the sole non-init notebook when there is exactly one', () => { + const result = getSingleNonInitNotebook(singleNotebookProject); + + assert.isDefined(result); + assert.strictEqual(result!.id, 'notebook-1'); + assert.strictEqual(result!.name, 'First Notebook'); + }); + + test('returns undefined when there are multiple non-init notebooks', () => { + assert.isUndefined(getSingleNonInitNotebook(multiNotebookProject)); + }); + + test('returns undefined when there are zero notebooks', () => { + const empty: DeepnoteProject = { + ...singleNotebookProject, + project: { ...singleNotebookProject.project, notebooks: [] } + }; + + assert.isUndefined(getSingleNonInitNotebook(empty)); + }); + + test('excludes init notebook from the non-init count', () => { + const projectWithInitAndOneNotebook: DeepnoteProject = { + ...singleNotebookProject, + project: { + ...singleNotebookProject.project, + initNotebookId: 'init-nb', + notebooks: [ + { + id: 'init-nb', + name: 'Init', + blocks: [], + executionMode: 'block', + isModule: false + }, + { + id: 'notebook-1', + name: 'Real', + blocks: [], + executionMode: 'block', + isModule: false + } + ] + } + }; + + const result = getSingleNonInitNotebook(projectWithInitAndOneNotebook); + + assert.isDefined(result); + assert.strictEqual(result!.id, 'notebook-1'); + }); + }); + suite('integration scenarios', () => { test('should create valid tree structure hierarchy', () => { - // Create parent project file + // Create parent project file (multi-notebook so it stays a projectFile row) const projectContext: DeepnoteTreeItemContext = { filePath: '/workspace/research-project.deepnote', projectId: 'research-123' @@ -648,7 +968,7 @@ suite('DeepnoteTreeItem', () => { const projectItem = new DeepnoteTreeItem( DeepnoteTreeItemType.ProjectFile, projectContext, - mockProject, + multiNotebookProject, TreeItemCollapsibleState.Expanded ); @@ -737,7 +1057,7 @@ suite('DeepnoteTreeItem', () => { new DeepnoteTreeItem( DeepnoteTreeItemType.ProjectFile, context, - mockProject, + multiNotebookProject, TreeItemCollapsibleState.Collapsed ) ); From c872be8ebf220dfb38dc2359f5d2f6aae49e6022 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 20 Apr 2026 10:15:01 +0000 Subject: [PATCH 43/47] refactor(deepnote): Key operations by projectId for improved server management - Updated `DeepnoteServerStarter` to serialize concurrent operations based on `projectId`, allowing sibling `.deepnote` files to share a single server instance. - Modified method signatures and internal logic to replace file paths with project IDs, enhancing context handling and preventing race conditions. - Adjusted related interfaces and types to reflect the new project-based architecture. - Enhanced unit tests to validate the new behavior, ensuring correct server reuse across sibling files sharing the same project. --- .../deepnote/deepnoteServerStarter.node.ts | 118 +++-- .../deepnoteServerStarter.unit.test.ts | 143 +++++- .../deepnoteEnvironmentManager.node.ts | 6 +- .../deepnoteEnvironmentsView.node.ts | 75 +-- .../deepnoteEnvironmentsView.unit.test.ts | 147 +++--- .../deepnoteExtensionSidecarWriter.node.ts | 120 +---- ...eepnoteExtensionSidecarWriter.unit.test.ts | 202 ++------ .../deepnoteNotebookEnvironmentMapper.node.ts | 110 ----- .../deepnoteProjectEnvironmentMapper.node.ts | 168 +++++++ ...pnoteProjectEnvironmentMapper.unit.test.ts | 333 +++++++++++++ src/kernels/deepnote/types.ts | 63 +-- .../deepnoteKernelAutoSelector.node.ts | 109 +++-- ...epnoteKernelAutoSelector.node.unit.test.ts | 440 ++++-------------- .../deepnoteKernelStatusIndicator.node.ts | 10 +- .../deepnote/deepnoteKernelStatusIndicator.ts | 10 +- .../snapshots/environmentCapture.node.ts | 22 +- src/notebooks/serviceRegistry.node.ts | 12 +- .../deepnote/deepnoteProjectIdResolver.ts | 36 ++ .../deepnote/deepnoteServerUtils.node.ts | 6 +- 19 files changed, 1148 insertions(+), 982 deletions(-) delete mode 100644 src/kernels/deepnote/environments/deepnoteNotebookEnvironmentMapper.node.ts create mode 100644 src/kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.node.ts create mode 100644 src/kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.unit.test.ts create mode 100644 src/platform/deepnote/deepnoteProjectIdResolver.ts diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 8c1d51717c..eb01c0ae71 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -94,7 +94,8 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension /** * Start a server for a kernel environment. - * Serializes concurrent operations on the same environment to prevent race conditions. + * Serializes concurrent operations on the same project to prevent race conditions. + * Keyed by `projectId` so sibling `.deepnote` files sharing a project reuse one server. */ public async startServer( interpreter: PythonEnvironment, @@ -102,14 +103,15 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension managedVenv: boolean, additionalPackages: string[], environmentId: string, + projectId: string, deepnoteFileUri: Uri, token?: CancellationToken ): Promise { - const fileKey = deepnoteFileUri.fsPath; + const projectKey = projectId; - let pendingOp = this.pendingOperations.get(fileKey); + let pendingOp = this.pendingOperations.get(projectKey); if (pendingOp) { - logger.info(`Waiting for pending operation on ${fileKey} to complete...`); + logger.info(`Waiting for pending operation on project ${projectKey} to complete...`); try { await pendingOp.promise; } catch { @@ -117,28 +119,28 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } - let existingContext = this.projectContexts.get(fileKey); + let existingContext = this.projectContexts.get(projectKey); if (existingContext != null) { const { environmentId: existingEnvironmentId, serverInfo: existingServerInfo } = existingContext; if (existingEnvironmentId === environmentId) { if (existingServerInfo != null && (await this.isServerRunning(existingServerInfo))) { logger.info( - `Deepnote server already running at ${existingServerInfo.url} for ${fileKey} (environmentId ${environmentId})` + `Deepnote server already running at ${existingServerInfo.url} for project ${projectKey} (environmentId ${environmentId})` ); return existingServerInfo; } - pendingOp = this.pendingOperations.get(fileKey); + pendingOp = this.pendingOperations.get(projectKey); if (pendingOp && pendingOp.type === 'start') { return await pendingOp.promise; } } else { logger.info( - `Stopping existing server for ${fileKey} with environmentId ${existingEnvironmentId} to start new one with environmentId ${environmentId}...` + `Stopping existing server for project ${projectKey} with environmentId ${existingEnvironmentId} to start new one with environmentId ${environmentId}...` ); - await this.stopServerForEnvironment(existingContext, deepnoteFileUri, token); + await this.stopServerForEnvironment(existingContext, projectKey, token); existingContext.environmentId = environmentId; } } else { @@ -147,7 +149,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension serverInfo: null }; - this.projectContexts.set(fileKey, newContext); + this.projectContexts.set(projectKey, newContext); existingContext = newContext; } @@ -160,11 +162,12 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension managedVenv, additionalPackages, environmentId, + projectKey, deepnoteFileUri, token ) }; - this.pendingOperations.set(fileKey, operation); + this.pendingOperations.set(projectKey, operation); try { const result = await operation.promise; @@ -172,29 +175,29 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension existingContext.serverInfo = result; return result; } finally { - if (this.pendingOperations.get(fileKey) === operation) { - this.pendingOperations.delete(fileKey); + if (this.pendingOperations.get(projectKey) === operation) { + this.pendingOperations.delete(projectKey); } } } /** - * Stop the deepnote-toolkit server for a kernel environment. + * Stop the deepnote-toolkit server for a project. */ - public async stopServer(deepnoteFileUri: Uri, token?: CancellationToken): Promise { + public async stopServer(projectId: string, token?: CancellationToken): Promise { Cancellation.throwIfCanceled(token); - const fileKey = deepnoteFileUri.fsPath; - const projectContext = this.projectContexts.get(fileKey) ?? null; + const projectKey = projectId; + const projectContext = this.projectContexts.get(projectKey) ?? null; if (projectContext == null) { - logger.warn(`No project context found for ${fileKey}, skipping stop server...`); + logger.warn(`No project context found for project ${projectKey}, skipping stop server...`); return; } - const pendingOp = this.pendingOperations.get(fileKey); + const pendingOp = this.pendingOperations.get(projectKey); if (pendingOp) { - logger.info(`Waiting for pending operation on ${fileKey} before stopping...`); + logger.info(`Waiting for pending operation on project ${projectKey} before stopping...`); try { await pendingOp.promise; } catch { @@ -206,15 +209,15 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const operation = { type: 'stop' as const, - promise: this.stopServerForEnvironment(projectContext, deepnoteFileUri, token) + promise: this.stopServerForEnvironment(projectContext, projectKey, token) }; - this.pendingOperations.set(fileKey, operation); + this.pendingOperations.set(projectKey, operation); try { await operation.promise; } finally { - if (this.pendingOperations.get(fileKey) === operation) { - this.pendingOperations.delete(fileKey); + if (this.pendingOperations.get(projectKey) === operation) { + this.pendingOperations.delete(projectKey); } } } @@ -235,11 +238,10 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension managedVenv: boolean, additionalPackages: string[], environmentId: string, + projectKey: string, deepnoteFileUri: Uri, token?: CancellationToken ): Promise { - const fileKey = deepnoteFileUri.fsPath; - Cancellation.throwIfCanceled(token); logger.info(`Ensuring deepnote-toolkit is installed in venv for environment ${environmentId}...`); @@ -258,13 +260,13 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); - logger.info(`Starting deepnote-toolkit server for ${fileKey} (environmentId ${environmentId})`); + logger.info(`Starting deepnote-toolkit server for project ${projectKey} (environmentId ${environmentId})`); this.outputChannel.appendLine(l10n.t('Starting Deepnote server...')); const extraEnv = await this.gatherSqlIntegrationEnvVars(deepnoteFileUri, environmentId, token); // Initialize output tracking for error reporting - this.serverOutputByFile.set(fileKey, { stdout: '', stderr: '' }); + this.serverOutputByFile.set(projectKey, { stdout: '', stderr: '' }); let serverInfo: DeepnoteServerInfo | undefined; try { @@ -275,8 +277,8 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension env: extraEnv }); } catch (error) { - const capturedOutput = this.serverOutputByFile.get(fileKey); - this.serverOutputByFile.delete(fileKey); + const capturedOutput = this.serverOutputByFile.get(projectKey); + this.serverOutputByFile.delete(projectKey); throw new DeepnoteServerStartupError( interpreter.uri.fsPath, @@ -291,17 +293,17 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension projectContext.serverInfo = serverInfo; // Set up output channel logging from the server process - this.monitorServerOutput(fileKey, serverInfo); + this.monitorServerOutput(projectKey, serverInfo); // Write lock file for orphan-cleanup tracking const serverPid = serverInfo.process.pid; if (serverPid) { await this.writeLockFile(serverPid); } else { - logger.warn(`Could not get PID for server process for ${fileKey}`); + logger.warn(`Could not get PID for server process for project ${projectKey}`); } - logger.info(`Deepnote server started successfully at ${serverInfo.url} for ${fileKey}`); + logger.info(`Deepnote server started successfully at ${serverInfo.url} for project ${projectKey}`); this.outputChannel.appendLine(l10n.t('✓ Deepnote server running at {0}', serverInfo.url)); return serverInfo; @@ -312,11 +314,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension */ private async stopServerForEnvironment( projectContext: ProjectContext, - deepnoteFileUri: Uri, + projectKey: string, token?: CancellationToken ): Promise { - const fileKey = deepnoteFileUri.fsPath; - Cancellation.throwIfCanceled(token); const { serverInfo } = projectContext; @@ -325,9 +325,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const serverPid = serverInfo.process.pid; try { - logger.info(`Stopping Deepnote server for ${fileKey}...`); + logger.info(`Stopping Deepnote server for project ${projectKey}...`); await stopServer(serverInfo); - this.outputChannel.appendLine(l10n.t('Deepnote server stopped for {0}', fileKey)); + this.outputChannel.appendLine(l10n.t('Deepnote server stopped for project {0}', projectKey)); } catch (ex) { logger.error('Error stopping Deepnote server', ex); } finally { @@ -341,9 +341,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); - this.serverOutputByFile.delete(fileKey); + this.serverOutputByFile.delete(projectKey); - this.disposeOutputListeners(fileKey); + this.disposeOutputListeners(projectKey); } /** @@ -373,10 +373,8 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension return extraEnv; } - const fileKey = deepnoteFileUri.fsPath; - logger.debug( - `DeepnoteServerStarter: Injecting SQL integration env vars for ${fileKey} with environmentId ${environmentId}` + `DeepnoteServerStarter: Injecting SQL integration env vars for ${deepnoteFileUri.fsPath} with environmentId ${environmentId}` ); try { const sqlEnvVars = await this.sqlIntegrationEnvVars.getEnvironmentVariables(deepnoteFileUri, token); @@ -396,20 +394,20 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension /** * Stream stdout/stderr from the server process to the VSCode output channel. */ - private monitorServerOutput(fileKey: string, serverInfo: DeepnoteServerInfo): void { + private monitorServerOutput(projectKey: string, serverInfo: DeepnoteServerInfo): void { const proc = serverInfo.process; - this.disposeOutputListeners(fileKey); + this.disposeOutputListeners(projectKey); const disposables: IDisposable[] = []; - this.disposablesByFile.set(fileKey, disposables); + this.disposablesByFile.set(projectKey, disposables); if (proc.stdout) { const stdout = proc.stdout; const onData = (data: Buffer) => { const text = data.toString(); - logger.trace(`Deepnote server (${fileKey}): ${text}`); + logger.trace(`Deepnote server (${projectKey}): ${text}`); this.outputChannel.appendLine(text); - const outputTracking = this.serverOutputByFile.get(fileKey); + const outputTracking = this.serverOutputByFile.get(projectKey); if (outputTracking) { outputTracking.stdout = (outputTracking.stdout + text).slice(-MAX_OUTPUT_TRACKING_LENGTH); } @@ -426,10 +424,10 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const stderr = proc.stderr; const onData = (data: Buffer) => { const text = data.toString(); - logger.warn(`Deepnote server stderr (${fileKey}): ${text}`); + logger.warn(`Deepnote server stderr (${projectKey}): ${text}`); this.outputChannel.appendLine(text); - const outputTracking = this.serverOutputByFile.get(fileKey); + const outputTracking = this.serverOutputByFile.get(projectKey); if (outputTracking) { outputTracking.stderr = (outputTracking.stderr + text).slice(-MAX_OUTPUT_TRACKING_LENGTH); } @@ -444,19 +442,19 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } /** - * Dispose all output listeners registered for a given file key. + * Dispose all output listeners registered for a given project key. * Per-listener try/catch ensures one failure doesn't prevent others from being disposed. */ - private disposeOutputListeners(fileKey: string): void { - const disposables = this.disposablesByFile.get(fileKey) ?? []; + private disposeOutputListeners(projectKey: string): void { + const disposables = this.disposablesByFile.get(projectKey) ?? []; for (const d of disposables) { try { d.dispose(); } catch (ex) { - logger.warn(`Error disposing output listener for ${fileKey}`, ex); + logger.warn(`Error disposing output listener for project ${projectKey}`, ex); } } - this.disposablesByFile.delete(fileKey); + this.disposablesByFile.delete(projectKey); } public async dispose(): Promise { @@ -473,17 +471,17 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const stopPromises: Promise[] = []; const pidsToCleanup: number[] = []; - for (const [key, ctx] of this.projectContexts.entries()) { + for (const [projectKey, ctx] of this.projectContexts.entries()) { if (ctx.serverInfo) { const pid = ctx.serverInfo.process.pid; if (pid) { pidsToCleanup.push(pid); } - logger.info(`Stopping Deepnote server for ${key}...`); + logger.info(`Stopping Deepnote server for project ${projectKey}...`); stopPromises.push( stopServer(ctx.serverInfo).catch((ex) => { - logger.error(`Error stopping Deepnote server for ${key}`, ex); + logger.error(`Error stopping Deepnote server for project ${projectKey}`, ex); }) ); } @@ -499,8 +497,8 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } // Snapshot the keys first: disposeOutputListeners mutates disposablesByFile via .delete(). - for (const fileKey of Array.from(this.disposablesByFile.keys())) { - this.disposeOutputListeners(fileKey); + for (const projectKey of Array.from(this.disposablesByFile.keys())) { + this.disposeOutputListeners(projectKey); } this.disposablesByFile.clear(); diff --git a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts index f90e3a8ec1..feb9cb2f27 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts @@ -1,20 +1,26 @@ import { assert } from 'chai'; import * as fakeTimers from '@sinonjs/fake-timers'; +import esmock from 'esmock'; +import * as sinon from 'sinon'; import { anything, instance, mock, when } from 'ts-mockito'; +import { Uri } from 'vscode'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; +import { createMockChildProcess } from './deepnoteTestHelpers.node'; import { DeepnoteServerStarter } from './deepnoteServerStarter.node'; import { IProcessServiceFactory } from '../../platform/common/process/types.node'; import { IAsyncDisposableRegistry, IOutputChannel } from '../../platform/common/types'; import { IDeepnoteToolkitInstaller } from './types'; import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; +import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; /** * Unit tests for DeepnoteServerStarter. * * Port allocation, server spawning, and health checks are delegated to * @deepnote/runtime-core's startServer/stopServer. These tests focus on the - * extension-specific layers: SQL env var gathering and lifecycle orchestration. + * extension-specific layers: SQL env var gathering, lifecycle orchestration, + * and project-id keyed reuse of a single server across sibling files. */ suite('DeepnoteServerStarter', () => { let serverStarter: DeepnoteServerStarter; @@ -69,7 +75,6 @@ suite('DeepnoteServerStarter', () => { ); const gatherEnvVars = getPrivateMethod(starterWithoutSql, 'gatherSqlIntegrationEnvVars'); - const { Uri } = await import('vscode'); const result = await gatherEnvVars(Uri.file('/test/file.deepnote'), 'env1'); assert.deepStrictEqual(result, {}); @@ -78,7 +83,7 @@ suite('DeepnoteServerStarter', () => { }); test('should return empty object when provider rejects with cancellation error', async () => { - const { CancellationError, Uri } = await import('vscode'); + const { CancellationError } = await import('vscode'); const cancelledProvider = mock(); when(cancelledProvider.getEnvironmentVariables(anything(), anything())).thenReject(new CancellationError()); @@ -132,7 +137,8 @@ suite('DeepnoteServerStarter', () => { resolveDeferred = resolve; }); - starter.pendingOperations.set('/test/inflight.deepnote', { + // Internal maps are now keyed by projectId, not fsPath + starter.pendingOperations.set('project-id-inflight', { type: 'stop', promise: deferred }); @@ -157,4 +163,133 @@ suite('DeepnoteServerStarter', () => { assert.strictEqual(starter.pendingOperations.size, 0); }); }); + + /** + * Verifies the core plan invariant: two sibling `.deepnote` files that share + * the same `projectId` must reuse a single underlying server process. We + * assert this by mocking `@deepnote/runtime-core`'s `startServer` and + * checking it's only called once, even though `startServer` is invoked + * twice with different `deepnoteFileUri`s but the same `projectId`. + */ + suite('shared server for siblings with same projectId', () => { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + let mockedStarterModule: any; + let runtimeCoreStartServer: sinon.SinonStub; + let runtimeCoreStopServer: sinon.SinonStub; + let localFetchStub: sinon.SinonStub; + + setup(async () => { + runtimeCoreStartServer = sinon.stub(); + runtimeCoreStopServer = sinon.stub().resolves(); + + mockedStarterModule = await esmock('./deepnoteServerStarter.node', { + '@deepnote/runtime-core': { + startServer: runtimeCoreStartServer, + stopServer: runtimeCoreStopServer + } + }); + + // Stub fetch so the existing-server health check path is testable. + localFetchStub = sinon.stub(globalThis, 'fetch'); + }); + + teardown(() => { + localFetchStub?.restore(); + esmock.purge(mockedStarterModule); + }); + + test('two deepnoteFileUris sharing one projectId reuse a single server', async () => { + const toolkitInstaller = mock(); + when(toolkitInstaller.ensureVenvAndToolkit(anything(), anything(), anything(), anything())).thenResolve({ + pythonInterpreter: { + id: '/venvs/env1/bin/python', + uri: Uri.file('/venvs/env1/bin/python') + } as PythonEnvironment, + toolkitVersion: '2.0.0' + }); + when(toolkitInstaller.installAdditionalPackages(anything(), anything(), anything())).thenResolve(); + + const agentSkillsManager = mock(); + when(agentSkillsManager.ensureSkillsUpdated(anything(), anything())).thenReturn(); + + const outputChannel = mock(); + when(outputChannel.appendLine(anything())).thenReturn(); + + const processServiceFactory = mock(); + const asyncRegistry = mock(); + when(asyncRegistry.push(anything())).thenReturn(); + + const mockedProcess = createMockChildProcess({ pid: 12345 }); + + const firstServerInfo = { + url: 'http://localhost:8899', + jupyterPort: 8899, + lspPort: 8900, + process: mockedProcess + }; + + runtimeCoreStartServer.resolves(firstServerInfo); + + // Force `isServerRunning` to return true on the second call, so the + // existing server context is reused instead of restarted. + localFetchStub.resolves({ ok: true }); + + const StarterCtor = mockedStarterModule.DeepnoteServerStarter; + const starter = new StarterCtor( + instance(processServiceFactory), + instance(toolkitInstaller), + instance(agentSkillsManager), + instance(outputChannel), + instance(asyncRegistry) + ); + + try { + const interpreter = { + id: '/usr/bin/python3', + uri: Uri.file('/usr/bin/python3') + } as PythonEnvironment; + + const venvPath = Uri.file('/venvs/env1'); + const projectId = 'shared-project-id'; + + // Sibling 1: file A sharing the project + const resultA = await starter.startServer( + interpreter, + venvPath, + true, // managedVenv + [], // additionalPackages + 'env1', // environmentId + projectId, + Uri.file('/workspace/a.deepnote') + ); + + // Sibling 2: file B sharing the SAME projectId + const resultB = await starter.startServer( + interpreter, + venvPath, + true, + [], + 'env1', + projectId, + Uri.file('/workspace/b.deepnote') + ); + + assert.strictEqual( + runtimeCoreStartServer.callCount, + 1, + 'Underlying @deepnote/runtime-core startServer should only be invoked once for two siblings sharing a projectId' + ); + assert.strictEqual(resultA, firstServerInfo); + assert.strictEqual(resultB, firstServerInfo, 'Second call should return the same ServerInfo'); + + // Only one project context should exist, keyed by projectId + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const projectContexts = (starter as any).projectContexts as Map; + assert.strictEqual(projectContexts.size, 1); + assert.strictEqual(projectContexts.has(projectId), true); + } finally { + await starter.dispose(); + } + }); + }); }); diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentManager.node.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentManager.node.ts index 4a156d750d..c5e2a6ec1e 100644 --- a/src/kernels/deepnote/environments/deepnoteEnvironmentManager.node.ts +++ b/src/kernels/deepnote/environments/deepnoteEnvironmentManager.node.ts @@ -24,7 +24,7 @@ export class DeepnoteEnvironmentManager implements IExtensionSyncActivationServi // private readonly notebookServerHandles = new Map(); private environments: Map = new Map(); - private environmentServers: Map = new Map(); + private environmentServers: Map = new Map(); private readonly _onDidChangeEnvironments = new EventEmitter(); public readonly onDidChangeEnvironments = this._onDidChangeEnvironments.event; private initializationPromise: Promise | undefined; @@ -207,8 +207,8 @@ export class DeepnoteEnvironmentManager implements IExtensionSyncActivationServi } // Stop the server if running - for (const fileKey of this.environmentServers.get(id) ?? []) { - await this.serverStarter.stopServer(fileKey, token); + for (const projectId of this.environmentServers.get(id) ?? []) { + await this.serverStarter.stopServer(projectId, token); Cancellation.throwIfCanceled(token); } diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts index 4d6c3ec7fa..d1ccf53885 100644 --- a/src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts +++ b/src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts @@ -14,6 +14,7 @@ import { IPythonApiProvider } from '../../../platform/api/types'; import { STANDARD_OUTPUT_CHANNEL } from '../../../platform/common/constants'; import { getDisplayPath } from '../../../platform/common/platform/fs-paths.node'; import { IDisposableRegistry, IOutputChannel } from '../../../platform/common/types'; +import { resolveProjectIdForNotebook } from '../../../platform/deepnote/deepnoteProjectIdResolver'; import { createDeepnoteServerConfigHandle } from '../../../platform/deepnote/deepnoteServerUtils.node'; import { DeepnoteToolkitMissingError } from '../../../platform/errors/deepnoteKernelErrors'; import { @@ -27,7 +28,7 @@ import { DeepnoteKernelConnectionMetadata, IDeepnoteEnvironmentManager, IDeepnoteKernelAutoSelector, - IDeepnoteNotebookEnvironmentMapper + IDeepnoteProjectEnvironmentMapper } from '../types'; import { CreateDeepnoteEnvironmentOptions, DeepnoteEnvironment } from './deepnoteEnvironment'; import { DeepnoteEnvironmentTreeDataProvider } from './deepnoteEnvironmentTreeDataProvider.node'; @@ -49,8 +50,8 @@ export class DeepnoteEnvironmentsView implements Disposable { @inject(IPythonApiProvider) private readonly pythonApiProvider: IPythonApiProvider, @inject(IDisposableRegistry) disposableRegistry: IDisposableRegistry, @inject(IDeepnoteKernelAutoSelector) private readonly kernelAutoSelector: IDeepnoteKernelAutoSelector, - @inject(IDeepnoteNotebookEnvironmentMapper) - private readonly notebookEnvironmentMapper: IDeepnoteNotebookEnvironmentMapper, + @inject(IDeepnoteProjectEnvironmentMapper) + private readonly projectEnvironmentMapper: IDeepnoteProjectEnvironmentMapper, @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider, @inject(IOutputChannel) @named(STANDARD_OUTPUT_CHANNEL) private readonly outputChannel: IOutputChannel ) { @@ -300,10 +301,13 @@ export class DeepnoteEnvironmentsView implements Disposable { cancellable: true }, async (_progress, token) => { - // Clean up notebook mappings referencing this env - const notebooks = this.notebookEnvironmentMapper.getNotebooksUsingEnvironment(environmentId); - for (const nb of notebooks) { - await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(nb); + // Wait for the mapper to finish migration so legacy entries are resolved + await this.projectEnvironmentMapper.waitForInitialization(); + + // Clean up project mappings referencing this env + const projectIds = this.projectEnvironmentMapper.getProjectsUsingEnvironment(environmentId); + for (const projectId of projectIds) { + await this.projectEnvironmentMapper.removeEnvironmentForProject(projectId); } // Dispose kernels from any open notebooks using this environment @@ -342,26 +346,33 @@ export class DeepnoteEnvironmentsView implements Disposable { // Check if this kernel is using the environment being deleted const connectionMetadata = kernel.kernelConnectionMetadata; - if (connectionMetadata.kind === 'startUsingDeepnoteKernel') { - const deepnoteMetadata = connectionMetadata as DeepnoteKernelConnectionMetadata; - const expectedHandle = createDeepnoteServerConfigHandle(environmentId, notebook.uri); + if (connectionMetadata.kind !== 'startUsingDeepnoteKernel') { + continue; + } - if (deepnoteMetadata.serverProviderHandle.handle === expectedHandle) { - logger.info( - `Disposing kernel for notebook ${getDisplayPath( - notebook.uri - )} as it uses deleted environment ${environmentId}` - ); + const projectId = await resolveProjectIdForNotebook(notebook); + if (!projectId) { + continue; + } - try { - // First, unselect the controller from the notebook UI - this.kernelAutoSelector.clearControllerForEnvironment(notebook, environmentId); + const deepnoteMetadata = connectionMetadata as DeepnoteKernelConnectionMetadata; + const expectedHandle = createDeepnoteServerConfigHandle(environmentId, projectId); - // Then dispose the kernel - await kernel.dispose(); - } catch (error) { - logger.error(`Failed to dispose kernel for ${getDisplayPath(notebook.uri)}`, error); - } + if (deepnoteMetadata.serverProviderHandle.handle === expectedHandle) { + logger.info( + `Disposing kernel for notebook ${getDisplayPath( + notebook.uri + )} as it uses deleted environment ${environmentId}` + ); + + try { + // First, unselect the controller from the notebook UI + await this.kernelAutoSelector.clearControllerForEnvironment(notebook, environmentId); + + // Then dispose the kernel + await kernel.dispose(); + } catch (error) { + logger.error(`Failed to dispose kernel for ${getDisplayPath(notebook.uri)}`, error); } } } @@ -370,11 +381,17 @@ export class DeepnoteEnvironmentsView implements Disposable { public async selectEnvironmentForNotebook({ notebook }: { notebook: NotebookDocument }): Promise { logger.info('Selecting environment for notebook:', notebook); - // Get base file URI (without query/fragment) - const baseFileUri = notebook.uri.with({ query: '', fragment: '' }); + const projectId = await resolveProjectIdForNotebook(notebook); + if (!projectId) { + void window.showWarningMessage(l10n.t('Could not resolve Deepnote project id for this notebook')); + return; + } + + // Wait for the mapper to finish migration so legacy entries are resolved + await this.projectEnvironmentMapper.waitForInitialization(); // Get current environment selection - const currentEnvironmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri); + const currentEnvironmentId = this.projectEnvironmentMapper.getEnvironmentForProject(projectId); const currentEnvironment = currentEnvironmentId ? this.environmentManager.getEnvironment(currentEnvironmentId) : undefined; @@ -471,8 +488,8 @@ export class DeepnoteEnvironmentsView implements Disposable { cancellable: true }, async (progress, token) => { - // Update the notebook-to-environment mapping - await this.notebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, selectedEnvironmentId); + // Update the project-to-environment mapping + await this.projectEnvironmentMapper.setEnvironmentForProject(projectId, selectedEnvironmentId); // Force rebuild the controller with the new environment // This clears cached metadata and creates a fresh controller. diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts index 0407420162..0485b664fd 100644 --- a/src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts +++ b/src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts @@ -3,7 +3,7 @@ import * as sinon from 'sinon'; import { anything, capture, instance, mock, when, verify, deepEqual, resetCalls } from 'ts-mockito'; import { CancellationToken, Disposable, NotebookDocument, ProgressOptions, Uri } from 'vscode'; import { DeepnoteEnvironmentsView } from './deepnoteEnvironmentsView.node'; -import { IDeepnoteEnvironmentManager, IDeepnoteKernelAutoSelector, IDeepnoteNotebookEnvironmentMapper } from '../types'; +import { IDeepnoteEnvironmentManager, IDeepnoteKernelAutoSelector, IDeepnoteProjectEnvironmentMapper } from '../types'; import { IPythonApiProvider } from '../../../platform/api/types'; import { IDisposableRegistry, IOutputChannel } from '../../../platform/common/types'; import { IKernelProvider } from '../../../kernels/types'; @@ -22,7 +22,7 @@ suite('DeepnoteEnvironmentsView', () => { let mockPythonApiProvider: IPythonApiProvider; let mockDisposableRegistry: IDisposableRegistry; let mockKernelAutoSelector: IDeepnoteKernelAutoSelector; - let mockNotebookEnvironmentMapper: IDeepnoteNotebookEnvironmentMapper; + let mockProjectEnvironmentMapper: IDeepnoteProjectEnvironmentMapper; let mockKernelProvider: IKernelProvider; let mockOutputChannel: IOutputChannel; let disposables: Disposable[] = []; @@ -40,10 +40,13 @@ suite('DeepnoteEnvironmentsView', () => { mockPythonApiProvider = mock(); mockDisposableRegistry = mock(); mockKernelAutoSelector = mock(); - mockNotebookEnvironmentMapper = mock(); + mockProjectEnvironmentMapper = mock(); mockKernelProvider = mock(); mockOutputChannel = mock(); + // Default: mapper initialization completes synchronously + when(mockProjectEnvironmentMapper.waitForInitialization()).thenResolve(); + // Mock onDidChangeEnvironments to return a disposable event when(mockConfigManager.onDidChangeEnvironments).thenReturn((_listener: () => void) => { return { @@ -59,7 +62,7 @@ suite('DeepnoteEnvironmentsView', () => { instance(mockPythonApiProvider), instance(mockDisposableRegistry), instance(mockKernelAutoSelector), - instance(mockNotebookEnvironmentMapper), + instance(mockProjectEnvironmentMapper), instance(mockKernelProvider), instance(mockOutputChannel) ); @@ -431,11 +434,11 @@ suite('DeepnoteEnvironmentsView', () => { setup(() => { resetCalls(mockConfigManager); - resetCalls(mockNotebookEnvironmentMapper); + resetCalls(mockProjectEnvironmentMapper); resetCalls(mockedVSCodeNamespaces.window); }); - test('should successfully delete environment with notebooks using it', async () => { + test('should successfully delete environment with projects using it', async () => { // Mock environment exists when(mockConfigManager.getEnvironment(testEnvironmentId)).thenReturn(testEnvironment); @@ -444,16 +447,14 @@ suite('DeepnoteEnvironmentsView', () => { Promise.resolve('Delete') ); - // Mock notebooks using this environment - const notebook1Uri = Uri.file('/workspace/notebook1.deepnote'); - const notebook2Uri = Uri.file('/workspace/notebook2.deepnote'); - when(mockNotebookEnvironmentMapper.getNotebooksUsingEnvironment(testEnvironmentId)).thenReturn([ - notebook1Uri, - notebook2Uri + // Mock projects using this environment + when(mockProjectEnvironmentMapper.getProjectsUsingEnvironment(testEnvironmentId)).thenReturn([ + 'proj-1', + 'proj-2' ]); // Mock removing environment mappings - when(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(anything())).thenResolve(); + when(mockProjectEnvironmentMapper.removeEnvironmentForProject(anything())).thenResolve(); // Mock window.withProgress to execute the callback when(mockedVSCodeNamespaces.window.withProgress(anything(), anything())).thenCall( @@ -489,11 +490,11 @@ suite('DeepnoteEnvironmentsView', () => { // Verify API calls verify(mockConfigManager.getEnvironment(testEnvironmentId)).once(); verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).once(); - verify(mockNotebookEnvironmentMapper.getNotebooksUsingEnvironment(testEnvironmentId)).once(); + verify(mockProjectEnvironmentMapper.getProjectsUsingEnvironment(testEnvironmentId)).once(); - // Verify environment mappings were removed for both notebooks - verify(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(notebook1Uri)).once(); - verify(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(notebook2Uri)).once(); + // Verify environment mappings were removed for both projects + verify(mockProjectEnvironmentMapper.removeEnvironmentForProject('proj-1')).once(); + verify(mockProjectEnvironmentMapper.removeEnvironmentForProject('proj-2')).once(); // Verify environment deletion verify(mockConfigManager.deleteEnvironment(testEnvironmentId, anything())).once(); @@ -511,27 +512,30 @@ suite('DeepnoteEnvironmentsView', () => { Promise.resolve('Delete') ); - // Mock notebooks using this environment - when(mockNotebookEnvironmentMapper.getNotebooksUsingEnvironment(testEnvironmentId)).thenReturn([]); - when(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(anything())).thenResolve(); + // Mock projects using this environment - none + when(mockProjectEnvironmentMapper.getProjectsUsingEnvironment(testEnvironmentId)).thenReturn([]); + when(mockProjectEnvironmentMapper.removeEnvironmentForProject(anything())).thenResolve(); // Mock open notebooks with kernels const openNotebook1 = { uri: Uri.file('/workspace/open-notebook1.deepnote'), notebookType: 'deepnote', - isClosed: false + isClosed: false, + metadata: { deepnoteProjectId: 'proj-open-1' } } as any; const openNotebook2 = { uri: Uri.file('/workspace/open-notebook2.deepnote'), notebookType: 'jupyter-notebook', - isClosed: false + isClosed: false, + metadata: {} } as any; const openNotebook3 = { uri: Uri.file('/workspace/open-notebook3.deepnote'), notebookType: 'deepnote', - isClosed: false + isClosed: false, + metadata: { deepnoteProjectId: 'proj-open-3' } } as any; // Mock workspace.notebookDocuments @@ -541,12 +545,12 @@ suite('DeepnoteEnvironmentsView', () => { openNotebook3 ]); - // Mock kernels + // Mock kernels — handles now use (envId, projectId) const mockKernel1 = { kernelConnectionMetadata: { kind: 'startUsingDeepnoteKernel', serverProviderHandle: { - handle: createDeepnoteServerConfigHandle(testEnvironmentId, openNotebook1.uri) + handle: createDeepnoteServerConfigHandle(testEnvironmentId, 'proj-open-1') } }, dispose: sinon.stub().resolves() @@ -556,7 +560,7 @@ suite('DeepnoteEnvironmentsView', () => { kernelConnectionMetadata: { kind: 'startUsingDeepnoteKernel', serverProviderHandle: { - handle: createDeepnoteServerConfigHandle('different-env-id', openNotebook3.uri) + handle: createDeepnoteServerConfigHandle('different-env-id', 'proj-open-3') } }, dispose: sinon.stub().resolves() @@ -567,6 +571,9 @@ suite('DeepnoteEnvironmentsView', () => { when(mockKernelProvider.get(openNotebook2)).thenReturn(undefined); // No kernel for jupyter notebook when(mockKernelProvider.get(openNotebook3)).thenReturn(mockKernel3 as any); + // Mock clearControllerForEnvironment (now async) + when(mockKernelAutoSelector.clearControllerForEnvironment(anything(), anything())).thenResolve(); + // Mock window.withProgress when(mockedVSCodeNamespaces.window.withProgress(anything(), anything())).thenCall( (_options: ProgressOptions, callback: Function) => { @@ -612,16 +619,14 @@ suite('DeepnoteEnvironmentsView', () => { Promise.resolve('Delete') ); - // Mock notebooks using this environment - const notebook1Uri = Uri.file('/workspace/notebook1.deepnote'); - const notebook2Uri = Uri.file('/workspace/notebook2.deepnote'); - when(mockNotebookEnvironmentMapper.getNotebooksUsingEnvironment(testExternalEnvironmentId)).thenReturn([ - notebook1Uri, - notebook2Uri + // Mock projects using this environment + when(mockProjectEnvironmentMapper.getProjectsUsingEnvironment(testExternalEnvironmentId)).thenReturn([ + 'proj-1', + 'proj-2' ]); // Mock removing environment mappings - when(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(anything())).thenResolve(); + when(mockProjectEnvironmentMapper.removeEnvironmentForProject(anything())).thenResolve(); // Mock open notebooks when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); @@ -660,11 +665,11 @@ suite('DeepnoteEnvironmentsView', () => { // Verify API calls - same as for managed venv verify(mockConfigManager.getEnvironment(testExternalEnvironmentId)).once(); verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).once(); - verify(mockNotebookEnvironmentMapper.getNotebooksUsingEnvironment(testExternalEnvironmentId)).once(); + verify(mockProjectEnvironmentMapper.getProjectsUsingEnvironment(testExternalEnvironmentId)).once(); - // Verify environment mappings were removed for both notebooks - verify(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(notebook1Uri)).once(); - verify(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(notebook2Uri)).once(); + // Verify environment mappings were removed for both projects + verify(mockProjectEnvironmentMapper.removeEnvironmentForProject('proj-1')).once(); + verify(mockProjectEnvironmentMapper.removeEnvironmentForProject('proj-2')).once(); // Verify environment deletion - the manager is responsible for checking managedVenv verify(mockConfigManager.deleteEnvironment(testExternalEnvironmentId, anything())).once(); @@ -682,15 +687,16 @@ suite('DeepnoteEnvironmentsView', () => { Promise.resolve('Delete') ); - // Mock notebooks using this environment - when(mockNotebookEnvironmentMapper.getNotebooksUsingEnvironment(testExternalEnvironmentId)).thenReturn([]); - when(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(anything())).thenResolve(); + // Mock projects using this environment + when(mockProjectEnvironmentMapper.getProjectsUsingEnvironment(testExternalEnvironmentId)).thenReturn([]); + when(mockProjectEnvironmentMapper.removeEnvironmentForProject(anything())).thenResolve(); // Mock open notebooks with kernels const openNotebook1 = { uri: Uri.file('/workspace/open-notebook1.deepnote'), notebookType: 'deepnote', - isClosed: false + isClosed: false, + metadata: { deepnoteProjectId: 'proj-ext-1' } } as any; when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([openNotebook1]); @@ -700,7 +706,7 @@ suite('DeepnoteEnvironmentsView', () => { kernelConnectionMetadata: { kind: 'startUsingDeepnoteKernel', serverProviderHandle: { - handle: createDeepnoteServerConfigHandle(testExternalEnvironmentId, openNotebook1.uri) + handle: createDeepnoteServerConfigHandle(testExternalEnvironmentId, 'proj-ext-1') } }, dispose: sinon.stub().resolves() @@ -708,6 +714,9 @@ suite('DeepnoteEnvironmentsView', () => { when(mockKernelProvider.get(openNotebook1)).thenReturn(mockKernel1 as any); + // Mock clearControllerForEnvironment (now async) + when(mockKernelAutoSelector.clearControllerForEnvironment(anything(), anything())).thenResolve(); + // Mock window.withProgress when(mockedVSCodeNamespaces.window.withProgress(anything(), anything())).thenCall( (_options: ProgressOptions, callback: Function) => { @@ -799,21 +808,27 @@ suite('DeepnoteEnvironmentsView', () => { lastUsedAt: new Date() }; + const testProjectId = 'test-project-id'; + setup(() => { resetCalls(mockConfigManager); - resetCalls(mockNotebookEnvironmentMapper); + resetCalls(mockProjectEnvironmentMapper); resetCalls(mockKernelAutoSelector); resetCalls(mockKernelProvider); resetCalls(mockedVSCodeNamespaces.window); + + // Mapper init + when(mockProjectEnvironmentMapper.waitForInitialization()).thenResolve(); }); test('should successfully switch to a different environment', async () => { - // Mock active notebook + // Mock active notebook with projectId in metadata so resolver succeeds const notebookUri = Uri.file('/workspace/notebook.deepnote'); const mockNotebook = { uri: notebookUri, notebookType: 'deepnote', - cellCount: 5 + cellCount: 5, + metadata: { deepnoteProjectId: testProjectId } }; const mockNotebookEditor = { notebook: mockNotebook @@ -821,9 +836,8 @@ suite('DeepnoteEnvironmentsView', () => { when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(mockNotebookEditor as any); - // Mock current environment mapping - const baseFileUri = notebookUri.with({ query: '', fragment: '' }); - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri)).thenReturn( + // Mock current environment mapping (by projectId) + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn( currentEnvironment.id ); when(mockConfigManager.getEnvironment(currentEnvironment.id)).thenReturn(currentEnvironment); @@ -832,7 +846,6 @@ suite('DeepnoteEnvironmentsView', () => { when(mockConfigManager.listEnvironments()).thenReturn([currentEnvironment, newEnvironment]); // Mock environment status - when(mockConfigManager.getEnvironment(currentEnvironment.id)).thenReturn(currentEnvironment); when(mockConfigManager.getEnvironment(newEnvironment.id)).thenReturn(newEnvironment); // Mock user selecting the new environment @@ -857,8 +870,8 @@ suite('DeepnoteEnvironmentsView', () => { } ); - // Mock environment mapping update - when(mockNotebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, newEnvironment.id)).thenResolve(); + // Mock environment mapping update (by projectId) + when(mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, newEnvironment.id)).thenResolve(); // Mock controller rebuild when(mockKernelAutoSelector.rebuildController(mockNotebook as any, anything(), anything())).thenResolve(); @@ -867,20 +880,18 @@ suite('DeepnoteEnvironmentsView', () => { when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenResolve(undefined); // Execute the command - await view.selectEnvironmentForNotebook({ notebook: mockNotebook as NotebookDocument }); + await view.selectEnvironmentForNotebook({ notebook: mockNotebook as unknown as NotebookDocument }); // Verify API calls - verify(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri)).once(); - verify(mockConfigManager.getEnvironment(currentEnvironment.id)).once(); + verify(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).once(); verify(mockConfigManager.listEnvironments()).once(); - verify(mockConfigManager.getEnvironment(currentEnvironment.id)).once(); verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).once(); verify(mockKernelProvider.get(mockNotebook as any)).once(); verify(mockKernelProvider.getKernelExecution(mockKernel as any)).once(); // Verify environment switch verify(mockedVSCodeNamespaces.window.withProgress(anything(), anything())).once(); - verify(mockNotebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, newEnvironment.id)).once(); + verify(mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, newEnvironment.id)).once(); verify(mockKernelAutoSelector.rebuildController(mockNotebook as any, anything(), anything())).once(); // Verify success message was shown @@ -893,7 +904,8 @@ suite('DeepnoteEnvironmentsView', () => { const mockNotebook = { uri: notebookUri, notebookType: 'deepnote', - cellCount: 5 + cellCount: 5, + metadata: { deepnoteProjectId: testProjectId } }; const mockNotebookEditor = { notebook: mockNotebook @@ -902,8 +914,7 @@ suite('DeepnoteEnvironmentsView', () => { when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(mockNotebookEditor as any); // Mock current environment mapping (managed) - const baseFileUri = notebookUri.with({ query: '', fragment: '' }); - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri)).thenReturn( + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn( currentEnvironment.id ); when(mockConfigManager.getEnvironment(currentEnvironment.id)).thenReturn(currentEnvironment); @@ -942,7 +953,7 @@ suite('DeepnoteEnvironmentsView', () => { // Mock environment mapping update when( - mockNotebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, newExternalEnvironment.id) + mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, newExternalEnvironment.id) ).thenResolve(); // Mock controller rebuild @@ -952,12 +963,12 @@ suite('DeepnoteEnvironmentsView', () => { when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenResolve(undefined); // Execute the command - await view.selectEnvironmentForNotebook({ notebook: mockNotebook as NotebookDocument }); + await view.selectEnvironmentForNotebook({ notebook: mockNotebook as unknown as NotebookDocument }); // Verify environment switch to external environment verify(mockedVSCodeNamespaces.window.withProgress(anything(), anything())).once(); verify( - mockNotebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, newExternalEnvironment.id) + mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, newExternalEnvironment.id) ).once(); verify(mockKernelAutoSelector.rebuildController(mockNotebook as any, anything(), anything())).once(); @@ -971,7 +982,8 @@ suite('DeepnoteEnvironmentsView', () => { const mockNotebook = { uri: notebookUri, notebookType: 'deepnote', - cellCount: 5 + cellCount: 5, + metadata: { deepnoteProjectId: testProjectId } }; const mockNotebookEditor = { notebook: mockNotebook @@ -980,8 +992,7 @@ suite('DeepnoteEnvironmentsView', () => { when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(mockNotebookEditor as any); // Mock current environment mapping (external) - const baseFileUri = notebookUri.with({ query: '', fragment: '' }); - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri)).thenReturn( + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn( currentExternalEnvironment.id ); when(mockConfigManager.getEnvironment(currentExternalEnvironment.id)).thenReturn( @@ -1016,7 +1027,7 @@ suite('DeepnoteEnvironmentsView', () => { ); // Mock environment mapping update - when(mockNotebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, newEnvironment.id)).thenResolve(); + when(mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, newEnvironment.id)).thenResolve(); // Mock controller rebuild when(mockKernelAutoSelector.rebuildController(mockNotebook as any, anything(), anything())).thenResolve(); @@ -1025,11 +1036,11 @@ suite('DeepnoteEnvironmentsView', () => { when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenResolve(undefined); // Execute the command - await view.selectEnvironmentForNotebook({ notebook: mockNotebook as NotebookDocument }); + await view.selectEnvironmentForNotebook({ notebook: mockNotebook as unknown as NotebookDocument }); // Verify environment switch from external to managed verify(mockedVSCodeNamespaces.window.withProgress(anything(), anything())).once(); - verify(mockNotebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, newEnvironment.id)).once(); + verify(mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, newEnvironment.id)).once(); verify(mockKernelAutoSelector.rebuildController(mockNotebook as any, anything(), anything())).once(); // Verify success message was shown diff --git a/src/kernels/deepnote/environments/deepnoteExtensionSidecarWriter.node.ts b/src/kernels/deepnote/environments/deepnoteExtensionSidecarWriter.node.ts index 20a39162c0..1781a21ee7 100644 --- a/src/kernels/deepnote/environments/deepnoteExtensionSidecarWriter.node.ts +++ b/src/kernels/deepnote/environments/deepnoteExtensionSidecarWriter.node.ts @@ -1,11 +1,10 @@ import { inject, injectable } from 'inversify'; -import * as YAML from 'yaml'; -import { env, NotebookDocument, Uri, workspace } from 'vscode'; +import { env, Uri, workspace } from 'vscode'; import { IExtensionSyncActivationService } from '../../../platform/activation/types'; import { IDisposableRegistry } from '../../../platform/common/types'; import { logger } from '../../../platform/logging'; -import { IDeepnoteEnvironmentManager, IDeepnoteNotebookEnvironmentMapper } from '../types'; +import { IDeepnoteEnvironmentManager, IDeepnoteProjectEnvironmentMapper } from '../types'; const SIDECAR_FILENAME = 'deepnote.json'; @@ -45,13 +44,11 @@ interface SidecarFile { */ @injectable() export class DeepnoteExtensionSidecarWriter implements IExtensionSyncActivationService { - /** Reverse map: notebookUri.fsPath → projectId (populated from sidecar + set calls). */ - private readonly fsPathToProjectId = new Map(); /** Serializes sidecar writes to avoid read-modify-write races. */ private writeQueue: Promise = Promise.resolve(); constructor( - @inject(IDeepnoteNotebookEnvironmentMapper) private readonly mapper: IDeepnoteNotebookEnvironmentMapper, + @inject(IDeepnoteProjectEnvironmentMapper) private readonly mapper: IDeepnoteProjectEnvironmentMapper, @inject(IDeepnoteEnvironmentManager) private readonly environmentManager: IDeepnoteEnvironmentManager, @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry ) {} @@ -60,8 +57,7 @@ export class DeepnoteExtensionSidecarWriter implements IExtensionSyncActivationS this.disposables.push( this.mapper.onDidSetEnvironment((e) => this.handleSetEnvironment(e)), this.mapper.onDidRemoveEnvironment((e) => this.handleRemoveEnvironment(e)), - this.environmentManager.onDidChangeEnvironments(() => this.handleEnvironmentsChanged()), - workspace.onDidOpenNotebookDocument((doc) => this.handleNotebookOpened(doc)) + this.environmentManager.onDidChangeEnvironments(() => this.handleEnvironmentsChanged()) ); // Sync existing mappings so the sidecar is up-to-date for existing users @@ -106,64 +102,8 @@ export class DeepnoteExtensionSidecarWriter implements IExtensionSyncActivationS } } - /** - * When a notebook is opened after activation, check if it already has - * a mapping and add it to the sidecar. - */ - private async handleNotebookOpened(doc: NotebookDocument): Promise { - if (doc.notebookType !== 'deepnote') { - return; - } - + private async handleRemoveEnvironment({ projectId }: { projectId: string }): Promise { try { - const notebookUri = doc.uri.with({ query: '', fragment: '' }); - const projectId = doc.metadata?.deepnoteProjectId as string | undefined; - if (!projectId) { - return; - } - - const environmentId = this.mapper.getEnvironmentForNotebook(notebookUri); - if (!environmentId) { - return; - } - - const environment = this.environmentManager.getEnvironment(environmentId); - if (!environment) { - return; - } - - this.fsPathToProjectId.set(notebookUri.fsPath, projectId); - - await this.enqueueWrite(async (sidecar) => { - const existing = sidecar.mappings[projectId]; - if ( - existing?.environmentId === environmentId && - existing?.venvPath === environment.venvPath.fsPath && - existing?.pythonInterpreter === environment.pythonInterpreter.uri.fsPath - ) { - return false; - } - sidecar.mappings[projectId] = { - environmentId, - venvPath: environment.venvPath.fsPath, - pythonInterpreter: environment.pythonInterpreter.uri.fsPath - }; - return true; - }); - } catch (error) { - logger.warn('[SidecarWriter] Failed to handle notebook opened', error); - } - } - - private async handleRemoveEnvironment({ notebookUri }: { notebookUri: Uri }): Promise { - try { - const projectId = this.fsPathToProjectId.get(notebookUri.fsPath) ?? this.resolveProjectId(notebookUri); - if (!projectId) { - return; - } - - this.fsPathToProjectId.delete(notebookUri.fsPath); - await this.enqueueWrite(async (sidecar) => { if (!(projectId in sidecar.mappings)) { return false; @@ -177,25 +117,18 @@ export class DeepnoteExtensionSidecarWriter implements IExtensionSyncActivationS } private async handleSetEnvironment({ - notebookUri, + projectId, environmentId }: { - notebookUri: Uri; + projectId: string; environmentId: string; }): Promise { try { - const projectId = this.resolveProjectId(notebookUri); - if (!projectId) { - return; - } - const environment = this.environmentManager.getEnvironment(environmentId); if (!environment) { return; } - this.fsPathToProjectId.set(notebookUri.fsPath, projectId); - await this.enqueueWrite(async (sidecar) => { sidecar.mappings[projectId] = { environmentId, @@ -210,49 +143,42 @@ export class DeepnoteExtensionSidecarWriter implements IExtensionSyncActivationS } /** - * On activation, iterate all persisted mapper entries (not just open - * notebooks) and write their mappings to the sidecar. For each entry - * we read the `.deepnote` file to extract the project ID. + * On activation, iterate all persisted mapper entries and write their + * mappings to the sidecar. Mappings are already keyed by project id from + * the new mapper, so no per-file YAML reads are required here. */ private async syncExistingMappings(): Promise { try { await this.environmentManager.waitForInitialization(); + await this.mapper.waitForInitialization(); const allMappings = this.mapper.getAllMappings(); if (allMappings.size === 0) { return; } - // Collect all project IDs outside the write queue to avoid holding the lock during I/O. const entries: Array<{ - fsPath: string; projectId: string; environmentId: string; venvPath: string; pythonInterpreter: string; }> = []; - for (const [fsPath, environmentId] of allMappings) { + + for (const [projectId, environmentId] of allMappings) { try { const environment = this.environmentManager.getEnvironment(environmentId); if (!environment) { continue; } - const projectId = await this.readProjectIdFromFile(Uri.file(fsPath)); - if (!projectId) { - continue; - } - - this.fsPathToProjectId.set(fsPath, projectId); entries.push({ - fsPath, projectId, environmentId, venvPath: environment.venvPath.fsPath, pythonInterpreter: environment.pythonInterpreter.uri.fsPath }); } catch (entryError) { - logger.warn(`[SidecarWriter] Failed to process mapping for ${fsPath}`, entryError); + logger.warn(`[SidecarWriter] Failed to process mapping for project ${projectId}`, entryError); } } @@ -304,16 +230,6 @@ export class DeepnoteExtensionSidecarWriter implements IExtensionSyncActivationS return Uri.joinPath(folder.uri, getEditorSettingsFolder(), SIDECAR_FILENAME); } - private async readProjectIdFromFile(fileUri: Uri): Promise { - try { - const raw = await workspace.fs.readFile(fileUri); - const parsed = YAML.parse(Buffer.from(raw).toString('utf-8')) as { project?: { id?: string } } | undefined; - return parsed?.project?.id; - } catch { - return undefined; - } - } - private async readSidecar(): Promise { const uri = this.getSidecarUri(); if (!uri) { @@ -333,14 +249,6 @@ export class DeepnoteExtensionSidecarWriter implements IExtensionSyncActivationS return { mappings: {} }; } - private resolveProjectId(notebookUri: Uri): string | undefined { - const doc = workspace.notebookDocuments.find( - (d) => - d.notebookType === 'deepnote' && d.uri.with({ query: '', fragment: '' }).fsPath === notebookUri.fsPath - ); - return doc?.metadata?.deepnoteProjectId as string | undefined; - } - private async writeSidecar(sidecar: SidecarFile): Promise { const uri = this.getSidecarUri(); if (!uri) { diff --git a/src/kernels/deepnote/environments/deepnoteExtensionSidecarWriter.unit.test.ts b/src/kernels/deepnote/environments/deepnoteExtensionSidecarWriter.unit.test.ts index 5db130e25d..9299b5e9c1 100644 --- a/src/kernels/deepnote/environments/deepnoteExtensionSidecarWriter.unit.test.ts +++ b/src/kernels/deepnote/environments/deepnoteExtensionSidecarWriter.unit.test.ts @@ -5,7 +5,7 @@ import { EventEmitter, NotebookDocument, Uri, WorkspaceFolder } from 'vscode'; import type { IDisposableRegistry } from '../../../platform/common/types'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; -import { IDeepnoteEnvironmentManager, IDeepnoteNotebookEnvironmentMapper } from '../types'; +import { IDeepnoteEnvironmentManager, IDeepnoteProjectEnvironmentMapper } from '../types'; import { DeepnoteEnvironment } from './deepnoteEnvironment'; import { DeepnoteExtensionSidecarWriter } from './deepnoteExtensionSidecarWriter.node'; @@ -38,36 +38,14 @@ function makeEnvironment(overrides: Partial & { id: string }; } -function makeDeepnoteYaml(projectId: string): string { - return `version: '1.0'\nproject:\n id: ${projectId}\n name: Test\n notebooks: []\n`; -} - -function createMockNotebook(opts: { uri: Uri; projectId?: string; notebookType?: string }): NotebookDocument { - return { - uri: opts.uri, - notebookType: opts.notebookType ?? 'deepnote', - metadata: { deepnoteProjectId: opts.projectId ?? 'project-1' }, - isDirty: false, - isUntitled: false, - isClosed: false, - version: 1, - cellCount: 0, - cellAt: () => { - throw new Error('Not implemented'); - }, - getCells: () => [], - save: async () => true - } as unknown as NotebookDocument; -} - suite('DeepnoteExtensionSidecarWriter', () => { let writer: DeepnoteExtensionSidecarWriter; let disposables: IDisposableRegistry; - let mockMapper: IDeepnoteNotebookEnvironmentMapper; + let mockMapper: IDeepnoteProjectEnvironmentMapper; let mockEnvironmentManager: IDeepnoteEnvironmentManager; - let onDidSetEnvironment: EventEmitter<{ notebookUri: Uri; environmentId: string }>; - let onDidRemoveEnvironment: EventEmitter<{ notebookUri: Uri }>; + let onDidSetEnvironment: EventEmitter<{ projectId: string; environmentId: string }>; + let onDidRemoveEnvironment: EventEmitter<{ projectId: string }>; let onDidChangeEnvironments: EventEmitter; let onDidOpenNotebookDocument: EventEmitter; @@ -75,8 +53,6 @@ suite('DeepnoteExtensionSidecarWriter', () => { let writeFileCallCount: number; let createDirectoryUris: Uri[]; let readFileContent: string; - /** Per-file read responses keyed by fsPath — used for .deepnote YAML files. */ - let fileContents: Map; const workspaceUri = Uri.file('/workspace'); @@ -86,13 +62,12 @@ suite('DeepnoteExtensionSidecarWriter', () => { writeFileCallCount = 0; createDirectoryUris = []; readFileContent = ''; - fileContents = new Map(); disposables = []; // Set up event emitters - onDidSetEnvironment = new EventEmitter<{ notebookUri: Uri; environmentId: string }>(); - onDidRemoveEnvironment = new EventEmitter<{ notebookUri: Uri }>(); + onDidSetEnvironment = new EventEmitter<{ projectId: string; environmentId: string }>(); + onDidRemoveEnvironment = new EventEmitter<{ projectId: string }>(); onDidChangeEnvironments = new EventEmitter(); onDidOpenNotebookDocument = new EventEmitter(); disposables.push( @@ -103,10 +78,11 @@ suite('DeepnoteExtensionSidecarWriter', () => { ); // Set up mapper mock - mockMapper = mock(); + mockMapper = mock(); when(mockMapper.onDidSetEnvironment).thenReturn(onDidSetEnvironment.event); when(mockMapper.onDidRemoveEnvironment).thenReturn(onDidRemoveEnvironment.event); when(mockMapper.getAllMappings()).thenReturn(new Map()); + when(mockMapper.waitForInitialization()).thenResolve(); // Set up environment manager mock mockEnvironmentManager = mock(); @@ -137,13 +113,7 @@ suite('DeepnoteExtensionSidecarWriter', () => { function setupMockFs() { const mockFs = mock(); - when(mockFs.readFile(anything())).thenCall((uri: Uri) => { - // Check per-file map first (for .deepnote YAML files) - const perFile = fileContents.get(uri.fsPath); - if (perFile !== undefined) { - return Promise.resolve(Buffer.from(perFile, 'utf-8')); - } - // Fall back to global sidecar content + when(mockFs.readFile(anything())).thenCall(() => { if (!readFileContent) { return Promise.reject(new Error('File not found')); } @@ -169,10 +139,6 @@ suite('DeepnoteExtensionSidecarWriter', () => { } test('set mapping writes sidecar with correct projectId, environmentId, and venvPath', async () => { - const notebookUri = Uri.file('/workspace/project.deepnote'); - const notebook = createMockNotebook({ uri: notebookUri, projectId: 'proj-abc' }); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - const env = makeEnvironment({ id: 'env-1', venvPath: Uri.file('/home/user/.venvs/my-env') @@ -181,7 +147,7 @@ suite('DeepnoteExtensionSidecarWriter', () => { writer.activate(); - onDidSetEnvironment.fire({ notebookUri, environmentId: 'env-1' }); + onDidSetEnvironment.fire({ projectId: 'proj-abc', environmentId: 'env-1' }); await waitFor(() => writeFileCallCount > 0); @@ -194,23 +160,19 @@ suite('DeepnoteExtensionSidecarWriter', () => { }); test('remove mapping removes entry from sidecar', async () => { - const notebookUri = Uri.file('/workspace/project.deepnote'); - const notebook = createMockNotebook({ uri: notebookUri, projectId: 'proj-abc' }); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - const env = makeEnvironment({ id: 'env-1' }); when(mockEnvironmentManager.getEnvironment('env-1')).thenReturn(env); writer.activate(); // First set, then remove - onDidSetEnvironment.fire({ notebookUri, environmentId: 'env-1' }); + onDidSetEnvironment.fire({ projectId: 'proj-abc', environmentId: 'env-1' }); await waitFor(() => writeFileCallCount >= 1); // Set the sidecar content so the next read picks it up readFileContent = writtenContent!; - onDidRemoveEnvironment.fire({ notebookUri }); + onDidRemoveEnvironment.fire({ projectId: 'proj-abc' }); await waitFor(() => writeFileCallCount >= 2); const sidecar = parseSidecar(); @@ -219,12 +181,6 @@ suite('DeepnoteExtensionSidecarWriter', () => { }); test('multiple projects accumulate entries in a single sidecar', async () => { - const uri1 = Uri.file('/workspace/project1.deepnote'); - const uri2 = Uri.file('/workspace/project2.deepnote'); - const nb1 = createMockNotebook({ uri: uri1, projectId: 'proj-1' }); - const nb2 = createMockNotebook({ uri: uri2, projectId: 'proj-2' }); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([nb1, nb2]); - const env1 = makeEnvironment({ id: 'env-1', venvPath: Uri.file('/venvs/env1') }); const env2 = makeEnvironment({ id: 'env-2', venvPath: Uri.file('/venvs/env2') }); when(mockEnvironmentManager.getEnvironment('env-1')).thenReturn(env1); @@ -232,11 +188,11 @@ suite('DeepnoteExtensionSidecarWriter', () => { writer.activate(); - onDidSetEnvironment.fire({ notebookUri: uri1, environmentId: 'env-1' }); + onDidSetEnvironment.fire({ projectId: 'proj-1', environmentId: 'env-1' }); await waitFor(() => writeFileCallCount >= 1); readFileContent = writtenContent!; - onDidSetEnvironment.fire({ notebookUri: uri2, environmentId: 'env-2' }); + onDidSetEnvironment.fire({ projectId: 'proj-2', environmentId: 'env-2' }); await waitFor(() => writeFileCallCount >= 2); const sidecar = parseSidecar(); @@ -247,10 +203,6 @@ suite('DeepnoteExtensionSidecarWriter', () => { }); test('error reading sidecar does not throw', async () => { - const notebookUri = Uri.file('/workspace/project.deepnote'); - const notebook = createMockNotebook({ uri: notebookUri, projectId: 'proj-abc' }); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - const env = makeEnvironment({ id: 'env-1' }); when(mockEnvironmentManager.getEnvironment('env-1')).thenReturn(env); @@ -258,7 +210,7 @@ suite('DeepnoteExtensionSidecarWriter', () => { writer.activate(); // Should not throw even though readFile fails - onDidSetEnvironment.fire({ notebookUri, environmentId: 'env-1' }); + onDidSetEnvironment.fire({ projectId: 'proj-abc', environmentId: 'env-1' }); await waitFor(() => writeFileCallCount >= 1); // Still writes successfully with a fresh sidecar @@ -267,10 +219,6 @@ suite('DeepnoteExtensionSidecarWriter', () => { }); test('error writing sidecar does not throw', async () => { - const notebookUri = Uri.file('/workspace/project.deepnote'); - const notebook = createMockNotebook({ uri: notebookUri, projectId: 'proj-abc' }); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - const env = makeEnvironment({ id: 'env-1' }); when(mockEnvironmentManager.getEnvironment('env-1')).thenReturn(env); @@ -283,7 +231,7 @@ suite('DeepnoteExtensionSidecarWriter', () => { writer.activate(); // Should not throw - onDidSetEnvironment.fire({ notebookUri, environmentId: 'env-1' }); + onDidSetEnvironment.fire({ projectId: 'proj-abc', environmentId: 'env-1' }); // Give time for the async operation to complete await new Promise((resolve) => setTimeout(resolve, 200)); @@ -292,17 +240,13 @@ suite('DeepnoteExtensionSidecarWriter', () => { }); test('environment changed refreshes sidecar with updated venvPath', async () => { - const notebookUri = Uri.file('/workspace/project.deepnote'); - const notebook = createMockNotebook({ uri: notebookUri, projectId: 'proj-abc' }); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - const env = makeEnvironment({ id: 'env-1', venvPath: Uri.file('/venvs/old-path') }); when(mockEnvironmentManager.getEnvironment('env-1')).thenReturn(env); writer.activate(); // Set initial mapping - onDidSetEnvironment.fire({ notebookUri, environmentId: 'env-1' }); + onDidSetEnvironment.fire({ projectId: 'proj-abc', environmentId: 'env-1' }); await waitFor(() => writeFileCallCount >= 1); readFileContent = writtenContent!; @@ -318,16 +262,12 @@ suite('DeepnoteExtensionSidecarWriter', () => { }); test('environment deleted removes entry on environments changed', async () => { - const notebookUri = Uri.file('/workspace/project.deepnote'); - const notebook = createMockNotebook({ uri: notebookUri, projectId: 'proj-abc' }); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - const env = makeEnvironment({ id: 'env-1' }); when(mockEnvironmentManager.getEnvironment('env-1')).thenReturn(env); writer.activate(); - onDidSetEnvironment.fire({ notebookUri, environmentId: 'env-1' }); + onDidSetEnvironment.fire({ projectId: 'proj-abc', environmentId: 'env-1' }); await waitFor(() => writeFileCallCount >= 1); readFileContent = writtenContent!; @@ -344,28 +284,21 @@ suite('DeepnoteExtensionSidecarWriter', () => { test('no-op when no workspace folder is open', async () => { when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn(undefined); - const notebookUri = Uri.file('/workspace/project.deepnote'); - const notebook = createMockNotebook({ uri: notebookUri, projectId: 'proj-abc' }); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - const env = makeEnvironment({ id: 'env-1' }); when(mockEnvironmentManager.getEnvironment('env-1')).thenReturn(env); writer.activate(); - onDidSetEnvironment.fire({ notebookUri, environmentId: 'env-1' }); + onDidSetEnvironment.fire({ projectId: 'proj-abc', environmentId: 'env-1' }); await new Promise((resolve) => setTimeout(resolve, 200)); assert.strictEqual(writeFileCallCount, 0, 'Should not write when no workspace folder'); }); - test('activation syncs existing mappings to sidecar by reading .deepnote files', async () => { - const fsPath = '/workspace/project.deepnote'; - - // Mapper has a persisted entry — notebook is NOT open - when(mockMapper.getAllMappings()).thenReturn(new Map([[fsPath, 'env-existing']])); - fileContents.set(fsPath, makeDeepnoteYaml('proj-existing')); + test('activation syncs existing project-keyed mappings to sidecar', async () => { + // Mapper already has project-keyed entries (no need to read any .deepnote files) + when(mockMapper.getAllMappings()).thenReturn(new Map([['proj-existing', 'env-existing']])); const env = makeEnvironment({ id: 'env-existing', venvPath: Uri.file('/venvs/existing') }); when(mockEnvironmentManager.getEnvironment('env-existing')).thenReturn(env); @@ -382,18 +315,13 @@ suite('DeepnoteExtensionSidecarWriter', () => { }); }); - test('activation syncs multiple projects including closed notebooks', async () => { - const path1 = '/workspace/proj1.deepnote'; - const path2 = '/workspace/proj2.deepnote'; - + test('activation syncs multiple project-keyed mappings', async () => { when(mockMapper.getAllMappings()).thenReturn( new Map([ - [path1, 'env-1'], - [path2, 'env-2'] + ['proj-1', 'env-1'], + ['proj-2', 'env-2'] ]) ); - fileContents.set(path1, makeDeepnoteYaml('proj-1')); - fileContents.set(path2, makeDeepnoteYaml('proj-2')); const env1 = makeEnvironment({ id: 'env-1', venvPath: Uri.file('/venvs/env1') }); const env2 = makeEnvironment({ id: 'env-2', venvPath: Uri.file('/venvs/env2') }); @@ -411,23 +339,17 @@ suite('DeepnoteExtensionSidecarWriter', () => { }); }); - test('activation skips entries whose .deepnote file cannot be read', async () => { - const goodPath = '/workspace/good.deepnote'; - const badPath = '/workspace/missing.deepnote'; - + test('activation skips entries whose environment does not exist', async () => { when(mockMapper.getAllMappings()).thenReturn( new Map([ - [goodPath, 'env-1'], - [badPath, 'env-2'] + ['proj-good', 'env-1'], + ['proj-missing', 'env-missing'] ]) ); - fileContents.set(goodPath, makeDeepnoteYaml('proj-good')); - // badPath not in fileContents → readFile will reject const env1 = makeEnvironment({ id: 'env-1', venvPath: Uri.file('/venvs/env1') }); - const env2 = makeEnvironment({ id: 'env-2', venvPath: Uri.file('/venvs/env2') }); when(mockEnvironmentManager.getEnvironment('env-1')).thenReturn(env1); - when(mockEnvironmentManager.getEnvironment('env-2')).thenReturn(env2); + // env-missing returns undefined from the environment manager writer.activate(); @@ -439,77 +361,13 @@ suite('DeepnoteExtensionSidecarWriter', () => { assert.isUndefined(sidecar.mappings['proj-missing']); }); - test('opening a notebook with existing mapping writes to sidecar', async () => { - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); - - writer.activate(); - - // Wait for initial sync (no-op since no notebooks open) - await new Promise((resolve) => setTimeout(resolve, 100)); - - const notebookUri = Uri.file('/workspace/project.deepnote'); - const notebook = createMockNotebook({ uri: notebookUri, projectId: 'proj-opened' }); - - const env = makeEnvironment({ id: 'env-1', venvPath: Uri.file('/venvs/env1') }); - when(mockEnvironmentManager.getEnvironment('env-1')).thenReturn(env); - when(mockMapper.getEnvironmentForNotebook(anything())).thenReturn('env-1'); - - onDidOpenNotebookDocument.fire(notebook); - - await waitFor(() => writeFileCallCount >= 1); - - const sidecar = parseSidecar(); - assert.deepStrictEqual(sidecar.mappings['proj-opened'], { - environmentId: 'env-1', - venvPath: '/venvs/env1', - pythonInterpreter: '/usr/bin/python3' - }); - }); - - test('no-op when notebook has no projectId in metadata', async () => { - const notebookUri = Uri.file('/workspace/project.deepnote'); - // Notebook without projectId - const notebook = { - uri: notebookUri, - notebookType: 'deepnote', - metadata: {}, - isDirty: false, - isUntitled: false, - isClosed: false, - version: 1, - cellCount: 0, - cellAt: () => { - throw new Error('Not implemented'); - }, - getCells: () => [], - save: async () => true - } as unknown as NotebookDocument; - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - - const env = makeEnvironment({ id: 'env-1' }); - when(mockEnvironmentManager.getEnvironment('env-1')).thenReturn(env); - - writer.activate(); - - onDidSetEnvironment.fire({ notebookUri, environmentId: 'env-1' }); - - await new Promise((resolve) => setTimeout(resolve, 200)); - - assert.strictEqual(writeFileCallCount, 0, 'Should not write when no projectId'); - }); - test('creates the editor settings folder before writing', async () => { - const notebookUri = Uri.file('/workspace/project.deepnote'); - const notebook = createMockNotebook({ uri: notebookUri, projectId: 'proj-abc' }); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - const env = makeEnvironment({ id: 'env-1' }); when(mockEnvironmentManager.getEnvironment('env-1')).thenReturn(env); writer.activate(); - onDidSetEnvironment.fire({ notebookUri, environmentId: 'env-1' }); + onDidSetEnvironment.fire({ projectId: 'proj-abc', environmentId: 'env-1' }); await waitFor(() => writeFileCallCount > 0); diff --git a/src/kernels/deepnote/environments/deepnoteNotebookEnvironmentMapper.node.ts b/src/kernels/deepnote/environments/deepnoteNotebookEnvironmentMapper.node.ts deleted file mode 100644 index 61070f8217..0000000000 --- a/src/kernels/deepnote/environments/deepnoteNotebookEnvironmentMapper.node.ts +++ /dev/null @@ -1,110 +0,0 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - -import { injectable, inject } from 'inversify'; -import { EventEmitter, Uri, Memento } from 'vscode'; - -import { IDisposableRegistry, IExtensionContext } from '../../../platform/common/types'; -import { logger } from '../../../platform/logging'; - -/** - * Manages the mapping between notebooks and their selected environments - * Stores selections in workspace state for persistence across sessions - */ -@injectable() -export class DeepnoteNotebookEnvironmentMapper { - private static readonly STORAGE_KEY = 'deepnote.notebookEnvironmentMappings'; - private readonly workspaceState: Memento; - private mappings: Map; // notebookUri.fsPath -> environmentId - - private readonly _onDidRemoveEnvironment = new EventEmitter<{ notebookUri: Uri }>(); - private readonly _onDidSetEnvironment = new EventEmitter<{ notebookUri: Uri; environmentId: string }>(); - public readonly onDidRemoveEnvironment = this._onDidRemoveEnvironment.event; - public readonly onDidSetEnvironment = this._onDidSetEnvironment.event; - - constructor( - @inject(IExtensionContext) context: IExtensionContext, - @inject(IDisposableRegistry) disposables: IDisposableRegistry - ) { - this.workspaceState = context.workspaceState; - this.mappings = new Map(); - this.loadMappings(); - disposables.push(this._onDidSetEnvironment, this._onDidRemoveEnvironment); - } - - /** - * Get the environment ID selected for a notebook - * @param notebookUri The notebook URI (without query/fragment) - * @returns Environment ID, or undefined if not set - */ - public getEnvironmentForNotebook(notebookUri: Uri): string | undefined { - const key = notebookUri.fsPath; - return this.mappings.get(key); - } - - /** - * Set the environment for a notebook - * @param notebookUri The notebook URI (without query/fragment) - * @param environmentId The environment ID - */ - public async setEnvironmentForNotebook(notebookUri: Uri, environmentId: string): Promise { - const key = notebookUri.fsPath; - this.mappings.set(key, environmentId); - await this.saveMappings(); - logger.info(`Mapped notebook ${notebookUri.fsPath} to environment ${environmentId}`); - this._onDidSetEnvironment.fire({ notebookUri, environmentId }); - } - - /** - * Remove the environment mapping for a notebook - * @param notebookUri The notebook URI (without query/fragment) - */ - public async removeEnvironmentForNotebook(notebookUri: Uri): Promise { - const key = notebookUri.fsPath; - this.mappings.delete(key); - await this.saveMappings(); - logger.info(`Removed environment mapping for notebook ${notebookUri.fsPath}`); - this._onDidRemoveEnvironment.fire({ notebookUri }); - } - - /** - * Get all notebooks using a specific environment - * @param environmentId The environment ID - * @returns Array of notebook URIs - */ - public getNotebooksUsingEnvironment(environmentId: string): Uri[] { - const notebooks: Uri[] = []; - for (const [notebookPath, configId] of this.mappings.entries()) { - if (configId === environmentId) { - notebooks.push(Uri.file(notebookPath)); - } - } - return notebooks; - } - - /** - * Get all notebook-to-environment mappings - */ - public getAllMappings(): ReadonlyMap { - return new Map(this.mappings); - } - - /** - * Load mappings from workspace state - */ - private loadMappings(): void { - const stored = this.workspaceState.get>(DeepnoteNotebookEnvironmentMapper.STORAGE_KEY); - if (stored) { - this.mappings = new Map(Object.entries(stored)); - logger.info(`Loaded ${this.mappings.size} notebook-environment mappings`); - } - } - - /** - * Save mappings to workspace state - */ - private async saveMappings(): Promise { - const obj = Object.fromEntries(this.mappings.entries()); - await this.workspaceState.update(DeepnoteNotebookEnvironmentMapper.STORAGE_KEY, obj); - } -} diff --git a/src/kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.node.ts b/src/kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.node.ts new file mode 100644 index 0000000000..3aa467becc --- /dev/null +++ b/src/kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.node.ts @@ -0,0 +1,168 @@ +import { injectable, inject } from 'inversify'; +import { EventEmitter, Memento, Uri } from 'vscode'; + +import { IDisposableRegistry, IExtensionContext } from '../../../platform/common/types'; +import { resolveProjectIdForFile } from '../../../platform/deepnote/deepnoteProjectIdResolver'; +import { logger } from '../../../platform/logging'; +import { IDeepnoteProjectEnvironmentMapper } from '../types'; + +/** + * Manages the mapping between Deepnote projects and their selected environments. + * Stores selections in workspace state keyed by `project.id` so that sibling + * `.deepnote` files sharing a project automatically inherit the same environment + * and share a single runtime. + */ +@injectable() +export class DeepnoteProjectEnvironmentMapper implements IDeepnoteProjectEnvironmentMapper { + private static readonly LEGACY_STORAGE_KEY = 'deepnote.notebookEnvironmentMappings'; + private static readonly STORAGE_KEY = 'deepnote.projectEnvironmentMappings'; + + private readonly _onDidRemoveEnvironment = new EventEmitter<{ projectId: string }>(); + private readonly _onDidSetEnvironment = new EventEmitter<{ projectId: string; environmentId: string }>(); + + public readonly onDidRemoveEnvironment = this._onDidRemoveEnvironment.event; + public readonly onDidSetEnvironment = this._onDidSetEnvironment.event; + + private readonly workspaceState: Memento; + private mappings: Map = new Map(); // projectId -> environmentId + private readonly initializationPromise: Promise; + + constructor( + @inject(IExtensionContext) context: IExtensionContext, + @inject(IDisposableRegistry) disposables: IDisposableRegistry + ) { + this.workspaceState = context.workspaceState; + disposables.push(this._onDidSetEnvironment, this._onDidRemoveEnvironment); + + this.initializationPromise = this.initialize().catch((error) => { + logger.error('Failed to initialize DeepnoteProjectEnvironmentMapper', error); + }); + } + + public getAllMappings(): ReadonlyMap { + return new Map(this.mappings); + } + + public getEnvironmentForProject(projectId: string): string | undefined { + return this.mappings.get(projectId); + } + + public getProjectsUsingEnvironment(environmentId: string): string[] { + const projects: string[] = []; + for (const [projectId, mappedEnvironmentId] of this.mappings.entries()) { + if (mappedEnvironmentId === environmentId) { + projects.push(projectId); + } + } + + return projects; + } + + public async removeEnvironmentForProject(projectId: string): Promise { + await this.waitForInitialization(); + + if (!this.mappings.has(projectId)) { + return; + } + + this.mappings.delete(projectId); + await this.saveMappings(); + + logger.info(`Removed environment mapping for project ${projectId}`); + this._onDidRemoveEnvironment.fire({ projectId }); + } + + public async setEnvironmentForProject(projectId: string, environmentId: string): Promise { + await this.waitForInitialization(); + + this.mappings.set(projectId, environmentId); + await this.saveMappings(); + + logger.info(`Mapped project ${projectId} to environment ${environmentId}`); + this._onDidSetEnvironment.fire({ projectId, environmentId }); + } + + public async waitForInitialization(): Promise { + await this.initializationPromise; + } + + /** + * Load existing project-keyed mappings and run the one-shot migration from + * the legacy notebook-URI keyed storage. + */ + private async initialize(): Promise { + const stored = this.workspaceState.get>(DeepnoteProjectEnvironmentMapper.STORAGE_KEY); + if (stored) { + this.mappings = new Map(Object.entries(stored)); + logger.info(`Loaded ${this.mappings.size} project-environment mappings`); + } + + await this.migrateLegacyMappings(); + } + + /** + * Migrate entries from the legacy `deepnote.notebookEnvironmentMappings` + * workspace-state key (fsPath → environmentId) to the new project-id keyed + * key. The migration is a one-shot: after resolving project ids from each + * `.deepnote` file, the legacy key is cleared. + * + * Entries whose project id cannot be resolved (file missing, unparseable) + * are logged and skipped. + */ + private async migrateLegacyMappings(): Promise { + const legacyStored = this.workspaceState.get>( + DeepnoteProjectEnvironmentMapper.LEGACY_STORAGE_KEY + ); + + if (!legacyStored || Object.keys(legacyStored).length === 0) { + return; + } + + logger.info( + `Migrating ${ + Object.keys(legacyStored).length + } legacy notebook-keyed environment mappings to project-keyed storage` + ); + + // Clear the legacy key up front so the migration is strictly one-shot + // even if a later step fails. Worst case the user re-picks an env for + // a project — acceptable, and better than re-running migration on a + // later activation and overwriting project-keyed entries the user set + // in between. + await this.workspaceState.update(DeepnoteProjectEnvironmentMapper.LEGACY_STORAGE_KEY, undefined); + + const migratedEntries: Array<{ projectId: string; environmentId: string }> = []; + + for (const [fsPath, environmentId] of Object.entries(legacyStored)) { + try { + const projectId = await resolveProjectIdForFile(Uri.file(fsPath)); + if (!projectId) { + logger.warn(`Skipping legacy environment mapping for ${fsPath}: project id could not be resolved`); + continue; + } + + // Last-writer-wins if multiple siblings mapped to different envs + this.mappings.set(projectId, environmentId); + migratedEntries.push({ projectId, environmentId }); + } catch (error) { + logger.warn(`Failed to migrate legacy environment mapping for ${fsPath}`, error); + } + } + + await this.saveMappings(); + + logger.info(`Migrated ${migratedEntries.length} environment mappings to project-keyed storage`); + + for (const entry of migratedEntries) { + this._onDidSetEnvironment.fire(entry); + } + } + + /** + * Save mappings to workspace state + */ + private async saveMappings(): Promise { + const obj = Object.fromEntries(this.mappings.entries()); + await this.workspaceState.update(DeepnoteProjectEnvironmentMapper.STORAGE_KEY, obj); + } +} diff --git a/src/kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.unit.test.ts b/src/kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.unit.test.ts new file mode 100644 index 0000000000..8f12230bcc --- /dev/null +++ b/src/kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.unit.test.ts @@ -0,0 +1,333 @@ +import { assert } from 'chai'; +import esmock from 'esmock'; +import * as sinon from 'sinon'; +import { instance, mock, when } from 'ts-mockito'; +import { Memento, Uri } from 'vscode'; + +import { IDisposableRegistry, IExtensionContext } from '../../../platform/common/types'; +import { resetVSCodeMocks } from '../../../test/vscode-mock'; +import { DeepnoteProjectEnvironmentMapper } from './deepnoteProjectEnvironmentMapper.node'; + +const LEGACY_STORAGE_KEY = 'deepnote.notebookEnvironmentMappings'; +const STORAGE_KEY = 'deepnote.projectEnvironmentMappings'; + +/** + * Simple in-memory Memento used to exercise the mapper's save/load logic with + * real state transitions (rather than stubbing ts-mockito on `update` and + * `get`, which makes assertions awkward). + */ +class InMemoryMemento implements Memento { + private readonly store = new Map(); + + public get(key: string): T | undefined; + public get(key: string, defaultValue: T): T; + public get(key: string, defaultValue?: T): T | undefined { + return (this.store.has(key) ? (this.store.get(key) as T) : defaultValue) as T | undefined; + } + + public keys(): readonly string[] { + return Array.from(this.store.keys()); + } + + public async update(key: string, value: unknown): Promise { + if (value === undefined) { + this.store.delete(key); + } else { + this.store.set(key, value); + } + } +} + +suite('DeepnoteProjectEnvironmentMapper', () => { + let state: InMemoryMemento; + let context: IExtensionContext; + let disposables: IDisposableRegistry; + + setup(() => { + resetVSCodeMocks(); + state = new InMemoryMemento(); + + const mockContext = mock(); + when(mockContext.workspaceState).thenReturn(state); + context = instance(mockContext); + + disposables = []; + }); + + teardown(() => { + for (const disposable of disposables) { + disposable.dispose(); + } + }); + + suite('CRUD methods', () => { + test('getEnvironmentForProject returns undefined when not set', async () => { + const mapper = new DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + assert.strictEqual(mapper.getEnvironmentForProject('proj-1'), undefined); + }); + + test('setEnvironmentForProject persists and is readable via getEnvironmentForProject', async () => { + const mapper = new DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + await mapper.setEnvironmentForProject('proj-1', 'env-abc'); + + assert.strictEqual(mapper.getEnvironmentForProject('proj-1'), 'env-abc'); + assert.deepStrictEqual(state.get>(STORAGE_KEY), { 'proj-1': 'env-abc' }); + }); + + test('removeEnvironmentForProject clears mapping and persists', async () => { + const mapper = new DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + await mapper.setEnvironmentForProject('proj-1', 'env-abc'); + await mapper.removeEnvironmentForProject('proj-1'); + + assert.strictEqual(mapper.getEnvironmentForProject('proj-1'), undefined); + assert.deepStrictEqual(state.get>(STORAGE_KEY), {}); + }); + + test('removeEnvironmentForProject is a no-op when no mapping exists', async () => { + const mapper = new DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + const events: { projectId: string }[] = []; + disposables.push(mapper.onDidRemoveEnvironment((e) => events.push(e))); + + await mapper.removeEnvironmentForProject('proj-missing'); + + assert.strictEqual(events.length, 0, 'Should not fire event when mapping did not exist'); + }); + + test('getProjectsUsingEnvironment returns all projects pointing to the env', async () => { + const mapper = new DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + await mapper.setEnvironmentForProject('proj-1', 'env-A'); + await mapper.setEnvironmentForProject('proj-2', 'env-B'); + await mapper.setEnvironmentForProject('proj-3', 'env-A'); + + const projects = mapper.getProjectsUsingEnvironment('env-A').sort(); + assert.deepStrictEqual(projects, ['proj-1', 'proj-3']); + + assert.deepStrictEqual(mapper.getProjectsUsingEnvironment('env-none'), []); + }); + + test('getAllMappings returns an independent snapshot', async () => { + const mapper = new DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + await mapper.setEnvironmentForProject('proj-1', 'env-A'); + + const snapshot = mapper.getAllMappings(); + assert.deepStrictEqual( + Array.from(snapshot.entries()), + [['proj-1', 'env-A']], + 'Snapshot should reflect current state' + ); + + // Mutating after snapshot should not affect the returned copy + await mapper.setEnvironmentForProject('proj-2', 'env-B'); + assert.strictEqual(snapshot.has('proj-2'), false, 'Snapshot must not observe later mutations'); + }); + + test('persisted mappings are loaded on startup', async () => { + await state.update(STORAGE_KEY, { 'proj-pre': 'env-pre' }); + + const mapper = new DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + assert.strictEqual(mapper.getEnvironmentForProject('proj-pre'), 'env-pre'); + }); + }); + + suite('events', () => { + test('onDidSetEnvironment fires with { projectId, environmentId }', async () => { + const mapper = new DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + const events: { projectId: string; environmentId: string }[] = []; + disposables.push(mapper.onDidSetEnvironment((e) => events.push(e))); + + await mapper.setEnvironmentForProject('proj-1', 'env-A'); + + assert.deepStrictEqual(events, [{ projectId: 'proj-1', environmentId: 'env-A' }]); + }); + + test('onDidRemoveEnvironment fires with { projectId } after remove', async () => { + const mapper = new DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + await mapper.setEnvironmentForProject('proj-1', 'env-A'); + + const events: { projectId: string }[] = []; + disposables.push(mapper.onDidRemoveEnvironment((e) => events.push(e))); + + await mapper.removeEnvironmentForProject('proj-1'); + + assert.deepStrictEqual(events, [{ projectId: 'proj-1' }]); + }); + }); + + suite('legacy migration', () => { + // Migration reads the legacy fsPath -> envId mapping and resolves each + // fsPath to a project id via `resolveProjectIdForFile`. The helper + // reads YAML and parses it strictly, which is cumbersome to stub via + // the vscode `workspace.fs.readFile` path — so we use `esmock` to + // replace the resolver import entirely for these tests. + let MapperModule: any; + let resolveProjectIdForFileStub: sinon.SinonStub; + + setup(async () => { + resolveProjectIdForFileStub = sinon.stub(); + + MapperModule = await esmock('./deepnoteProjectEnvironmentMapper.node', { + '../../../platform/deepnote/deepnoteProjectIdResolver': { + resolveProjectIdForFile: (uri: Uri) => resolveProjectIdForFileStub(uri), + // The mapper doesn't use this symbol, but esmock's module + // replacement forces us to provide a full replacement. + resolveProjectIdForNotebook: () => undefined + } + }); + }); + + teardown(() => { + esmock.purge(MapperModule); + }); + + test('migrates legacy fsPath-keyed entries to projectId-keyed storage', async () => { + const goodPath = '/workspace/good.deepnote'; + const missingPath = '/workspace/gone.deepnote'; + + await state.update(LEGACY_STORAGE_KEY, { + [goodPath]: 'env-migrated', + [missingPath]: 'env-orphan' + }); + + resolveProjectIdForFileStub.callsFake(async (uri: Uri) => { + if (uri.fsPath === goodPath) { + return 'proj-migrated'; + } + return undefined; // simulates YAML read / parse failure + }); + + const mapper = new MapperModule.DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + // Resolved entry is now keyed by project id + assert.strictEqual(mapper.getEnvironmentForProject('proj-migrated'), 'env-migrated'); + // Unresolved entry is skipped — fsPath must not leak into the new storage + assert.strictEqual(mapper.getEnvironmentForProject(missingPath), undefined); + + // Legacy key is cleared + assert.strictEqual( + state.get(LEGACY_STORAGE_KEY), + undefined, + 'Legacy workspace-state key should be removed after migration' + ); + + // New key has only the migrated entry + const stored = state.get>(STORAGE_KEY); + assert.deepStrictEqual(stored, { 'proj-migrated': 'env-migrated' }); + }); + + test('fires onDidSetEnvironment for each migrated entry', async () => { + const path1 = '/workspace/a.deepnote'; + const path2 = '/workspace/b.deepnote'; + + await state.update(LEGACY_STORAGE_KEY, { + [path1]: 'env-1', + [path2]: 'env-2' + }); + + resolveProjectIdForFileStub.callsFake(async (uri: Uri) => { + if (uri.fsPath === path1) return 'proj-1'; + if (uri.fsPath === path2) return 'proj-2'; + return undefined; + }); + + const events: { projectId: string; environmentId: string }[] = []; + // Construct the mapper and attach the listener synchronously before + // awaiting initialization so the migration events are captured. + const mapper = new MapperModule.DeepnoteProjectEnvironmentMapper(context, disposables); + disposables.push( + mapper.onDidSetEnvironment((e: { projectId: string; environmentId: string }) => events.push(e)) + ); + + await mapper.waitForInitialization(); + + const sorted = events.slice().sort((a, b) => a.projectId.localeCompare(b.projectId)); + assert.deepStrictEqual(sorted, [ + { projectId: 'proj-1', environmentId: 'env-1' }, + { projectId: 'proj-2', environmentId: 'env-2' } + ]); + }); + + test('legacy migration is a no-op when there are no legacy entries', async () => { + // No legacy key set at all — mapper must initialize cleanly + const mapper = new MapperModule.DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + assert.deepStrictEqual(Array.from(mapper.getAllMappings().entries()), []); + assert.strictEqual(state.get(LEGACY_STORAGE_KEY), undefined); + assert.strictEqual(resolveProjectIdForFileStub.called, false); + }); + + test('legacy migration respects last-writer-wins for duplicate project ids', async () => { + const pathA = '/workspace/a.deepnote'; + const pathB = '/workspace/b.deepnote'; + + await state.update(LEGACY_STORAGE_KEY, { + [pathA]: 'env-first', + [pathB]: 'env-second' + }); + + resolveProjectIdForFileStub.callsFake(async () => 'proj-shared'); + + const mapper = new MapperModule.DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + // Last one wins — the exact value depends on iteration order, but + // the critical invariant is no crash on collision and some value + // is retained. + const envForShared = mapper.getEnvironmentForProject('proj-shared'); + assert.oneOf( + envForShared, + ['env-first', 'env-second'], + 'One of the two entries should win; migration must not crash on collision' + ); + }); + + test('already-migrated workspace-state is loaded alongside a no-op legacy migration', async () => { + await state.update(STORAGE_KEY, { 'proj-loaded': 'env-loaded' }); + + const mapper = new MapperModule.DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + assert.strictEqual(mapper.getEnvironmentForProject('proj-loaded'), 'env-loaded'); + }); + + test('legacy key is cleared before any resolver work so a crash mid-migration cannot re-run it', async () => { + const path1 = '/workspace/a.deepnote'; + await state.update(LEGACY_STORAGE_KEY, { [path1]: 'env-1' }); + + let legacyStillPresentWhenResolverRan: boolean | undefined; + resolveProjectIdForFileStub.callsFake(async () => { + legacyStillPresentWhenResolverRan = state.get(LEGACY_STORAGE_KEY) !== undefined; + return 'proj-1'; + }); + + const mapper = new MapperModule.DeepnoteProjectEnvironmentMapper(context, disposables); + await mapper.waitForInitialization(); + + assert.strictEqual( + legacyStillPresentWhenResolverRan, + false, + 'Legacy key must be cleared before the resolver runs so a crash during migration cannot cause a re-run' + ); + assert.strictEqual(mapper.getEnvironmentForProject('proj-1'), 'env-1'); + }); + }); +}); diff --git a/src/kernels/deepnote/types.ts b/src/kernels/deepnote/types.ts index a0c17a31ba..a8a17627f2 100644 --- a/src/kernels/deepnote/types.ts +++ b/src/kernels/deepnote/types.ts @@ -157,7 +157,8 @@ export interface IDeepnoteServerStarter { * @param venvPath The path to the venv * @param managedVenv Whether the venv is managed by this extension (created by us) * @param environmentId The environment ID (for server management) - * @param deepnoteFileUri The URI of the .deepnote file + * @param projectId The Deepnote project id (shared across sibling `.deepnote` files) + * @param deepnoteFileUri The URI of the `.deepnote` file (used for working directory + SQL env vars) * @param token Cancellation token to cancel the operation * @returns Connection information (URL, port, etc.) */ @@ -167,17 +168,17 @@ export interface IDeepnoteServerStarter { managedVenv: boolean, additionalPackages: string[], environmentId: string, + projectId: string, deepnoteFileUri: vscode.Uri, token?: vscode.CancellationToken ): Promise; /** - * Stops the deepnote-toolkit server for a kernel environment. - * @param environmentId The environment ID + * Stops the deepnote-toolkit server for a project. + * @param projectId The Deepnote project id * @param token Cancellation token to cancel the operation */ - // stopServer(environmentId: string, token?: vscode.CancellationToken): Promise; - stopServer(deepnoteFileUri: vscode.Uri, token?: vscode.CancellationToken): Promise; + stopServer(projectId: string, token?: vscode.CancellationToken): Promise; /** * Disposes all server processes and resources. @@ -214,7 +215,7 @@ export interface IDeepnoteKernelAutoSelector { * @param notebook The notebook document * @param environmentId The environment ID */ - clearControllerForEnvironment(notebook: vscode.NotebookDocument, environmentId: string): void; + clearControllerForEnvironment(notebook: vscode.NotebookDocument, environmentId: string): Promise; /** * Ensure an environment is configured for the notebook before execution. @@ -320,50 +321,56 @@ export interface IDeepnoteEnvironmentManager { dispose(): void; } -export const IDeepnoteNotebookEnvironmentMapper = Symbol('IDeepnoteNotebookEnvironmentMapper'); -export interface IDeepnoteNotebookEnvironmentMapper { +export const IDeepnoteProjectEnvironmentMapper = Symbol('IDeepnoteProjectEnvironmentMapper'); +export interface IDeepnoteProjectEnvironmentMapper { + /** + * Get all project-to-environment mappings + * @returns Map of projectId → environmentId + */ + getAllMappings(): ReadonlyMap; + /** - * Get the environment ID selected for a notebook - * @param notebookUri The notebook URI (without query/fragment) + * Get the environment ID selected for a project + * @param projectId The Deepnote project id * @returns Environment ID, or undefined if not set */ - getEnvironmentForNotebook(notebookUri: vscode.Uri): string | undefined; + getEnvironmentForProject(projectId: string): string | undefined; /** - * Set the environment for a notebook - * @param notebookUri The notebook URI (without query/fragment) + * Get all projects using a specific environment * @param environmentId The environment ID + * @returns Array of project ids */ - setEnvironmentForNotebook(notebookUri: vscode.Uri, environmentId: string): Promise; + getProjectsUsingEnvironment(environmentId: string): string[]; /** - * Remove the environment mapping for a notebook - * @param notebookUri The notebook URI (without query/fragment) + * Event fired when an environment mapping is removed for a project */ - removeEnvironmentForNotebook(notebookUri: vscode.Uri): Promise; + onDidRemoveEnvironment: vscode.Event<{ projectId: string }>; /** - * Get all notebooks using a specific environment - * @param environmentId The environment ID - * @returns Array of notebook URIs + * Event fired when an environment is set for a project */ - getNotebooksUsingEnvironment(environmentId: string): vscode.Uri[]; + onDidSetEnvironment: vscode.Event<{ projectId: string; environmentId: string }>; /** - * Get all notebook-to-environment mappings - * @returns Map of notebookUri.fsPath → environmentId + * Remove the environment mapping for a project + * @param projectId The Deepnote project id */ - getAllMappings(): ReadonlyMap; + removeEnvironmentForProject(projectId: string): Promise; /** - * Event fired when an environment mapping is removed for a notebook + * Set the environment for a project + * @param projectId The Deepnote project id + * @param environmentId The environment ID */ - onDidRemoveEnvironment: vscode.Event<{ notebookUri: vscode.Uri }>; + setEnvironmentForProject(projectId: string, environmentId: string): Promise; /** - * Event fired when an environment is set for a notebook + * Wait until the mapper has finished loading mappings and running the + * one-shot migration from the legacy notebook-URI keyed storage. */ - onDidSetEnvironment: vscode.Event<{ notebookUri: vscode.Uri; environmentId: string }>; + waitForInitialization(): Promise; } export const IDeepnoteLspClientManager = Symbol('IDeepnoteLspClientManager'); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index d1330ae14f..3a7e07d312 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -30,7 +30,7 @@ import { IDeepnoteEnvironmentManager, IDeepnoteKernelAutoSelector, IDeepnoteLspClientManager, - IDeepnoteNotebookEnvironmentMapper, + IDeepnoteProjectEnvironmentMapper, IDeepnoteServerProvider, IDeepnoteServerStarter, IDeepnoteToolkitInstaller @@ -50,6 +50,7 @@ import { JVSC_EXTENSION_ID, STANDARD_OUTPUT_CHANNEL } from '../../platform/commo import { getDisplayPath } from '../../platform/common/platform/fs-paths.node'; import { IConfigurationService, IDisposableRegistry, IOutputChannel } from '../../platform/common/types'; import { disposeAsync } from '../../platform/common/utils'; +import { resolveProjectIdForNotebook } from '../../platform/deepnote/deepnoteProjectIdResolver'; import { createDeepnoteServerConfigHandle } from '../../platform/deepnote/deepnoteServerUtils.node'; import { DeepnoteKernelError, DeepnoteToolkitMissingError } from '../../platform/errors/deepnoteKernelErrors'; import { logger } from '../../platform/logging'; @@ -77,7 +78,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, private readonly notebookEnvironmentsIds = new Map(); // Track per-notebook placeholder controllers for notebooks without configured environments private readonly placeholderControllers = new Map(); - // Track server handles per PROJECT (baseFileUri) - one server per project + // Track server handles per PROJECT (projectId) - one server per project, shared across sibling files private readonly projectServerHandles = new Map(); // Track projects where we need to run init notebook (set during controller setup) private readonly projectsPendingInitNotebook = new Map< @@ -102,8 +103,8 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, @inject(IDeepnoteRequirementsHelper) private readonly requirementsHelper: IDeepnoteRequirementsHelper, @inject(IDeepnoteEnvironmentManager) private readonly environmentManager: IDeepnoteEnvironmentManager, @inject(IDeepnoteServerStarter) private readonly serverStarter: IDeepnoteServerStarter, - @inject(IDeepnoteNotebookEnvironmentMapper) - private readonly notebookEnvironmentMapper: IDeepnoteNotebookEnvironmentMapper, + @inject(IDeepnoteProjectEnvironmentMapper) + private readonly projectEnvironmentMapper: IDeepnoteProjectEnvironmentMapper, @inject(IOutputChannel) @named(STANDARD_OUTPUT_CHANNEL) private readonly outputChannel: IOutputChannel, @inject(IDeepnoteToolkitInstaller) private readonly toolkitInstaller: IDeepnoteToolkitInstaller ) {} @@ -391,10 +392,18 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, ): Promise { const baseFileUri = notebook.uri.with({ query: '', fragment: '' }); const notebookKey = notebook.uri.toString(); - const projectKey = baseFileUri.fsPath; logger.info(`Switching controller environment for ${getDisplayPath(notebook.uri)}`); + const projectId = await resolveProjectIdForNotebook(notebook); + if (!projectId) { + logger.error(`Cannot resolve Deepnote project id for ${getDisplayPath(notebook.uri)}`); + return; + } + + // Wait for the mapper to finish migration so legacy entries are resolved + await this.projectEnvironmentMapper.waitForInitialization(); + // Check if any cells are executing and log a warning const kernel = this.kernelProvider.get(notebook); if (kernel) { @@ -411,10 +420,10 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, this.notebookConnectionMetadata.delete(notebookKey); // Clear old server handle - new environment will register a new handle - const oldServerHandle = this.projectServerHandles.get(projectKey); + const oldServerHandle = this.projectServerHandles.get(projectId); if (oldServerHandle) { logger.info(`Clearing old server handle from tracking: ${oldServerHandle}`); - this.projectServerHandles.delete(projectKey); + this.projectServerHandles.delete(projectId); } // Stop existing LSP clients so new ones can be created with fresh environment @@ -425,11 +434,11 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Update the controller with new environment's metadata // Because we use notebook-based controller IDs, addOrUpdate will call updateConnection() // on the existing controller instead of creating a new one - const environmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri); + const environmentId = this.projectEnvironmentMapper.getEnvironmentForProject(projectId); const environment = environmentId ? this.environmentManager.getEnvironment(environmentId) : undefined; if (environment == null) { - await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(baseFileUri); + await this.projectEnvironmentMapper.removeEnvironmentForProject(projectId); logger.error(`No environment found for notebook ${getDisplayPath(notebook.uri)}`); return; } @@ -439,7 +448,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, environment, baseFileUri, notebookKey, - projectKey, + projectId, progress, token ); @@ -452,14 +461,23 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, progress: { report(value: { message?: string; increment?: number }): void }, token: CancellationToken ): Promise { - // baseFileUri identifies the PROJECT (without query/fragment) + // baseFileUri identifies the FILE on disk (without query/fragment) const baseFileUri = notebook.uri.with({ query: '', fragment: '' }); // notebookKey uniquely identifies THIS NOTEBOOK (includes query with notebook ID) const notebookKey = notebook.uri.toString(); - // projectKey identifies the PROJECT for server tracking - const projectKey = baseFileUri.fsPath; - const environmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri); + const projectId = await resolveProjectIdForNotebook(notebook); + if (!projectId) { + logger.warn(`Cannot resolve Deepnote project id for ${getDisplayPath(notebook.uri)}`); + await this.selectPlaceholderController(notebook); + + return false; + } + + // Wait for the mapper to finish migration so legacy entries are resolved + await this.projectEnvironmentMapper.waitForInitialization(); + + const environmentId = this.projectEnvironmentMapper.getEnvironmentForProject(projectId); if (environmentId == null) { await this.selectPlaceholderController(notebook); @@ -471,7 +489,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, if (environment == null) { logger.info(`No environment found for notebook ${getDisplayPath(notebook.uri)}`); - await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(baseFileUri); + await this.projectEnvironmentMapper.removeEnvironmentForProject(projectId); await this.selectPlaceholderController(notebook); return false; @@ -482,7 +500,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, environment, baseFileUri, notebookKey, - projectKey, + projectId, progress, token ); @@ -495,7 +513,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, configuration: DeepnoteEnvironment, baseFileUri: Uri, notebookKey: string, - projectKey: string, + projectId: string, progress: { report(value: { message?: string; increment?: number }): void }, progressToken: CancellationToken ): Promise { @@ -553,6 +571,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, configuration.managedVenv, configuration.packages ?? [], configuration.id, + projectId, baseFileUri, progressToken ); @@ -568,12 +587,12 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const serverProviderHandle: JupyterServerProviderHandle = { extensionId: JVSC_EXTENSION_ID, id: 'deepnote-server', - handle: createDeepnoteServerConfigHandle(configuration.id, baseFileUri) + handle: createDeepnoteServerConfigHandle(configuration.id, projectId) }; // Register the server with the provider (one server per PROJECT) this.serverProvider.registerServer(serverProviderHandle.handle, serverInfo); - this.projectServerHandles.set(projectKey, serverProviderHandle.handle); + this.projectServerHandles.set(projectId, serverProviderHandle.handle); const lspInterpreterUri = this.getVenvInterpreterUri(configuration.venvPath); @@ -669,11 +688,10 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, this.notebookControllers.set(notebookKey, controller); // Prepare init notebook execution - const projectId = notebook.metadata?.deepnoteProjectId; const notebookIdForProject = notebook.metadata?.deepnoteNotebookId; - const project = projectId - ? (this.notebookManager.getOriginalProject(projectId, notebookIdForProject) as DeepnoteFile | undefined) - : undefined; + const project = this.notebookManager.getOriginalProject(projectId, notebookIdForProject) as + | DeepnoteFile + | undefined; if (project) { // Only create requirements.txt if requirements have changed from what's on disk @@ -689,8 +707,8 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, logger.info(`Skipping requirements.txt creation for project ${projectId} (no changes detected)`); } - if (project.project.initNotebookId && !this.notebookManager.hasInitNotebookBeenRun(projectId!)) { - this.projectsPendingInitNotebook.set(projectId!, { notebook, project }); + if (project.project.initNotebookId && !this.notebookManager.hasInitNotebookBeenRun(projectId)) { + this.projectsPendingInitNotebook.set(projectId, { notebook, project }); logger.info(`Init notebook will run automatically when kernel starts for project ${projectId}`); } } @@ -794,13 +812,21 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, const baseFileUri = notebook.uri.with({ query: '', fragment: '' }); const notebookKey = notebook.uri.toString(); - const projectKey = baseFileUri.fsPath; - const existingEnvironmentId = this.notebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri); + const projectId = await resolveProjectIdForNotebook(notebook); + if (!projectId) { + logger.warn(`Cannot resolve Deepnote project id for ${getDisplayPath(notebook.uri)}`); + return false; + } + + // Wait for the mapper to finish migration so legacy entries are resolved + await this.projectEnvironmentMapper.waitForInitialization(); + + const existingEnvironmentId = this.projectEnvironmentMapper.getEnvironmentForProject(projectId); // No environment configured - need to pick one if (!existingEnvironmentId) { - return this.pickAndSetupEnvironment(notebook, baseFileUri, notebookKey, projectKey, token); + return this.pickAndSetupEnvironment(notebook, baseFileUri, notebookKey, projectId, token); } const environment = this.environmentManager.getEnvironment(existingEnvironmentId); @@ -808,9 +834,9 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Environment no longer exists - remove stale mapping and pick a new one if (!environment) { logger.info(`Removing stale environment mapping for ${getDisplayPath(notebook.uri)}`); - await this.notebookEnvironmentMapper.removeEnvironmentForNotebook(baseFileUri); + await this.projectEnvironmentMapper.removeEnvironmentForProject(projectId); - return this.pickAndSetupEnvironment(notebook, baseFileUri, notebookKey, projectKey, token); + return this.pickAndSetupEnvironment(notebook, baseFileUri, notebookKey, projectId, token); } const existingController = this.notebookControllers.get(notebookKey); @@ -831,7 +857,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, environment, baseFileUri, notebookKey, - projectKey, + projectId, token ); } @@ -848,7 +874,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, )}, triggering setup` ); - return this.setupKernelForEnvironment(notebook, environment, baseFileUri, notebookKey, projectKey, token); + return this.setupKernelForEnvironment(notebook, environment, baseFileUri, notebookKey, projectId, token); } /** @@ -858,7 +884,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, notebook: NotebookDocument, baseFileUri: Uri, notebookKey: string, - projectKey: string, + projectId: string, token: CancellationToken ): Promise { Cancellation.throwIfCanceled(token); @@ -874,14 +900,14 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, Cancellation.throwIfCanceled(token); - await this.notebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, selectedEnvironment.id); + await this.projectEnvironmentMapper.setEnvironmentForProject(projectId, selectedEnvironment.id); const result = await this.setupKernelForEnvironment( notebook, selectedEnvironment, baseFileUri, notebookKey, - projectKey, + projectId, token ); @@ -900,7 +926,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, environment: DeepnoteEnvironment, baseFileUri: Uri, notebookKey: string, - projectKey: string, + projectId: string, token: CancellationToken ): Promise { try { @@ -916,7 +942,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, environment, baseFileUri, notebookKey, - projectKey, + projectId, progress, progressToken ); @@ -948,13 +974,18 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, * Clear the controller selection for a notebook using a specific environment. * This is used when deleting an environment to unselect its controller from any open notebooks. */ - public clearControllerForEnvironment(notebook: NotebookDocument, environmentId: string): void { + public async clearControllerForEnvironment(notebook: NotebookDocument, environmentId: string): Promise { const selectedController = this.controllerRegistration.getSelected(notebook); if (!selectedController || selectedController.connection.kind !== 'startUsingDeepnoteKernel') { return; } - const expectedHandle = createDeepnoteServerConfigHandle(environmentId, notebook.uri); + const projectId = await resolveProjectIdForNotebook(notebook); + if (!projectId) { + return; + } + + const expectedHandle = createDeepnoteServerConfigHandle(environmentId, projectId); if (selectedController.connection.serverProviderHandle.handle === expectedHandle) { // Unselect the controller by setting affinity to Default diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts index 824635343c..bcab271a14 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -4,9 +4,10 @@ import { anything, instance, mock, verify, when } from 'ts-mockito'; import { DeepnoteKernelAutoSelector } from './deepnoteKernelAutoSelector.node'; import { createMockChildProcess } from '../../kernels/deepnote/deepnoteTestHelpers.node'; import { + DEEPNOTE_NOTEBOOK_TYPE, IDeepnoteEnvironmentManager, IDeepnoteLspClientManager, - IDeepnoteNotebookEnvironmentMapper, + IDeepnoteProjectEnvironmentMapper, IDeepnoteServerProvider, IDeepnoteServerStarter, IDeepnoteToolkitInstaller @@ -41,7 +42,7 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { let mockRequirementsHelper: IDeepnoteRequirementsHelper; let mockEnvironmentManager: IDeepnoteEnvironmentManager; let mockServerStarter: IDeepnoteServerStarter; - let mockNotebookEnvironmentMapper: IDeepnoteNotebookEnvironmentMapper; + let mockProjectEnvironmentMapper: IDeepnoteProjectEnvironmentMapper; let mockOutputChannel: IOutputChannel; let mockToolkitInstaller: IDeepnoteToolkitInstaller; @@ -53,6 +54,8 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { let mockNewController: IVSCodeNotebookController; let sandbox: sinon.SinonSandbox; + const testProjectId = 'project-123'; + setup(() => { resetVSCodeMocks(); sandbox = sinon.createSandbox(); @@ -72,17 +75,20 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { mockEnvironmentManager = mock(); mockServerStarter = mock(); mockToolkitInstaller = mock(); - mockNotebookEnvironmentMapper = mock(); + mockProjectEnvironmentMapper = mock(); mockOutputChannel = mock(); + // Mapper init resolves immediately in all tests unless overridden + when(mockProjectEnvironmentMapper.waitForInitialization()).thenResolve(); + mockProgress = { report: sandbox.stub() }; mockCancellationToken = mock(); // Create mock notebook mockNotebook = { uri: Uri.parse('file:///test/notebook.deepnote?notebook=123'), - notebookType: 'deepnote', - metadata: { deepnoteProjectId: 'project-123' }, + notebookType: DEEPNOTE_NOTEBOOK_TYPE, + metadata: { deepnoteProjectId: testProjectId }, // Add minimal required properties for NotebookDocument version: 1, isDirty: false, @@ -137,7 +143,7 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { instance(mockRequirementsHelper), instance(mockEnvironmentManager), instance(mockServerStarter), - instance(mockNotebookEnvironmentMapper), + instance(mockProjectEnvironmentMapper), instance(mockOutputChannel), instance(mockToolkitInstaller) ); @@ -161,8 +167,8 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // Create mock environment const mockEnvironment = createMockEnvironment('test-env-id', 'Test Environment'); - // Mock environment mapper and manager - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn('test-env-id'); + // Mock project environment mapper and manager + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn('test-env-id'); when(mockEnvironmentManager.getEnvironment('test-env-id')).thenReturn(mockEnvironment); when(mockKernelProvider.get(mockNotebook)).thenReturn(instance(mockKernel)); @@ -198,8 +204,8 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // Create mock environment const mockEnvironment = createMockEnvironment('test-env-id', 'Test Environment'); - // Mock environment mapper and manager - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn('test-env-id'); + // Mock project environment mapper and manager + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn('test-env-id'); when(mockEnvironmentManager.getEnvironment('test-env-id')).thenReturn(mockEnvironment); // Stub ensureKernelSelectedWithConfiguration to verify it's called @@ -233,8 +239,8 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // Create mock environment const mockEnvironment = createMockEnvironment('test-env-id', 'Test Environment'); - // Mock environment mapper and manager - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn('test-env-id'); + // Mock project environment mapper and manager + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn('test-env-id'); when(mockEnvironmentManager.getEnvironment('test-env-id')).thenReturn(mockEnvironment); // Stub ensureKernelSelectedWithConfiguration to verify delegation @@ -264,8 +270,8 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // Create mock environment const mockEnvironment = createMockEnvironment('test-env-id', 'Test Environment'); - // Mock environment mapper and manager - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn('test-env-id'); + // Mock project environment mapper and manager + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn('test-env-id'); when(mockEnvironmentManager.getEnvironment('test-env-id')).thenReturn(mockEnvironment); // Stub ensureKernelSelectedWithConfiguration to verify it receives the token @@ -343,9 +349,9 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { }); suite('ensureKernelSelected', () => { - test('should return false when no environment ID is assigned to the notebook', async () => { - // Mock environment mapper to return null (no environment assigned) - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn(undefined); + test('should return false when no environment ID is assigned to the project', async () => { + // Mock environment mapper to return undefined (no environment assigned) + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn(undefined); // Stub ensureKernelSelectedWithConfiguration to track if it gets called const ensureKernelSelectedStub = sandbox.stub(selector, 'ensureKernelSelectedWithConfiguration').resolves(); @@ -367,7 +373,7 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { false, 'ensureKernelSelectedWithConfiguration should not be called' ); - verify(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).once(); + verify(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).once(); }); test('should return false and remove mapping when environment is not found', async () => { @@ -375,13 +381,13 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { const environmentId = 'missing-env-id'; // Mock environment mapper to return an ID - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn(environmentId); + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn(environmentId); - // Mock environment manager to return null (environment not found) + // Mock environment manager to return undefined (environment not found) when(mockEnvironmentManager.getEnvironment(environmentId)).thenReturn(undefined); // Mock remove environment mapping - when(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(anything())).thenResolve(); + when(mockProjectEnvironmentMapper.removeEnvironmentForProject(testProjectId)).thenResolve(); // Stub ensureKernelSelectedWithConfiguration to track if it gets called const ensureKernelSelectedStub = sandbox.stub(selector, 'ensureKernelSelectedWithConfiguration').resolves(); @@ -403,21 +409,20 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { false, 'ensureKernelSelectedWithConfiguration should not be called' ); - verify(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).once(); + verify(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).once(); verify(mockEnvironmentManager.getEnvironment(environmentId)).once(); - verify(mockNotebookEnvironmentMapper.removeEnvironmentForNotebook(anything())).once(); + verify(mockProjectEnvironmentMapper.removeEnvironmentForProject(testProjectId)).once(); }); test('should return true and call ensureKernelSelectedWithConfiguration when environment is found', async () => { // Arrange const baseFileUri = mockNotebook.uri.with({ query: '', fragment: '' }); const notebookKey = mockNotebook.uri.toString(); - const projectKey = baseFileUri.fsPath; const environmentId = 'test-env-id'; const mockEnvironment = createMockEnvironment(environmentId, 'Test Environment'); // Mock environment mapper to return an ID - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).thenReturn(environmentId); + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn(environmentId); // Mock environment manager to return the environment when(mockEnvironmentManager.getEnvironment(environmentId)).thenReturn(mockEnvironment); @@ -443,19 +448,89 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { 'ensureKernelSelectedWithConfiguration should be called once' ); - // Verify it was called with correct arguments + // Verify it was called with correct arguments (notebook, env, baseFileUri, notebookKey, projectId, progress, token) const callArgs = ensureKernelSelectedStub.firstCall.args; assert.strictEqual(callArgs[0], mockNotebook, 'First arg should be notebook'); assert.strictEqual(callArgs[1], mockEnvironment, 'Second arg should be environment'); assert.strictEqual(callArgs[2].toString(), baseFileUri.toString(), 'Third arg should be baseFileUri'); assert.strictEqual(callArgs[3], notebookKey, 'Fourth arg should be notebookKey'); - assert.strictEqual(callArgs[4], projectKey, 'Fifth arg should be projectKey'); + assert.strictEqual(callArgs[4], testProjectId, 'Fifth arg should be projectId'); assert.strictEqual(callArgs[5], mockProgress, 'Sixth arg should be progress'); assert.strictEqual(callArgs[6], instance(mockCancellationToken), 'Seventh arg should be token'); - verify(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(anything())).once(); + verify(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).once(); verify(mockEnvironmentManager.getEnvironment(environmentId)).once(); }); + + test('sibling notebooks sharing a projectId resolve to the same environment mapping', async () => { + // Primary plan use-case: two distinct .deepnote files that share the same + // project.id must both resolve to the same environment via the mapper. + const environmentId = 'shared-env'; + const mockEnvironment = createMockEnvironment(environmentId, 'Shared Environment'); + + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn(environmentId); + when(mockEnvironmentManager.getEnvironment(environmentId)).thenReturn(mockEnvironment); + + // Stub the heavy downstream setup call so we can check dispatching only + const ensureKernelSelectedStub = sandbox.stub(selector, 'ensureKernelSelectedWithConfiguration').resolves(); + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); + + const siblingNotebook1 = { + uri: Uri.parse('file:///test/sibling-one.deepnote?notebook=1'), + notebookType: DEEPNOTE_NOTEBOOK_TYPE, + metadata: { deepnoteProjectId: testProjectId }, + version: 1, + isDirty: false, + isUntitled: false, + isClosed: false, + cellCount: 0, + cellAt: () => { + throw new Error('Not implemented'); + }, + getCells: () => [], + save: async () => true + } as unknown as NotebookDocument; + + const siblingNotebook2 = { + uri: Uri.parse('file:///test/sibling-two.deepnote?notebook=2'), + notebookType: DEEPNOTE_NOTEBOOK_TYPE, + metadata: { deepnoteProjectId: testProjectId }, + version: 1, + isDirty: false, + isUntitled: false, + isClosed: false, + cellCount: 0, + cellAt: () => { + throw new Error('Not implemented'); + }, + getCells: () => [], + save: async () => true + } as unknown as NotebookDocument; + + const result1 = await selector.ensureKernelSelected( + siblingNotebook1, + mockProgress, + instance(mockCancellationToken) + ); + const result2 = await selector.ensureKernelSelected( + siblingNotebook2, + mockProgress, + instance(mockCancellationToken) + ); + + assert.strictEqual(result1, true); + assert.strictEqual(result2, true); + + // Mapper is queried by projectId only, so both siblings see the same env + verify(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).twice(); + + // Both calls passed the same projectId and environment through to configuration + assert.strictEqual(ensureKernelSelectedStub.callCount, 2); + for (const call of ensureKernelSelectedStub.getCalls()) { + assert.strictEqual(call.args[1], mockEnvironment, 'Same env is passed through for sibling'); + assert.strictEqual(call.args[4], testProjectId, 'Same projectId is passed through for sibling'); + } + }); }); // Priority 1 Tests - Critical for environment switching @@ -507,64 +582,10 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // Priority 1 Integration Tests - Critical for environment switching suite('Priority 1: Integration Tests (IT-1, IT-8)', () => { test('IT-1: Full environment switch flow is validated by existing tests', () => { - // IT-1 requires testing the full environment switch flow: - // 1. Notebook mapped to environment B - // 2. New controller for B created and selected - // 3. Old controller for A left alive (not disposed) to handle queued executions - // 4. Can execute cell successfully on B - // - // THIS IS VALIDATED BY EXISTING TESTS: - // - // 1. "should switch from one environment to another" (line 260) - // - Simulates switching from env-a to env-b - // - Validates rebuildController flow with environment change - // - // 2. "should NOT dispose old controller..." (line 178) - // - Validates that old controller is NOT disposed - // - This prevents "DISPOSED" errors for queued cell executions - // - Old controller will be garbage collected naturally - // - // 3. "should clear cached controller and metadata" (line 109) - // - Validates state clearing before rebuild - // - Ensures clean state for new environment - // - // 4. "should unregister old server handle" (line 151) - // - Validates server cleanup during switch - // - // Full integration testing with actual cell execution requires a running VS Code - // instance and is better suited for E2E tests. These unit tests validate all the - // critical invariants that make environment switching work correctly. - assert.ok(true, 'IT-1 requirements validated by existing rebuildController tests'); }); test('IT-8: Execute cell immediately after switch validated by disposal order tests', () => { - // IT-8 requires: "Execute cell immediately after environment switch" - // Verify: - // 1. Cell executes successfully - // 2. No "controller disposed" error - // 3. Output shows new environment - // - // THIS IS VALIDATED BY THE NON-DISPOSAL APPROACH: - // - // The test on line 178 validates that old controllers are NOT disposed. - // - // This prevents the "controller disposed" error because: - // - VS Code may have queued cell executions that reference the old controller - // - If we disposed the old controller, those executions would fail with "DISPOSED" error - // - By leaving the old controller alive, queued executions complete successfully - // - New cell executions use the new controller (it's now preferred) - // - The old controller will be garbage collected when no longer referenced - // - // The implementation at deepnoteKernelAutoSelector.node.ts:306-315 does this: - // // IMPORTANT: We do NOT dispose the old controller here - // // Reason: VS Code may have queued cell executions that reference the old controller - // // If we dispose it immediately, those queued executions will fail with "DISPOSED" error - // // Instead, we let the old controller stay alive - it will be garbage collected eventually - // - // Full integration testing with actual cell execution requires a running VS Code - // instance with real kernel execution, which is better suited for E2E tests. - assert.ok(true, 'IT-8 requirements validated by INV-1 and INV-2 controller disposal tests'); }); }); @@ -572,85 +593,12 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // Priority 2 Tests - High importance for environment switching suite('Priority 2: State Management (UT-2)', () => { test('Implementation verifies INV-9: cached state cleared before rebuild', () => { - // UT-2 requires verifying that rebuildController() clears cached state: - // 1. notebookControllers.delete() called before ensureKernelSelected() - // 2. notebookConnectionMetadata.delete() called before ensureKernelSelected() - // 3. Old server unregistered from provider - // - // THIS IS VALIDATED BY EXISTING TESTS AND IMPLEMENTATION: - // - // 1. "should clear cached controller and metadata" test (line 109) - // - Tests the cache clearing behavior during rebuild - // - Validates INV-9: Connection metadata cache cleared before creating new metadata - // - // 2. "should unregister old server handle" test (line 151) - // - Validates server cleanup during rebuild - // - Ensures old server is unregistered from provider - // - // THE ACTUAL IMPLEMENTATION at deepnoteKernelAutoSelector.node.ts:269-291: - // - // // Clear cached state - // this.notebookControllers.delete(notebookKey); - // this.notebookConnectionMetadata.delete(notebookKey); - // - // // Unregister old server - // const oldServerHandle = this.notebookServerHandles.get(notebookKey); - // if (oldServerHandle) { - // this.serverProvider.unregisterServer(oldServerHandle); - // this.notebookServerHandles.delete(notebookKey); - // } - // - // These operations happen BEFORE calling ensureKernelSelected() to create the new controller, - // ensuring clean state for the environment switch. - assert.ok(true, 'UT-2 is validated by existing tests and implementation (INV-9)'); }); }); suite('Priority 2: Server Concurrency (UT-7)', () => { test('Implementation verifies INV-8: concurrent startServer() calls are serialized', () => { - // UT-7 requires testing that concurrent startServer() calls for the same environment: - // 1. Second call waits for first to complete - // 2. Only one server process started - // 3. Both calls return same serverInfo - // - // THIS BEHAVIOR IS IMPLEMENTED IN deepnoteServerStarter.node.ts:82-91: - // - // // Wait for any pending operations on this environment to complete - // const pendingOp = this.pendingOperations.get(environmentId); - // if (pendingOp) { - // logger.info(`Waiting for pending operation on environment ${environmentId}...`); - // try { - // await pendingOp; - // } catch { - // // Ignore errors from previous operations - // } - // } - // - // And then tracks new operations at lines 103-114: - // - // // Start the operation and track it - // const operation = this.startServerForEnvironment(...); - // this.pendingOperations.set(environmentId, operation); - // - // try { - // const result = await operation; - // return result; - // } finally { - // // Remove from pending operations when done - // if (this.pendingOperations.get(environmentId) === operation) { - // this.pendingOperations.delete(environmentId); - // } - // } - // - // This ensures INV-8: Only one startServer() operation per environmentId can be in - // flight at a time. The second concurrent call will wait for the first to complete, - // then check if the server is already running (line 94-100) and return the existing - // serverInfo, preventing duplicate server processes and port conflicts. - // - // Creating a unit test for this would require complex async mocking and race condition - // simulation. The implementation's use of pendingOperations map provides the guarantee. - assert.ok(true, 'UT-7 is validated by implementation using pendingOperations map (INV-8)'); }); }); @@ -658,77 +606,10 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // Priority 2 Integration Tests suite('Priority 2: Integration Tests (IT-2, IT-6)', () => { test('IT-2: Switch while cells executing is handled by warning flow', () => { - // IT-2 requires: "Switch environment while cells are running" - // Verify: - // 1. Warning shown about executing cells - // 2. Switch completes - // 3. Running cell may fail (acceptable) - // 4. New cells execute on new environment - // - // THIS IS VALIDATED BY IMPLEMENTATION: - // - // 1. User warning in deepnoteEnvironmentsView.ts:542-561: - // - Checks kernel.pendingCells before switch - // - Shows warning dialog to user if cells executing - // - User can proceed or cancel - // - // 2. Logging in deepnoteKernelAutoSelector.node.ts:269-276: - // - Checks kernel.pendingCells during rebuildController - // - Logs warning if cells are executing - // - Proceeds with rebuild (non-blocking) - // - // The implementation allows switches during execution (with warnings) because: - // - Blocking would create a poor user experience - // - Running cells may fail, which is acceptable - // - New cells will use the new environment - // - Controller disposal order (INV-2) ensures no "disposed controller" error - // - // Full integration testing would require: - // - Real notebook with executing cells - // - Real kernel execution - // - Timing-sensitive test (start execution, then immediately switch) - // - Better suited for E2E tests - assert.ok(true, 'IT-2 is validated by warning implementation and INV-2'); }); test('IT-6: Server start failure during switch should show error to user', () => { - // IT-6 requires: "Environment switch fails due to server error" - // Verify: - // 1. Error shown to user - // 2. Notebook still usable (ideally on old environment A) - // 3. No controller leak - // 4. Can retry switch - // - // CURRENT IMPLEMENTATION BEHAVIOR: - // - // 1. If startServer() fails, the error propagates from ensureKernelSelectedWithConfiguration() - // (deepnoteKernelAutoSelector.node.ts:450-467) - // - // 2. The error is caught and shown to user in the UI layer - // - // 3. Controller handling in rebuildController() (lines 306-315): - // - Old controller is stored before rebuild - // - Old controller is NEVER disposed (even on success) - // - This means notebook can still use old controller for queued executions - // - // POTENTIAL IMPROVEMENT (noted in test plan): - // The test plan identifies this as a gap in "Known Gaps and Future Improvements": - // - "No atomic rollback: If switch fails mid-way, state may be inconsistent" - // - Recommended: "Implement rollback mechanism: Restore old controller if switch fails" - // - // Currently, if server start fails: - // - Old controller is NOT disposed (good - notebook still has a controller) - // - Cached state WAS cleared (lines 279-282) - // - So getSelected() may not return the old controller from cache - // - // RECOMMENDED FUTURE IMPROVEMENT: - // Wrap ensureKernelSelected() in try-catch in rebuildController(): - // - On success: dispose old controller as usual - // - On failure: restore cached state for old controller - // - // For now, this test documents the current behavior and the known limitation. - assert.ok( true, 'IT-6 behavior is partially implemented - error shown, but rollback not implemented (known gap)' @@ -739,10 +620,6 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { // REAL TDD Tests - These should FAIL if bugs exist suite('Bug Detection: Kernel Selection', () => { test('BUG-1: Should prefer environment-specific kernel over .env kernel', () => { - // REAL TEST: This will FAIL if the wrong kernel is selected - // - // The selectKernelSpec method is now extracted and testable! - const envId = 'env123'; const kernelSpecs: IJupyterKernelSpec[] = [ createMockKernelSpec('.env', '.env Python', 'python'), @@ -752,7 +629,6 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { const selected = selector.selectKernelSpec(kernelSpecs, envId); - // CRITICAL ASSERTION: Should select environment-specific kernel, NOT .env assert.strictEqual( selected?.name, `deepnote-${envId}`, @@ -761,21 +637,14 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { }); test('BUG-1b: Current implementation falls back to Python kernel (documents expected behavior)', () => { - // This test documents that the current implementation DOES have fallback logic - // - // EXPECTED BEHAVIOR (current): Fall back to generic Python kernel when env-specific kernel not found - // This is a design decision - we don't want to block users if the environment-specific kernel isn't ready yet - const envId = 'env123'; const kernelSpecs: IJupyterKernelSpec[] = [ createMockKernelSpec('.env', '.env Python', 'python'), createMockKernelSpec('python3', 'Python 3', 'python') ]; - // Should fall back to a Python kernel (this is the current behavior) const selected = selector.selectKernelSpec(kernelSpecs, envId); - // Should have selected a fallback kernel (either .env or python3) assert.ok(selected, 'Should select a fallback kernel'); assert.strictEqual(selected.language, 'python', 'Fallback should be a Python kernel'); }); @@ -793,14 +662,12 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { }); test('Kernel selection: Should fall back to python3 when env kernel missing', () => { - // Documents current fallback behavior - falls back to python3 when env kernel missing const envId = 'my-env'; const kernelSpecs: IJupyterKernelSpec[] = [ createMockKernelSpec('python3', 'Python 3', 'python'), createMockKernelSpec('javascript', 'JavaScript', 'javascript') ]; - // Should fall back to python3 (current behavior) const selected = selector.selectKernelSpec(kernelSpecs, envId); assert.strictEqual(selected.name, 'python3', 'Should fall back to python3'); @@ -809,103 +676,8 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { suite('Bug Detection: Controller Disposal', () => { test('BUG-2: Old controller is NOT disposed to prevent queued execution errors', async () => { - // This test documents the fix for the DISPOSED error - // - // SCENARIO: User switches environments and has queued cell executions - // - // THE FIX: We do NOT dispose the old controller at all (lines 306-315) - // - Line 281: notebookControllers.delete(notebookKey) removes controller from cache - // - Lines 306-315: Old controller is left alive (NOT disposed) - // - VS Code may have queued cell executions that reference the old controller - // - Those executions will complete successfully using the old controller - // - New executions will use the new controller (it's now preferred) - // - The old controller will be garbage collected when no longer referenced - // - // This prevents the "notebook controller is DISPOSED" error that happened when: - // 1. User queues cell execution (references old controller) - // 2. User switches environments (creates new controller, disposes old one) - // 3. Queued execution tries to run (BOOM - old controller is disposed) - assert.ok(true, 'Old controller is never disposed - prevents DISPOSED errors for queued executions'); }); - - test.skip('BUG-2b: Old controller should only be disposed AFTER new controller is in cache', async () => { - // This test is skipped because _testOnly_setController method doesn't exist in the implementation - // REAL TEST: This will FAIL if disposal happens too early - // - // Setup: Create a scenario where we have an old controller and create a new one - const baseFileUri = mockNotebook.uri.with({ query: '', fragment: '' }); - // const notebookKey = baseFileUri.fsPath; - const newEnv = createMockEnvironment('env-new', 'New Environment', true); - - // Track call order - const callOrder: string[] = []; - - // Setup old controller that tracks when dispose() is called - const oldController = mock(); - when(oldController.id).thenReturn('deepnote-config-kernel-env-old'); - when(oldController.controller).thenReturn({} as any); - when(oldController.dispose()).thenCall(() => { - callOrder.push('OLD_CONTROLLER_DISPOSED'); - return undefined; - }); - - // CRITICAL: Use test helper to set up initial controller in cache - // This simulates the state where a controller already exists before environment switch - // selector._testOnly_setController(notebookKey, instance(oldController)); - - // Setup new controller - const newController = mock(); - when(newController.id).thenReturn('deepnote-config-kernel-env-new'); - when(newController.controller).thenReturn({} as any); - - // Setup mocks - when(mockNotebookEnvironmentMapper.getEnvironmentForNotebook(baseFileUri)).thenReturn('env-new'); - when(mockEnvironmentManager.getEnvironment('env-new')).thenReturn(newEnv); - when(mockPythonExtensionChecker.isPythonExtensionInstalled).thenReturn(true); - - // Mock controller registration to track when new controller is added - when(mockControllerRegistration.addOrUpdate(anything(), anything())).thenCall(() => { - callOrder.push('NEW_CONTROLLER_ADDED_TO_REGISTRATION'); - return [instance(newController)]; - }); - - // CRITICAL TEST: We need to verify that within rebuildController: - // 1. ensureKernelSelected creates and caches new controller (NEW_CONTROLLER_ADDED_TO_REGISTRATION) - // 2. Only THEN is old controller disposed (OLD_CONTROLLER_DISPOSED) - // - // If OLD_CONTROLLER_DISPOSED happens before NEW_CONTROLLER_ADDED_TO_REGISTRATION, - // then there's a window where no valid controller exists! - - await selector.rebuildController(mockNotebook, mockProgress, instance(mockCancellationToken)); - - // ASSERTION: If implementation is correct, call order should be: - // 1. NEW_CONTROLLER_ADDED_TO_REGISTRATION (from ensureKernelSelected) - // 2. OLD_CONTROLLER_DISPOSED (from rebuildController after new controller is ready) - // - // This test will FAIL if: - // - dispose() is called before new controller is registered - // - or if dispose() is never called - - if (callOrder.length > 0) { - const newControllerIndex = callOrder.indexOf('NEW_CONTROLLER_ADDED_TO_REGISTRATION'); - const oldDisposeIndex = callOrder.indexOf('OLD_CONTROLLER_DISPOSED'); - - if (newControllerIndex !== -1 && oldDisposeIndex !== -1) { - assert.ok( - newControllerIndex < oldDisposeIndex, - `BUG DETECTED: Old controller disposed before new controller was registered! Order: ${callOrder.join( - ' -> ' - )}` - ); - } else { - // This is OK - test might not have reached disposal due to mocking limitations - assert.ok(true, 'Test did not reach disposal phase due to mocking complexity'); - } - } else { - assert.ok(true, 'Test did not capture call order due to mocking complexity'); - } - }); }); suite('Requirements Optimization', () => { @@ -938,7 +710,6 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { const requirements = ['pandas', 123, 'numpy', null, 'scipy', undefined]; const result = computeRequirementsHash(requirements); - // Should only include string entries assert.strictEqual(result, 'numpy|pandas|scipy'); }); @@ -967,7 +738,6 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { const requirements = ['pandas', 'numpy', 'pandas', 'scipy', 'numpy']; const result = computeRequirementsHash(requirements); - // Should have each requirement only once assert.strictEqual(result, 'numpy|pandas|scipy'); }); @@ -1001,18 +771,14 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { suite('getExistingRequirementsHash', () => { test('parsing logic correctness', () => { - // Test the parsing logic directly by calling computeRequirementsHash - // with a parsed file-like array (mimics what getExistingRequirementsHash does) const fileLines = ['# This is a comment', 'pandas', '', ' numpy ', 'scipy', '# Another comment']; - // Filter out comments and empty lines (same logic as getExistingRequirementsHash) const requirements = fileLines .map((line) => line.trim()) .filter((line) => line.length > 0 && !line.startsWith('#')); const hash = computeRequirementsHash(requirements); - // Should have filtered and sorted correctly assert.strictEqual(hash, 'numpy|pandas|scipy'); }); }); diff --git a/src/notebooks/deepnote/deepnoteKernelStatusIndicator.node.ts b/src/notebooks/deepnote/deepnoteKernelStatusIndicator.node.ts index af1586f136..13ee63363b 100644 --- a/src/notebooks/deepnote/deepnoteKernelStatusIndicator.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelStatusIndicator.node.ts @@ -18,7 +18,7 @@ import { logger } from '../../platform/logging'; import { DEEPNOTE_NOTEBOOK_TYPE, IDeepnoteEnvironmentManager, - IDeepnoteNotebookEnvironmentMapper + IDeepnoteProjectEnvironmentMapper } from '../../kernels/deepnote/types'; import { IControllerRegistration, IVSCodeNotebookController } from '../controllers/types'; @@ -81,8 +81,8 @@ export class DeepnoteKernelStatusIndicator constructor( @inject(IDisposableRegistry) disposableRegistry: IDisposableRegistry, @inject(IControllerRegistration) private readonly controllerRegistration: IControllerRegistration, - @inject(IDeepnoteNotebookEnvironmentMapper) - private readonly environmentMapper: IDeepnoteNotebookEnvironmentMapper, + @inject(IDeepnoteProjectEnvironmentMapper) + private readonly environmentMapper: IDeepnoteProjectEnvironmentMapper, @inject(IDeepnoteEnvironmentManager) private readonly environmentManager: IDeepnoteEnvironmentManager ) { disposableRegistry.push(this); @@ -283,8 +283,8 @@ export class DeepnoteKernelStatusIndicator return undefined; } - const baseUri = notebook.uri.with({ query: '', fragment: '' }); - const environmentId = this.environmentMapper.getEnvironmentForNotebook(baseUri); + const projectId = notebook.metadata?.deepnoteProjectId as string | undefined; + const environmentId = projectId ? this.environmentMapper.getEnvironmentForProject(projectId) : undefined; const environmentName = environmentId ? this.environmentManager.getEnvironment(environmentId)?.name : undefined; const connection = controller.connection; diff --git a/src/notebooks/deepnote/deepnoteKernelStatusIndicator.ts b/src/notebooks/deepnote/deepnoteKernelStatusIndicator.ts index 6bc389b9dd..e432163d17 100644 --- a/src/notebooks/deepnote/deepnoteKernelStatusIndicator.ts +++ b/src/notebooks/deepnote/deepnoteKernelStatusIndicator.ts @@ -19,7 +19,7 @@ import { logger } from '../../platform/logging'; import { DEEPNOTE_NOTEBOOK_TYPE, IDeepnoteEnvironmentManager, - IDeepnoteNotebookEnvironmentMapper + IDeepnoteProjectEnvironmentMapper } from '../../kernels/deepnote/types'; import { IControllerRegistration, IVSCodeNotebookController } from '../controllers/types'; import { Commands } from '../../platform/common/constants'; @@ -85,8 +85,8 @@ export class DeepnoteKernelStatusIndicator constructor( @inject(IDisposableRegistry) disposableRegistry: IDisposableRegistry, @inject(IControllerRegistration) private readonly controllerRegistration: IControllerRegistration, - @inject(IDeepnoteNotebookEnvironmentMapper) - private readonly environmentMapper: IDeepnoteNotebookEnvironmentMapper, + @inject(IDeepnoteProjectEnvironmentMapper) + private readonly environmentMapper: IDeepnoteProjectEnvironmentMapper, @inject(IDeepnoteEnvironmentManager) private readonly environmentManager: IDeepnoteEnvironmentManager ) { disposableRegistry.push(this); @@ -290,8 +290,8 @@ export class DeepnoteKernelStatusIndicator return undefined; } - const baseUri = notebook.uri.with({ query: '', fragment: '' }); - const environmentId = this.environmentMapper.getEnvironmentForNotebook(baseUri); + const projectId = notebook.metadata?.deepnoteProjectId as string | undefined; + const environmentId = projectId ? this.environmentMapper.getEnvironmentForProject(projectId) : undefined; const environmentName = environmentId ? this.environmentManager.getEnvironment(environmentId)?.name : undefined; const connection = controller.connection; diff --git a/src/notebooks/deepnote/snapshots/environmentCapture.node.ts b/src/notebooks/deepnote/snapshots/environmentCapture.node.ts index dcd3c923cb..1f26ea3739 100644 --- a/src/notebooks/deepnote/snapshots/environmentCapture.node.ts +++ b/src/notebooks/deepnote/snapshots/environmentCapture.node.ts @@ -10,9 +10,10 @@ import { computeHash } from '../../../platform/common/crypto'; import { raceTimeout } from '../../../platform/common/utils/async'; import { logger } from '../../../platform/logging'; import { parsePipFreezeFile } from './pipFileParser'; -import { IDeepnoteEnvironmentManager, IDeepnoteNotebookEnvironmentMapper } from '../../../kernels/deepnote/types'; +import { IDeepnoteEnvironmentManager, IDeepnoteProjectEnvironmentMapper } from '../../../kernels/deepnote/types'; import { Uri } from 'vscode'; import { DeepnoteEnvironment } from '../../../kernels/deepnote/environments/deepnoteEnvironment'; +import { resolveProjectIdForFile } from '../../../platform/deepnote/deepnoteProjectIdResolver'; import * as path from '../../../platform/vscode-path/path'; const captureTimeoutInMilliseconds = 5_000; @@ -34,13 +35,13 @@ type PythonEnvironmentType = 'uv' | 'conda' | 'venv' | 'poetry' | 'system'; @injectable() export class EnvironmentCapture implements IEnvironmentCapture { constructor( - @inject(IDeepnoteNotebookEnvironmentMapper) - private readonly environmentMapper: IDeepnoteNotebookEnvironmentMapper, + @inject(IDeepnoteProjectEnvironmentMapper) + private readonly environmentMapper: IDeepnoteProjectEnvironmentMapper, @inject(IDeepnoteEnvironmentManager) private readonly environmentManager: IDeepnoteEnvironmentManager ) {} async captureEnvironment(notebookUri: Uri): Promise { - const deepnoteEnvironment = this.getEnvironmentForNotebook(notebookUri); + const deepnoteEnvironment = await this.resolveEnvironmentForNotebook(notebookUri); if (!deepnoteEnvironment) { logger.warn('[EnvironmentCapture] No Deepnote environment found for the given notebook'); @@ -184,8 +185,17 @@ export class EnvironmentCapture implements IEnvironmentCapture { return raceTimeout(captureTimeoutInMilliseconds, undefined, getVersion()); } - private getEnvironmentForNotebook(notebookUri: Uri): DeepnoteEnvironment | undefined { - const environmentId = this.environmentMapper.getEnvironmentForNotebook(notebookUri); + private async resolveEnvironmentForNotebook(notebookUri: Uri): Promise { + const baseFileUri = notebookUri.with({ query: '', fragment: '' }); + const projectId = await resolveProjectIdForFile(baseFileUri); + + if (!projectId) { + return undefined; + } + + await this.environmentMapper.waitForInitialization(); + + const environmentId = this.environmentMapper.getEnvironmentForProject(projectId); if (!environmentId) { return undefined; diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index f4ea82b469..f671334941 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -67,7 +67,7 @@ import { IDeepnoteKernelAutoSelector, IDeepnoteServerProvider, IDeepnoteEnvironmentManager, - IDeepnoteNotebookEnvironmentMapper, + IDeepnoteProjectEnvironmentMapper, IDeepnoteLspClientManager } from '../kernels/deepnote/types'; import { DeepnoteAgentSkillsManager } from '../kernels/deepnote/deepnoteAgentSkillsManager.node'; @@ -83,7 +83,7 @@ import { DeepnoteEnvironmentStorage } from '../kernels/deepnote/environments/dee import { DeepnoteEnvironmentsView } from '../kernels/deepnote/environments/deepnoteEnvironmentsView.node'; import { DeepnoteEnvironmentsActivationService } from '../kernels/deepnote/environments/deepnoteEnvironmentsActivationService'; import { DeepnoteExtensionSidecarWriter } from '../kernels/deepnote/environments/deepnoteExtensionSidecarWriter.node'; -import { DeepnoteNotebookEnvironmentMapper } from '../kernels/deepnote/environments/deepnoteNotebookEnvironmentMapper.node'; +import { DeepnoteProjectEnvironmentMapper } from '../kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.node'; import { DeepnoteNotebookCommandListener } from './deepnote/deepnoteNotebookCommandListener'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; @@ -255,10 +255,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea DeepnoteEnvironmentsActivationService ); - // Deepnote configuration selection - serviceManager.addSingleton( - IDeepnoteNotebookEnvironmentMapper, - DeepnoteNotebookEnvironmentMapper + // Deepnote configuration selection (project-id keyed) + serviceManager.addSingleton( + IDeepnoteProjectEnvironmentMapper, + DeepnoteProjectEnvironmentMapper ); // Sidecar file writer (exposes env mappings for external tools) diff --git a/src/platform/deepnote/deepnoteProjectIdResolver.ts b/src/platform/deepnote/deepnoteProjectIdResolver.ts new file mode 100644 index 0000000000..ad78a1b8ff --- /dev/null +++ b/src/platform/deepnote/deepnoteProjectIdResolver.ts @@ -0,0 +1,36 @@ +import { NotebookDocument, Uri } from 'vscode'; + +import { logger } from '../logging'; +import { readDeepnoteProjectFile } from '../../notebooks/deepnote/deepnoteProjectUtils'; + +/** + * Resolve the Deepnote project id for an open notebook document. + * Prefers the pre-populated `deepnoteProjectId` metadata set by the serializer; + * falls back to reading the YAML on disk if metadata is missing. + */ +export async function resolveProjectIdForNotebook(notebook: NotebookDocument): Promise { + const metadataProjectId = notebook.metadata?.deepnoteProjectId as string | undefined; + if (metadataProjectId) { + return metadataProjectId; + } + + const baseFileUri = notebook.uri.with({ query: '', fragment: '' }); + + return resolveProjectIdForFile(baseFileUri); +} + +/** + * Resolve the Deepnote project id for a `.deepnote` file on disk by parsing the + * YAML content. Swallows I/O/parse errors and returns `undefined` on failure. + */ +export async function resolveProjectIdForFile(fileUri: Uri): Promise { + try { + const parsed = await readDeepnoteProjectFile(fileUri); + + return parsed?.project?.id; + } catch (error) { + logger.warn(`Failed to resolve Deepnote project id for ${fileUri.fsPath}`, error); + + return undefined; + } +} diff --git a/src/platform/deepnote/deepnoteServerUtils.node.ts b/src/platform/deepnote/deepnoteServerUtils.node.ts index e8fea3659e..ff4697b016 100644 --- a/src/platform/deepnote/deepnoteServerUtils.node.ts +++ b/src/platform/deepnote/deepnoteServerUtils.node.ts @@ -1,5 +1,3 @@ -import { Uri } from 'vscode'; - -export function createDeepnoteServerConfigHandle(environmentId: string, deepnoteFileUri: Uri): string { - return `deepnote-config-server-${environmentId}-${deepnoteFileUri.fsPath}`; +export function createDeepnoteServerConfigHandle(environmentId: string, projectId: string): string { + return `deepnote-config-server-${environmentId}-${projectId}`; } From b415441189b37d389962a9228335d6f1a798fac8 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 23 Apr 2026 08:30:28 +0000 Subject: [PATCH 44/47] feat(deepnote): Improve snapshot handling with output settling logic - Introduced mechanisms to manage pending snapshot saves based on output activity, preventing premature saves during ongoing kernel output delivery. - Added new constants for output settling periods and maximum wait times to enhance snapshot reliability. - Updated `SnapshotService` to clear pending saves on notebook document changes and to handle output-bearing changes effectively. - Enhanced unit tests to validate the new fallback logic for snapshots, ensuring correct behavior when latest snapshots lack outputs. --- .../deepnote/snapshots/snapshotService.ts | 263 ++++++++++++++++-- .../snapshots/snapshotService.unit.test.ts | 234 +++++++++++++++- 2 files changed, 479 insertions(+), 18 deletions(-) diff --git a/src/notebooks/deepnote/snapshots/snapshotService.ts b/src/notebooks/deepnote/snapshots/snapshotService.ts index c20f5948f3..f2b66df6f3 100644 --- a/src/notebooks/deepnote/snapshots/snapshotService.ts +++ b/src/notebooks/deepnote/snapshots/snapshotService.ts @@ -10,7 +10,17 @@ import { } from '@deepnote/blocks'; import fastDeepEqual from 'fast-deep-equal'; import { inject, injectable, optional } from 'inversify'; -import { Disposable, FileType, NotebookCell, NotebookCellKind, RelativePattern, Uri, window, workspace } from 'vscode'; +import { + Disposable, + FileType, + NotebookCell, + NotebookCellKind, + NotebookDocumentChangeEvent, + RelativePattern, + Uri, + window, + workspace +} from 'vscode'; import { Utils } from 'vscode-uri'; import { IExtensionSyncActivationService } from '../../../platform/activation/types'; @@ -101,6 +111,31 @@ interface NotebookExecutionState { /** How long a written URI is considered "recent" and suppressed from file-change processing. */ const recentWriteExpirationMs = 2000; +/** + * After the cell execution queue completes, the kernel can still be delivering iopub + * messages (stream output, display_data, etc.) that land on `cell.outputs` AFTER + * `onDidCompleteQueueExecution` fires. Instead of reading `cell.outputs` immediately, + * we arm a timer and reset it every time `workspace.onDidChangeNotebookDocument` reports + * an output-bearing cell change. The save only runs once we observe this many ms + * without any output activity. + */ +const outputSettleQuietPeriodMs = 150; + +/** + * Hard cap on how long a pending snapshot save will be deferred by output activity. + * Bounds a misbehaving kernel (e.g. a background thread continuously printing) so it + * can't starve the snapshot pipeline. + */ +const outputSettleMaxWaitMs = 2000; + +/** Pending, output-settled snapshot save for a single notebook. */ +interface PendingSnapshotSave { + /** Wall-clock time when the queue-complete event first armed this pending save. */ + armedAt: number; + /** Active timer that will call `flushPendingSnapshotSave` when it fires. */ + timer: ReturnType; +} + class TimeoutError extends Error { constructor(message: string) { super(message); @@ -126,6 +161,11 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync private readonly converter = new DeepnoteDataConverter(); private readonly executionStates = new Map(); private readonly fileWrittenCallbacks: ((uri: Uri) => void)[] = []; + /** + * Notebooks with a snapshot save scheduled but not yet executed. Keyed by notebook URI. + * Resettable timer debounces the save against output-bearing notebook-document changes. + */ + private readonly pendingSnapshotSaves = new Map(); private readonly recentlyWrittenTimers = new Map>(); private readonly recentlyWrittenUris = new Set(); @@ -144,17 +184,29 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync clearTimeout(timer); } this.recentlyWrittenTimers.clear(); + + for (const pending of this.pendingSnapshotSaves.values()) { + clearTimeout(pending.timer); + } + this.pendingSnapshotSaves.clear(); } }); workspace.onDidCloseNotebookDocument( (notebook) => { - this.clearExecutionState(notebook.uri.toString()); + const uri = notebook.uri.toString(); + this.clearExecutionState(uri); + this.cancelPendingSnapshotSave(uri); }, this, this.disposables ); + // Persistent listener so we can detect the kernel still delivering output AFTER + // onDidCompleteQueueExecution has fired. When a save is pending for this notebook, + // each output-bearing change defers the save by `outputSettleQuietPeriodMs`. + workspace.onDidChangeNotebookDocument((event) => this.onNotebookDocumentChanged(event), this, this.disposables); + notebookCellExecutions.onDidChangeNotebookCellExecutionState( (e) => { logger.debug(`[Snapshot] Cell execution state changed: ${e.state} for cell ${e.cell.metadata?.id}`); @@ -166,7 +218,7 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync notebookCellExecutions.onDidCompleteQueueExecution( async (e) => { - logger.debug(`[Snapshot] Queue execution complete for ${e.notebookUri}`); + logger.debug(`[Snapshot] Queue execution complete for ${e.notebookUri}, arming settled-output save`); await this.onExecutionComplete(e.notebookUri); }, this, @@ -417,27 +469,47 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync ? `**/snapshots/*_${projectId}_${notebookId}_latest.snapshot.deepnote` : `**/snapshots/*_${projectId}_latest.snapshot.deepnote`; + let latestOutputs: Map | undefined; + let latestUri: Uri | undefined; + for (const folder of workspaceFolders) { logger.debug(`[Snapshot] Searching for latest snapshot with glob: ${latestGlob} in ${folder.uri.path}`); const latestPattern = new RelativePattern(folder, latestGlob); const latestFiles = await workspace.findFiles(latestPattern, null, 1); if (latestFiles.length > 0) { - logger.debug(`[Snapshot] Found latest snapshot: ${latestFiles[0].path}`); + latestUri = latestFiles[0]; + logger.debug(`[Snapshot] Found latest snapshot: ${latestUri.path}`); try { - return await this.parseSnapshotFile(latestFiles[0]); + latestOutputs = await this.parseSnapshotFile(latestUri); } catch (error) { - logger.error(`[Snapshot] Failed to parse snapshot: ${Utils.basename(latestFiles[0])}`, error); + logger.error(`[Snapshot] Failed to parse snapshot: ${Utils.basename(latestUri)}`, error); - await window.showErrorMessage(`Failed to read latest snapshot: ${Utils.basename(latestFiles[0])}`); + await window.showErrorMessage(`Failed to read latest snapshot: ${Utils.basename(latestUri)}`); return; } + + // If `latest` carries any outputs, it's authoritative. If it exists but is + // empty, fall through to the timestamped fallback in case `latest` was + // overwritten by a partial save that lost the outputs (see save race + // discussion in snapshotService). + if (this.hasAnyOutputs(latestOutputs)) { + return latestOutputs; + } + + logger.debug( + `[Snapshot] Latest snapshot has no outputs, looking for a timestamped snapshot with outputs` + ); + + break; } } - logger.debug(`[Snapshot] No latest snapshot found, looking for timestamped files`); + if (!latestUri) { + logger.debug(`[Snapshot] No latest snapshot found, looking for timestamped files`); + } // 2. Find timestamped snapshots across all workspace folders const timestampedGlob = notebookId @@ -465,17 +537,45 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync if (sortedFiles.length === 0) { logger.debug(`[Snapshot] No timestamped snapshots found`); - return; + // No timestamped fallback. If we already parsed an empty `latest`, return it + // so callers can still observe that the snapshot file existed. + return latestOutputs; } - const newestFile = sortedFiles[0]; + // Walk timestamped files newest-first and use the first one that contains outputs. + for (const candidate of sortedFiles) { + try { + const outputs = await this.parseSnapshotFile(candidate); + + if (this.hasAnyOutputs(outputs)) { + if (latestUri) { + logger.warn( + `[Snapshot] Falling back to timestamped snapshot ${Utils.basename(candidate)} ` + + `because ${Utils.basename(latestUri)} contained no outputs` + ); + } else { + logger.debug(`[Snapshot] Using timestamped snapshot: ${Utils.basename(candidate)}`); + } + + return outputs; + } + } catch (error) { + logger.error(`[Snapshot] Failed to parse snapshot: ${Utils.basename(candidate)}`, error); + // Try the next-most-recent timestamped file rather than bailing out. + } + } - logger.debug(`[Snapshot] Using timestamped snapshot: ${Utils.basename(newestFile)}`); + // None of the timestamped files contained outputs either; fall back to whatever + // we parsed from `latest` (possibly an empty map) or to the newest timestamped + // file's parse result so callers always get a stable answer. + if (latestOutputs) { + return latestOutputs; + } try { - return await this.parseSnapshotFile(newestFile); + return await this.parseSnapshotFile(sortedFiles[0]); } catch (error) { - logger.error(`[Snapshot] Failed to parse snapshot: ${Utils.basename(newestFile)}`, error); + logger.error(`[Snapshot] Failed to parse snapshot: ${Utils.basename(sortedFiles[0])}`, error); return; } @@ -633,6 +733,26 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync return state; } + /** + * Returns true if the parsed outputs map contains at least one block with at least + * one real output entry. A snapshot that maps a block id to an empty array is treated + * as "no outputs" because it offers nothing to merge into the cells and is the + * signature of a save race that captured the cell before its outputs landed. + */ + private hasAnyOutputs(outputs: Map | undefined): boolean { + if (!outputs) { + return false; + } + + for (const blockOutputs of outputs.values()) { + if (blockOutputs && blockOutputs.length > 0) { + return true; + } + } + + return false; + } + private handleCellExecutionStateChange(cell: NotebookCell, state: NotebookCellExecutionState): void { const notebookUri = cell.notebook.uri.toString(); const cellId = cell.metadata?.id as string | undefined; @@ -644,6 +764,19 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync if (state === NotebookCellExecutionState.Executing) { const startTime = Date.now(); + // A new execution is starting for this notebook. Invalidate any pending + // settled-output save from a previous run — the outputs it would capture + // are about to change, and the original armedAt would otherwise keep + // counting toward the max-wait cap across runs, potentially firing the + // save mid-execution with stale cell.outputs. The next + // onDidCompleteQueueExecution will arm a fresh save with a fresh armedAt. + if (this.pendingSnapshotSaves.has(notebookUri)) { + logger.debug( + `[Snapshot] Cell ${cellId} starting execution; cancelling pending save for ${notebookUri}` + ); + this.cancelPendingSnapshotSave(notebookUri); + } + this.recordCellExecutionStart(notebookUri, cellId, startTime); } else if (state === NotebookCellExecutionState.Idle) { const endTime = Date.now(); @@ -727,11 +860,96 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync private async onExecutionComplete(notebookUri: string): Promise { logger.debug(`[Snapshot] onExecutionComplete called for ${notebookUri}`); - // Wait for any pending cell state change events to be processed. - // This is needed because the queue completion event can fire before the - // last cell's Idle state change event has been processed (race condition). + // Wait for any pending cell state change events to be processed. The queue + // completion event can fire before the last cell's Idle state change event has. await this.waitForPendingCellStateChanges(notebookUri, 100); + // Don't write the snapshot yet. The kernel can still be delivering output via + // async NotebookCellExecution.appendOutput calls that land AFTER queue completion. + // Instead, ARM a save timer here; `onNotebookDocumentChanged` will keep resetting + // it as long as output-bearing notebook changes keep arriving. Once they stop, + // the timer fires and `flushPendingSnapshotSave` runs `performSnapshotSave`. + this.schedulePendingSnapshotSave(notebookUri); + } + + private schedulePendingSnapshotSave(notebookUri: string): void { + const existing = this.pendingSnapshotSaves.get(notebookUri); + const armedAt = existing?.armedAt ?? Date.now(); + + if (existing) { + clearTimeout(existing.timer); + } + + const delay = this.computeSettleDelay(armedAt); + const timer = setTimeout(() => this.flushPendingSnapshotSave(notebookUri), delay); + + this.pendingSnapshotSaves.set(notebookUri, { armedAt, timer }); + } + + private computeSettleDelay(armedAt: number): number { + const elapsed = Date.now() - armedAt; + const remainingBudget = Math.max(0, outputSettleMaxWaitMs - elapsed); + + return Math.min(outputSettleQuietPeriodMs, remainingBudget); + } + + private cancelPendingSnapshotSave(notebookUri: string): void { + const existing = this.pendingSnapshotSaves.get(notebookUri); + + if (existing) { + clearTimeout(existing.timer); + this.pendingSnapshotSaves.delete(notebookUri); + } + } + + private onNotebookDocumentChanged(event: NotebookDocumentChangeEvent): void { + const notebookUri = event.notebook.uri.toString(); + const pending = this.pendingSnapshotSaves.get(notebookUri); + + // We only care about output activity while a save is pending — no pending save, + // no work to do. + if (!pending) { + return; + } + + const hasOutputOrMetadataChange = event.cellChanges.some( + (change) => change.outputs !== undefined || change.metadata !== undefined + ); + + if (!hasOutputOrMetadataChange) { + return; + } + + // If we've already spent the whole max-wait budget, stop deferring — let the + // scheduled timer fire on its existing schedule. + if (Date.now() - pending.armedAt >= outputSettleMaxWaitMs) { + return; + } + + // Output is still landing. Reset the quiet timer so the save waits another + // `outputSettleQuietPeriodMs` before running. + clearTimeout(pending.timer); + const delay = this.computeSettleDelay(pending.armedAt); + pending.timer = setTimeout(() => this.flushPendingSnapshotSave(notebookUri), delay); + } + + private async flushPendingSnapshotSave(notebookUri: string): Promise { + const pending = this.pendingSnapshotSaves.get(notebookUri); + this.pendingSnapshotSaves.delete(notebookUri); + + if (pending) { + const settledIn = Date.now() - pending.armedAt; + logger.debug(`[Snapshot] Output settled in ${settledIn}ms, running save for ${notebookUri}`); + } + + try { + await this.performSnapshotSave(notebookUri); + } catch (error) { + logger.error(`[Snapshot] Failed to save snapshot for ${notebookUri}`, error); + } + } + + private async performSnapshotSave(notebookUri: string): Promise { if (!this.isSnapshotsEnabled()) { logger.debug(`[Snapshot] Snapshots not enabled, skipping`); @@ -793,6 +1011,19 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync metadata: cell.metadata, outputs: [...cell.outputs] })); + + // Diagnostic: log what we're about to persist so future staleness bugs are + // traceable in the user's output log without additional instrumentation. + for (let i = 0; i < cellData.length; i++) { + const cell = cellData[i]; + const blockId = (cell.metadata?.id as string | undefined) ?? '(no-id)'; + + logger.debug( + `[Snapshot] Pre-save cell[${i}] blockId=${blockId} outputs=${cell.outputs.length} ` + + `contentLen=${cell.value.length}` + ); + } + const blocks = this.converter.convertCellsToBlocks(cellData); const snapshotProject = structuredClone(originalProject) as DeepnoteFile; diff --git a/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts b/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts index 36666343a7..64fede3e2c 100644 --- a/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts +++ b/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts @@ -9,6 +9,7 @@ import { IEnvironmentCapture } from './environmentCapture.node'; import { SnapshotService } from './snapshotService'; import type { DeepnoteOutput } from '../../../platform/deepnote/deepnoteTypes'; import { IDisposableRegistry } from '../../../platform/common/types'; +import { NotebookCellExecutionState } from '../../../platform/notebooks/cellExecutionStateService'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; suite('SnapshotService', () => { @@ -684,6 +685,207 @@ project: assert.isDefined(result); assert.strictEqual(result!.size, 0); }); + + test('should fall back to timestamped snapshot when latest exists but has no outputs', async () => { + const workspaceFolder: WorkspaceFolder = { + uri: Uri.file('/workspace'), + name: 'workspace', + index: 0 + }; + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder]); + + const latestUri = Uri.file( + '/workspace/snapshots/project_test-project-id-123_notebook-1_latest.snapshot.deepnote' + ); + const timestampedNewerUri = Uri.file( + '/workspace/snapshots/project_test-project-id-123_notebook-1_2025-01-02T10-00-00.snapshot.deepnote' + ); + const timestampedOlderUri = Uri.file( + '/workspace/snapshots/project_test-project-id-123_notebook-1_2025-01-01T10-00-00.snapshot.deepnote' + ); + + // Mock findFiles: first call resolves latest, second resolves all matching timestamped + latest. + let findFilesCall = 0; + when(mockedVSCodeNamespaces.workspace.findFiles(anything(), anything(), anything())).thenCall(() => { + findFilesCall++; + + if (findFilesCall === 1) { + return Promise.resolve([latestUri]); + } + + return Promise.resolve([latestUri, timestampedNewerUri, timestampedOlderUri]); + }); + + // The latest yaml has no outputs (matches the save-bug shape). + const latestYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-02T00:00:00Z' +project: + id: test-project-id-123 + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + blockGroup: group-1 + type: code + sortingKey: 'a0' + content: print(1) +`; + // The newer timestamped snapshot has real outputs. + const newerTimestampedYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-02T00:00:00Z' +project: + id: test-project-id-123 + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + blockGroup: group-1 + type: code + sortingKey: 'a0' + content: print(1) + outputs: + - output_type: stream + name: stdout + text: 'from newer timestamped' +`; + const olderTimestampedYaml = newerTimestampedYaml.replace('from newer timestamped', 'from older'); + + const mockFs = mock(); + when(mockFs.readFile(anything())).thenCall((uri: Uri) => { + if (uri.toString() === latestUri.toString()) { + return Promise.resolve(new TextEncoder().encode(latestYaml) as any); + } + + if (uri.toString() === timestampedNewerUri.toString()) { + return Promise.resolve(new TextEncoder().encode(newerTimestampedYaml) as any); + } + + return Promise.resolve(new TextEncoder().encode(olderTimestampedYaml) as any); + }); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + const result = await service.readSnapshot(projectId, 'notebook-1'); + + assert.isDefined(result); + assert.strictEqual(result!.size, 1); + assert.deepStrictEqual(result!.get('block-1'), [ + { output_type: 'stream', name: 'stdout', text: 'from newer timestamped' } + ]); + }); + + test('should prefer latest snapshot when it already has outputs', async () => { + const workspaceFolder: WorkspaceFolder = { + uri: Uri.file('/workspace'), + name: 'workspace', + index: 0 + }; + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder]); + + const latestUri = Uri.file( + '/workspace/snapshots/project_test-project-id-123_notebook-1_latest.snapshot.deepnote' + ); + + when(mockedVSCodeNamespaces.workspace.findFiles(anything(), anything(), anything())).thenResolve([ + latestUri + ] as any); + + const latestYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-02T00:00:00Z' +project: + id: test-project-id-123 + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + blockGroup: group-1 + type: code + sortingKey: 'a0' + content: print(1) + outputs: + - output_type: stream + name: stdout + text: 'from latest' +`; + const mockFs = mock(); + when(mockFs.readFile(anything())).thenResolve(new TextEncoder().encode(latestYaml) as any); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + const result = await service.readSnapshot(projectId, 'notebook-1'); + + assert.isDefined(result); + assert.strictEqual(result!.size, 1); + assert.deepStrictEqual(result!.get('block-1'), [ + { output_type: 'stream', name: 'stdout', text: 'from latest' } + ]); + }); + + test('should return empty latest when neither latest nor any timestamped have outputs', async () => { + const workspaceFolder: WorkspaceFolder = { + uri: Uri.file('/workspace'), + name: 'workspace', + index: 0 + }; + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder]); + + const latestUri = Uri.file( + '/workspace/snapshots/project_test-project-id-123_notebook-1_latest.snapshot.deepnote' + ); + const timestampedUri = Uri.file( + '/workspace/snapshots/project_test-project-id-123_notebook-1_2025-01-01T10-00-00.snapshot.deepnote' + ); + + let findFilesCall = 0; + when(mockedVSCodeNamespaces.workspace.findFiles(anything(), anything(), anything())).thenCall(() => { + findFilesCall++; + + if (findFilesCall === 1) { + return Promise.resolve([latestUri]); + } + + return Promise.resolve([latestUri, timestampedUri]); + }); + + // Both files are missing the outputs field entirely. + const noOutputsYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-02T00:00:00Z' +project: + id: test-project-id-123 + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + blockGroup: group-1 + type: code + sortingKey: 'a0' + content: print(1) +`; + const mockFs = mock(); + when(mockFs.readFile(anything())).thenResolve(new TextEncoder().encode(noOutputsYaml) as any); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + const result = await service.readSnapshot(projectId, 'notebook-1'); + + // Both latest and timestamped have no outputs; we still resolve to a map (empty) + // so callers can distinguish "no snapshot file at all" from "snapshot exists, no outputs". + assert.isDefined(result); + assert.strictEqual(result!.size, 0); + }); }); suite('createSnapshot', () => { @@ -1149,8 +1351,11 @@ project: when(mockFs.copy(anything(), anything(), anything())).thenResolve(); when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); - // Call onExecutionComplete (which should auto-detect Run All) - await testServiceAny.onExecutionComplete(notebookUri); + // onExecutionComplete now just arms a debounced save. Drive the save + // directly via performSnapshotSave so the test doesn't race against the + // settle timer — the Run All / partial-run branch selection itself is + // what we care about here. + await testServiceAny.performSnapshotSave(notebookUri); // ASSERT: createSnapshot should be called (full snapshot, not just latest) assert.isTrue( @@ -1162,6 +1367,31 @@ project: 'updateLatestSnapshot should NOT be called when all code cells are executed' ); }); + + test('should cancel pending snapshot save when a cell starts a new execution', () => { + const testService = new SnapshotService(instance(mockEnvironmentCapture), mockDisposables); + const testServiceAny = testService as any; + + // Arm a pending save from a previous run. + testServiceAny.schedulePendingSnapshotSave(notebookUri); + assert.isTrue( + testServiceAny.pendingSnapshotSaves.has(notebookUri), + 'pending save should be armed after schedulePendingSnapshotSave' + ); + + // Simulate the notebook starting a new execution (cell becomes Executing). + const mockCell = { + notebook: { uri: Uri.parse(notebookUri) }, + metadata: { id: 'cell-1' } + }; + + testServiceAny.handleCellExecutionStateChange(mockCell, NotebookCellExecutionState.Executing); + + assert.isFalse( + testServiceAny.pendingSnapshotSaves.has(notebookUri), + 'pending save should be cancelled when a new execution starts' + ); + }); }); suite('captureEnvironmentBeforeExecution', () => { From 32671f0654872b28ce03c94563927542169a0247 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 23 Apr 2026 11:30:23 +0000 Subject: [PATCH 45/47] feat(deepnote): Enhance environment switching and notebook rebuilding logic - Updated `DeepnoteEnvironmentsView` to ensure all open sibling notebooks are rebuilt when switching project-scoped environments, improving consistency across notebooks. - Implemented rollback mechanisms to restore previous project environment mappings if a rebuild fails, enhancing reliability during environment transitions. - Refactored environment selection logic to use clearer variable names and improved error handling during notebook rebuilds. - Added unit tests to validate the new behavior, ensuring correct handling of environment switches and rollback scenarios. --- .../deepnoteEnvironmentsView.node.ts | 105 +++++++++++++-- .../deepnoteEnvironmentsView.unit.test.ts | 125 +++++++++++++++++- .../deepnoteProjectEnvironmentMapper.node.ts | 2 +- .../deepnote/deepnoteProjectUtils.ts | 9 +- .../deepnote/deepnoteTreeItem.unit.test.ts | 2 +- .../deepnote/snapshots/snapshotService.ts | 2 +- .../deepnote/deepnoteProjectFileReader.ts | 8 ++ .../deepnoteProjectFileReader.unit.test.ts} | 4 +- .../deepnote/deepnoteProjectIdResolver.ts | 4 +- 9 files changed, 233 insertions(+), 28 deletions(-) create mode 100644 src/platform/deepnote/deepnoteProjectFileReader.ts rename src/{notebooks/deepnote/deepnoteProjectUtils.unit.test.ts => platform/deepnote/deepnoteProjectFileReader.unit.test.ts} (96%) diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts index d1ccf53885..3e973f0ca2 100644 --- a/src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts +++ b/src/kernels/deepnote/environments/deepnoteEnvironmentsView.node.ts @@ -1,5 +1,6 @@ import { inject, injectable, named } from 'inversify'; import { + CancellationTokenSource, commands, Disposable, l10n, @@ -391,9 +392,9 @@ export class DeepnoteEnvironmentsView implements Disposable { await this.projectEnvironmentMapper.waitForInitialization(); // Get current environment selection - const currentEnvironmentId = this.projectEnvironmentMapper.getEnvironmentForProject(projectId); - const currentEnvironment = currentEnvironmentId - ? this.environmentManager.getEnvironment(currentEnvironmentId) + const previousEnvironmentId = this.projectEnvironmentMapper.getEnvironmentForProject(projectId); + const currentEnvironment = previousEnvironmentId + ? this.environmentManager.getEnvironment(previousEnvironmentId) : undefined; // Get all environments @@ -446,7 +447,7 @@ export class DeepnoteEnvironmentsView implements Disposable { } // Check if user selected the same environment - if (selectedEnvironmentId === currentEnvironmentId) { + if (selectedEnvironmentId === previousEnvironmentId) { logger.info(`User selected the same environment - no changes needed`); return; } else if (selectedEnvironmentId == null) { @@ -454,6 +455,30 @@ export class DeepnoteEnvironmentsView implements Disposable { return; } + const notebooksToRebuild: NotebookDocument[] = [notebook]; + const initiatingNotebookUri = notebook.uri.toString(); + + for (const openNotebook of workspace.notebookDocuments) { + if ( + openNotebook.notebookType !== 'deepnote' || + openNotebook.metadata?.deepnoteProjectId !== projectId || + openNotebook.uri.toString() === initiatingNotebookUri + ) { + continue; + } + + notebooksToRebuild.push(openNotebook); + } + + const nextEnvironmentId = selectedEnvironmentId; + const restoreProjectEnvironmentMapping = async () => { + if (previousEnvironmentId) { + await this.projectEnvironmentMapper.setEnvironmentForProject(projectId, previousEnvironmentId); + } else { + await this.projectEnvironmentMapper.removeEnvironmentForProject(projectId); + } + }; + // Check if any cells are currently executing using the kernel execution state // This is more reliable than checking executionSummary const kernel = this.kernelProvider.get(notebook); @@ -478,7 +503,7 @@ export class DeepnoteEnvironmentsView implements Disposable { } // User selected a different environment - switch to it - logger.info(`Switching notebook ${getDisplayPath(notebook.uri)} to environment ${selectedEnvironmentId}`); + logger.info(`Switching notebook ${getDisplayPath(notebook.uri)} to environment ${nextEnvironmentId}`); try { await window.withProgress( @@ -488,15 +513,71 @@ export class DeepnoteEnvironmentsView implements Disposable { cancellable: true }, async (progress, token) => { - // Update the project-to-environment mapping - await this.projectEnvironmentMapper.setEnvironmentForProject(projectId, selectedEnvironmentId); + const rebuiltNotebooks: NotebookDocument[] = []; + const rollbackNotebookUris = new Set(); + let mappingUpdated = false; + let notebookBeingRebuilt: NotebookDocument | undefined; - // Force rebuild the controller with the new environment - // This clears cached metadata and creates a fresh controller. - // await this.kernelAutoSelector.ensureKernelSelected(activeNotebook); - await this.kernelAutoSelector.rebuildController(notebook, progress, token); + try { + // Update the project-to-environment mapping before rebuilding so controller creation + // resolves the new project-scoped environment for every open notebook in the project. + await this.projectEnvironmentMapper.setEnvironmentForProject(projectId, nextEnvironmentId); + mappingUpdated = true; + + for (const targetNotebook of notebooksToRebuild) { + notebookBeingRebuilt = targetNotebook; + await this.kernelAutoSelector.rebuildController(targetNotebook, progress, token); + rebuiltNotebooks.push(targetNotebook); + } - logger.info(`Successfully switched to environment ${selectedEnvironmentId}`); + logger.info(`Successfully switched project ${projectId} to environment ${nextEnvironmentId}`); + } catch (error) { + if (mappingUpdated) { + await restoreProjectEnvironmentMapping(); + + const notebooksToRollBack: NotebookDocument[] = []; + const addRollbackNotebook = (candidate: NotebookDocument | undefined) => { + if (!candidate) { + return; + } + + const candidateKey = candidate.uri.toString(); + if (rollbackNotebookUris.has(candidateKey)) { + return; + } + + rollbackNotebookUris.add(candidateKey); + notebooksToRollBack.push(candidate); + }; + + addRollbackNotebook(notebookBeingRebuilt); + rebuiltNotebooks.forEach(addRollbackNotebook); + + const rollbackCancellation = new CancellationTokenSource(); + try { + for (const rollbackNotebook of notebooksToRollBack) { + try { + await this.kernelAutoSelector.rebuildController( + rollbackNotebook, + progress, + rollbackCancellation.token + ); + } catch (rollbackError) { + logger.error( + `Failed to roll back environment switch for ${getDisplayPath( + rollbackNotebook.uri + )}`, + rollbackError + ); + } + } + } finally { + rollbackCancellation.dispose(); + } + } + + throw error; + } } ); diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts index 0485b664fd..e1f9a79b3e 100644 --- a/src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts +++ b/src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts @@ -6,7 +6,7 @@ import { DeepnoteEnvironmentsView } from './deepnoteEnvironmentsView.node'; import { IDeepnoteEnvironmentManager, IDeepnoteKernelAutoSelector, IDeepnoteProjectEnvironmentMapper } from '../types'; import { IPythonApiProvider } from '../../../platform/api/types'; import { IDisposableRegistry, IOutputChannel } from '../../../platform/common/types'; -import { IKernelProvider } from '../../../kernels/types'; +import { IKernel, IKernelProvider, INotebookKernelExecution } from '../../../kernels/types'; import { DeepnoteEnvironment } from './deepnoteEnvironment'; import { PythonEnvironment } from '../../../platform/pythonEnvironments/info'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; @@ -809,6 +809,21 @@ suite('DeepnoteEnvironmentsView', () => { }; const testProjectId = 'test-project-id'; + const createNotebookDocument = (path: string, projectId: string): NotebookDocument => ({ + uri: Uri.file(path), + notebookType: 'deepnote', + version: 1, + isDirty: false, + isUntitled: false, + isClosed: false, + metadata: { deepnoteProjectId: projectId }, + cellCount: 0, + cellAt: () => { + throw new Error('Not implemented'); + }, + getCells: () => [], + save: async () => true + }); setup(() => { resetCalls(mockConfigManager); @@ -976,6 +991,114 @@ suite('DeepnoteEnvironmentsView', () => { verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); }); + test('should rebuild every open sibling notebook when switching a project-scoped environment', async () => { + // Catches: project-scoped environment switches only rebuilt the initiating notebook, leaving siblings stale. + + // Arrange + const initiatingNotebook = createNotebookDocument('/workspace/initiating.deepnote', testProjectId); + const siblingNotebook = createNotebookDocument('/workspace/sibling.deepnote', testProjectId); + const otherProjectNotebook = createNotebookDocument('/workspace/other.deepnote', 'different-project-id'); + const executionKernel = mock(); + const notebookExecution = mock(); + + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn( + currentEnvironment.id + ); + when(mockConfigManager.getEnvironment(currentEnvironment.id)).thenReturn(currentEnvironment); + when(mockConfigManager.listEnvironments()).thenReturn([currentEnvironment, newEnvironment]); + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenCall( + (items: ReadonlyArray<{ environmentId?: string }>) => { + return Promise.resolve(items.find((item) => item.environmentId === newEnvironment.id)); + } + ); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([ + initiatingNotebook, + siblingNotebook, + otherProjectNotebook + ]); + when(mockKernelProvider.get(initiatingNotebook)).thenReturn(instance(executionKernel)); + when(notebookExecution.pendingCells).thenReturn([]); + when(mockKernelProvider.getKernelExecution(instance(executionKernel))).thenReturn( + instance(notebookExecution) + ); + when(mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, newEnvironment.id)).thenResolve(); + when(mockKernelAutoSelector.rebuildController(initiatingNotebook, anything(), anything())).thenResolve(); + when(mockKernelAutoSelector.rebuildController(siblingNotebook, anything(), anything())).thenResolve(); + + // Act + await view.selectEnvironmentForNotebook({ notebook: initiatingNotebook }); + + // Assert + verify(mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, newEnvironment.id)).once(); + verify(mockKernelAutoSelector.rebuildController(initiatingNotebook, anything(), anything())).once(); + verify(mockKernelAutoSelector.rebuildController(siblingNotebook, anything(), anything())).once(); + verify(mockKernelAutoSelector.rebuildController(otherProjectNotebook, anything(), anything())).never(); + verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); + }); + + test('should restore the previous project mapping and roll back rebuilt notebooks when a sibling rebuild fails', async () => { + // Catches: the new project mapping was committed before rebuild and stayed wrong when a sibling rebuild failed. + + // Arrange + const initiatingNotebook = createNotebookDocument('/workspace/rollback-initiating.deepnote', testProjectId); + const siblingNotebook = createNotebookDocument('/workspace/rollback-sibling.deepnote', testProjectId); + const executionKernel = mock(); + const notebookExecution = mock(); + const rebuildCalls: string[] = []; + const rebuildError = new Error('rebuild failed'); + + when(mockProjectEnvironmentMapper.getEnvironmentForProject(testProjectId)).thenReturn( + currentEnvironment.id + ); + when(mockConfigManager.getEnvironment(currentEnvironment.id)).thenReturn(currentEnvironment); + when(mockConfigManager.listEnvironments()).thenReturn([currentEnvironment, newEnvironment]); + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenCall( + (items: ReadonlyArray<{ environmentId?: string }>) => { + return Promise.resolve(items.find((item) => item.environmentId === newEnvironment.id)); + } + ); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([initiatingNotebook, siblingNotebook]); + when(mockKernelProvider.get(initiatingNotebook)).thenReturn(instance(executionKernel)); + when(notebookExecution.pendingCells).thenReturn([]); + when(mockKernelProvider.getKernelExecution(instance(executionKernel))).thenReturn( + instance(notebookExecution) + ); + when(mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, newEnvironment.id)).thenResolve(); + when( + mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, currentEnvironment.id) + ).thenResolve(); + when(mockKernelAutoSelector.rebuildController(anything(), anything(), anything())).thenCall( + async (targetNotebook: NotebookDocument) => { + const notebookPath = targetNotebook.uri.toString(); + rebuildCalls.push(notebookPath); + + if ( + notebookPath === siblingNotebook.uri.toString() && + rebuildCalls.filter((call) => call === notebookPath).length === 1 + ) { + throw rebuildError; + } + } + ); + + // Act + await view.selectEnvironmentForNotebook({ notebook: initiatingNotebook }); + + // Assert + verify(mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, newEnvironment.id)).once(); + verify(mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, currentEnvironment.id)).once(); + verify(mockKernelAutoSelector.rebuildController(initiatingNotebook, anything(), anything())).twice(); + verify(mockKernelAutoSelector.rebuildController(siblingNotebook, anything(), anything())).atLeast(1); + assert.deepEqual(rebuildCalls, [ + initiatingNotebook.uri.toString(), + siblingNotebook.uri.toString(), + siblingNotebook.uri.toString(), + initiatingNotebook.uri.toString() + ]); + verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).never(); + verify(mockedVSCodeNamespaces.window.showErrorMessage(anything(), anything(), anything())).once(); + }); + test('should successfully switch from external to managed environment', async () => { // Mock active notebook const notebookUri = Uri.file('/workspace/notebook.deepnote'); diff --git a/src/kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.node.ts b/src/kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.node.ts index 3aa467becc..6049cdfe75 100644 --- a/src/kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.node.ts +++ b/src/kernels/deepnote/environments/deepnoteProjectEnvironmentMapper.node.ts @@ -106,7 +106,7 @@ export class DeepnoteProjectEnvironmentMapper implements IDeepnoteProjectEnviron * key. The migration is a one-shot: after resolving project ids from each * `.deepnote` file, the legacy key is cleared. * - * Entries whose project id cannot be resolved (file missing, unparseable) + * Entries whose project id cannot be resolved (file missing, unparsable) * are logged and skipped. */ private async migrateLegacyMappings(): Promise { diff --git a/src/notebooks/deepnote/deepnoteProjectUtils.ts b/src/notebooks/deepnote/deepnoteProjectUtils.ts index 976008ac6e..b8e67386c6 100644 --- a/src/notebooks/deepnote/deepnoteProjectUtils.ts +++ b/src/notebooks/deepnote/deepnoteProjectUtils.ts @@ -1,11 +1,4 @@ -import { deserializeDeepnoteFile, type DeepnoteFile } from '@deepnote/blocks'; -import { Uri, workspace } from 'vscode'; - -export async function readDeepnoteProjectFile(fileUri: Uri): Promise { - const fileContent = await workspace.fs.readFile(fileUri); - const yamlContent = new TextDecoder().decode(fileContent); - return deserializeDeepnoteFile(yamlContent); -} +export { readDeepnoteProjectFile } from '../../platform/deepnote/deepnoteProjectFileReader'; /** * Compute a hash of the requirements to detect changes. diff --git a/src/notebooks/deepnote/deepnoteTreeItem.unit.test.ts b/src/notebooks/deepnote/deepnoteTreeItem.unit.test.ts index e09042a0eb..40676a2b7e 100644 --- a/src/notebooks/deepnote/deepnoteTreeItem.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteTreeItem.unit.test.ts @@ -825,7 +825,7 @@ suite('DeepnoteTreeItem', () => { // instance rather than a DeepnoteTreeItem, so instance methods on the subclass // prototype are NOT reachable via the usual `item.updateVisualFields()` call. // We invoke the method by reading it off the class prototype and calling it on - // the instance — same behaviour, just side-steps the proxy's prototype chain loss. + // the instance — same behavior, just side-steps the proxy's prototype chain loss. function callUpdateVisualFields(item: DeepnoteTreeItem): void { const method = DeepnoteTreeItem.prototype.updateVisualFields; diff --git a/src/notebooks/deepnote/snapshots/snapshotService.ts b/src/notebooks/deepnote/snapshots/snapshotService.ts index f2b66df6f3..f0302f61da 100644 --- a/src/notebooks/deepnote/snapshots/snapshotService.ts +++ b/src/notebooks/deepnote/snapshots/snapshotService.ts @@ -112,7 +112,7 @@ interface NotebookExecutionState { const recentWriteExpirationMs = 2000; /** - * After the cell execution queue completes, the kernel can still be delivering iopub + * After the cell execution queue completes, the kernel can still be delivering IOPub * messages (stream output, display_data, etc.) that land on `cell.outputs` AFTER * `onDidCompleteQueueExecution` fires. Instead of reading `cell.outputs` immediately, * we arm a timer and reset it every time `workspace.onDidChangeNotebookDocument` reports diff --git a/src/platform/deepnote/deepnoteProjectFileReader.ts b/src/platform/deepnote/deepnoteProjectFileReader.ts new file mode 100644 index 0000000000..a0ebe5133f --- /dev/null +++ b/src/platform/deepnote/deepnoteProjectFileReader.ts @@ -0,0 +1,8 @@ +import { deserializeDeepnoteFile, type DeepnoteFile } from '@deepnote/blocks'; +import { Uri, workspace } from 'vscode'; + +export async function readDeepnoteProjectFile(fileUri: Uri): Promise { + const fileContent = await workspace.fs.readFile(fileUri); + const yamlContent = new TextDecoder().decode(fileContent); + return deserializeDeepnoteFile(yamlContent); +} diff --git a/src/notebooks/deepnote/deepnoteProjectUtils.unit.test.ts b/src/platform/deepnote/deepnoteProjectFileReader.unit.test.ts similarity index 96% rename from src/notebooks/deepnote/deepnoteProjectUtils.unit.test.ts rename to src/platform/deepnote/deepnoteProjectFileReader.unit.test.ts index 3ff362ed4d..54e85d6df9 100644 --- a/src/notebooks/deepnote/deepnoteProjectUtils.unit.test.ts +++ b/src/platform/deepnote/deepnoteProjectFileReader.unit.test.ts @@ -3,10 +3,10 @@ import * as sinon from 'sinon'; import { anything, instance, mock, when } from 'ts-mockito'; import { Uri, workspace } from 'vscode'; -import { readDeepnoteProjectFile } from './deepnoteProjectUtils'; +import { readDeepnoteProjectFile } from './deepnoteProjectFileReader'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; -suite('DeepnoteProjectUtils', () => { +suite('DeepnoteProjectFileReader', () => { let sandbox: sinon.SinonSandbox; setup(() => { diff --git a/src/platform/deepnote/deepnoteProjectIdResolver.ts b/src/platform/deepnote/deepnoteProjectIdResolver.ts index ad78a1b8ff..b39d3c56c8 100644 --- a/src/platform/deepnote/deepnoteProjectIdResolver.ts +++ b/src/platform/deepnote/deepnoteProjectIdResolver.ts @@ -1,7 +1,7 @@ import { NotebookDocument, Uri } from 'vscode'; import { logger } from '../logging'; -import { readDeepnoteProjectFile } from '../../notebooks/deepnote/deepnoteProjectUtils'; +import { readDeepnoteProjectFile } from './deepnoteProjectFileReader'; /** * Resolve the Deepnote project id for an open notebook document. @@ -29,7 +29,7 @@ export async function resolveProjectIdForFile(fileUri: Uri): Promise Date: Thu, 23 Apr 2026 14:56:01 +0000 Subject: [PATCH 46/47] feat(deepnote): Process already-open notebooks during activation - Enhanced `DeepnoteActivationService` to handle notebooks that are already open at activation time, ensuring they are processed correctly. - Updated unit tests to validate the new behavior, confirming that already-open deepnote notebooks are checked and split if necessary. - Refactored `DeepnoteAutoSplitter` to improve file URI handling and avoid collisions with existing files during the split operation. --- .../deepnote/deepnoteActivationService.ts | 6 + .../deepnoteActivationService.unit.test.ts | 69 +++++++- .../deepnote/deepnoteAutoSplitter.ts | 28 ++- .../deepnoteAutoSplitter.unit.test.ts | 159 ++++++++++++++++++ 4 files changed, 253 insertions(+), 9 deletions(-) create mode 100644 src/notebooks/deepnote/deepnoteAutoSplitter.unit.test.ts diff --git a/src/notebooks/deepnote/deepnoteActivationService.ts b/src/notebooks/deepnote/deepnoteActivationService.ts index b81db6a7ba..2eb47f6a4f 100644 --- a/src/notebooks/deepnote/deepnoteActivationService.ts +++ b/src/notebooks/deepnote/deepnoteActivationService.ts @@ -61,6 +61,12 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic void this.checkAndSplitIfNeeded(doc); }) ); + + // Process notebooks that are already open at activation time (e.g., restored windows) + for (const doc of workspace.notebookDocuments) { + void this.checkAndSplitIfNeeded(doc); + } + this.extensionContext.subscriptions.push( workspace.onDidChangeConfiguration((event) => { if (event.affectsConfiguration('deepnote.snapshots.enabled')) { diff --git a/src/notebooks/deepnote/deepnoteActivationService.unit.test.ts b/src/notebooks/deepnote/deepnoteActivationService.unit.test.ts index a8068d5f64..b534cd8954 100644 --- a/src/notebooks/deepnote/deepnoteActivationService.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteActivationService.unit.test.ts @@ -1,5 +1,7 @@ +import { serializeDeepnoteFile, type DeepnoteFile } from '@deepnote/blocks'; import { assert } from 'chai'; -import { anything, verify, when } from 'ts-mockito'; +import { anything, instance, mock, verify, when } from 'ts-mockito'; +import { Uri, workspace } from 'vscode'; import { DeepnoteActivationService } from './deepnoteActivationService'; import { DeepnoteNotebookManager } from './deepnoteNotebookManager'; @@ -164,6 +166,71 @@ suite('DeepnoteActivationService', () => { verify(mockedVSCodeNamespaces.window.showInformationMessage(anything(), anything(), anything())).once(); }); + + test('should process already-open deepnote notebooks at activation time', async () => { + resetVSCodeMocks(); + + when( + mockedVSCodeNamespaces.workspace.registerNotebookSerializer(anything(), anything(), anything()) + ).thenReturn({ dispose: () => undefined } as any); + when(mockedVSCodeNamespaces.workspace.onDidChangeConfiguration).thenReturn((() => ({ + dispose: () => undefined + })) as any); + + const deepnoteUri = Uri.file('/workspace/existing.deepnote'); + const otherUri = Uri.file('/workspace/existing.ipynb'); + const deepnoteFile: DeepnoteFile = { + version: '1.0.0', + metadata: { createdAt: '2025-01-01T00:00:00Z' }, + project: { + id: 'project-1', + name: 'Restored Project', + notebooks: [ + { + id: 'notebook-1', + name: 'Notebook 1', + blocks: [] + } + ] + } + }; + const yaml = serializeDeepnoteFile(deepnoteFile); + const alreadyOpenDeepnoteDoc = { notebookType: 'deepnote', uri: deepnoteUri } as any; + const alreadyOpenOtherDoc = { notebookType: 'jupyter-notebook', uri: otherUri } as any; + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([ + alreadyOpenDeepnoteDoc, + alreadyOpenOtherDoc + ]); + + const readCalls: string[] = []; + const mockFs = mock(); + + when(mockFs.readFile(anything())).thenCall((uri: Uri) => { + readCalls.push(uri.path); + + return Promise.resolve(new TextEncoder().encode(yaml)); + }); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + activationService = new DeepnoteActivationService( + mockExtensionContext, + manager, + mockIntegrationManager, + mockLogger + ); + + try { + activationService.activate(); + } catch { + // Activation may fail in the test environment, but iteration should still occur. + } + + // Allow the fire-and-forget checkAndSplitIfNeeded promises to run + await new Promise((resolve) => setImmediate(resolve)); + + assert.deepStrictEqual(readCalls, [deepnoteUri.path]); + }); }); suite('component initialization', () => { diff --git a/src/notebooks/deepnote/deepnoteAutoSplitter.ts b/src/notebooks/deepnote/deepnoteAutoSplitter.ts index caeaf655ed..6bb5487c6c 100644 --- a/src/notebooks/deepnote/deepnoteAutoSplitter.ts +++ b/src/notebooks/deepnote/deepnoteAutoSplitter.ts @@ -3,10 +3,9 @@ import { l10n, RelativePattern, Uri, window, workspace } from 'vscode'; import { logger } from '../../platform/logging'; import { + buildSiblingNotebookFileUri, buildSingleNotebookFile, - computeSnapshotHash, - getFileStem, - slugifyNotebookNameOrFallback + computeSnapshotHash } from './deepnoteNotebookFileFactory'; import { slugifyProjectName } from './snapshots/snapshotFiles'; @@ -43,15 +42,27 @@ export class DeepnoteAutoSplitter { `[AutoSplitter] Splitting ${nonInitNotebooks.length} notebooks from project ${deepnoteFile.project.id}` ); - const parentDir = Uri.joinPath(fileUri, '..'); - const originalStem = getFileStem(fileUri); const newFiles: Uri[] = []; + const reservedPaths = new Set(); + const exists = async (u: Uri): Promise => { + if (reservedPaths.has(u.path)) { + return true; + } + + try { + await workspace.fs.stat(u); + + return true; + } catch { + return false; + } + }; // Create a new file for each extra notebook for (const notebook of extraNotebooks) { - const notebookSlug = slugifyNotebookNameOrFallback(notebook.name); - const newFileName = `${originalStem}_${notebookSlug}.deepnote`; - const newFileUri = Uri.joinPath(parentDir, newFileName); + const newFileUri = await buildSiblingNotebookFileUri(fileUri, notebook.name, exists); + + reservedPaths.add(newFileUri.path); const newProject = await buildSingleNotebookFile(deepnoteFile, notebook); @@ -60,6 +71,7 @@ export class DeepnoteAutoSplitter { newFiles.push(newFileUri); + const newFileName = newFileUri.path.split('/').pop(); logger.info(`[AutoSplitter] Created ${newFileName} for notebook "${notebook.name}"`); } diff --git a/src/notebooks/deepnote/deepnoteAutoSplitter.unit.test.ts b/src/notebooks/deepnote/deepnoteAutoSplitter.unit.test.ts new file mode 100644 index 0000000000..b482e6709a --- /dev/null +++ b/src/notebooks/deepnote/deepnoteAutoSplitter.unit.test.ts @@ -0,0 +1,159 @@ +import { type DeepnoteFile } from '@deepnote/blocks'; +import { assert } from 'chai'; +import { anything, instance, mock, when } from 'ts-mockito'; +import { FileType, Uri, workspace } from 'vscode'; + +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; +import { DeepnoteAutoSplitter } from './deepnoteAutoSplitter'; + +function createDeepnoteFile(notebooks: Array<{ id: string; name: string }>, initNotebookId?: string): DeepnoteFile { + return { + metadata: { createdAt: '2025-01-01T00:00:00Z' }, + project: { + id: 'project-1', + initNotebookId, + name: 'Test Project', + notebooks: notebooks.map(({ id, name }) => ({ + blocks: [], + executionMode: 'block', + id, + name + })) + }, + version: '1.0.0' + }; +} + +interface FsMockSetup { + existingPaths: Set; + writeCalls: Uri[]; +} + +function setupFsMocks(existingPaths: string[] = []): FsMockSetup { + const setup: FsMockSetup = { + existingPaths: new Set(existingPaths), + writeCalls: [] + }; + const mockFs = mock(); + + when(mockFs.stat(anything())).thenCall((uri: Uri) => { + if (setup.existingPaths.has(uri.path)) { + return Promise.resolve({ type: FileType.File, ctime: 0, mtime: 0, size: 0 } as any); + } + + return Promise.reject(new Error('ENOENT')); + }); + when(mockFs.writeFile(anything(), anything())).thenCall((uri: Uri) => { + setup.writeCalls.push(uri); + setup.existingPaths.add(uri.path); + + return Promise.resolve(); + }); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn(undefined); + + return setup; +} + +suite('DeepnoteAutoSplitter', () => { + let splitter: DeepnoteAutoSplitter; + + setup(() => { + resetVSCodeMocks(); + splitter = new DeepnoteAutoSplitter(); + }); + + suite('splitIfNeeded', () => { + test('should not split when there is only one non-init notebook', async () => { + const fs = setupFsMocks(); + const file = createDeepnoteFile([{ id: 'nb-1', name: 'Notebook 1' }]); + + const result = await splitter.splitIfNeeded(Uri.file('/workspace/project.deepnote'), file); + + assert.strictEqual(result.wasSplit, false); + assert.deepStrictEqual(result.newFiles, []); + assert.strictEqual(fs.writeCalls.length, 0); + }); + + test('should suffix duplicate notebook names to avoid collisions', async () => { + const fs = setupFsMocks(); + const file = createDeepnoteFile([ + { id: 'nb-1', name: 'Analysis' }, + { id: 'nb-2', name: 'Analysis' }, + { id: 'nb-3', name: 'Analysis' } + ]); + + const result = await splitter.splitIfNeeded(Uri.file('/workspace/project.deepnote'), file); + + assert.strictEqual(result.wasSplit, true); + assert.strictEqual(result.newFiles.length, 2); + + const newFilePaths = result.newFiles.map((uri) => uri.path); + assert.deepStrictEqual(newFilePaths, [ + '/workspace/project_analysis.deepnote', + '/workspace/project_analysis_2.deepnote' + ]); + + // Verify each path was written exactly once and they are unique + assert.strictEqual(new Set(fs.writeCalls.map((uri) => uri.path)).size, fs.writeCalls.length); + }); + + test('should suffix notebook names that slugify to the same value', async () => { + const fs = setupFsMocks(); + const file = createDeepnoteFile([ + { id: 'nb-1', name: 'Primary' }, + { id: 'nb-2', name: 'My Notebook' }, + { id: 'nb-3', name: 'my-notebook' } + ]); + + const result = await splitter.splitIfNeeded(Uri.file('/workspace/project.deepnote'), file); + + assert.strictEqual(result.wasSplit, true); + assert.strictEqual(result.newFiles.length, 2); + + const newFilePaths = result.newFiles.map((uri) => uri.path); + assert.deepStrictEqual(newFilePaths, [ + '/workspace/project_my-notebook.deepnote', + '/workspace/project_my-notebook_2.deepnote' + ]); + + assert.strictEqual(new Set(fs.writeCalls.map((uri) => uri.path)).size, fs.writeCalls.length); + }); + + test('should suffix fallback-slug collisions when notebook names produce empty slugs', async () => { + const fs = setupFsMocks(); + const file = createDeepnoteFile([ + { id: 'nb-1', name: 'Primary' }, + { id: 'nb-2', name: '!!!' }, + { id: 'nb-3', name: '???' } + ]); + + const result = await splitter.splitIfNeeded(Uri.file('/workspace/project.deepnote'), file); + + assert.strictEqual(result.wasSplit, true); + assert.strictEqual(result.newFiles.length, 2); + + const newFilePaths = result.newFiles.map((uri) => uri.path); + assert.deepStrictEqual(newFilePaths, [ + '/workspace/project_notebook.deepnote', + '/workspace/project_notebook_2.deepnote' + ]); + + assert.strictEqual(new Set(fs.writeCalls.map((uri) => uri.path)).size, fs.writeCalls.length); + }); + + test('should avoid colliding with an existing sibling file on disk', async () => { + setupFsMocks(['/workspace/project_analysis.deepnote']); + const file = createDeepnoteFile([ + { id: 'nb-1', name: 'Primary' }, + { id: 'nb-2', name: 'Analysis' } + ]); + + const result = await splitter.splitIfNeeded(Uri.file('/workspace/project.deepnote'), file); + + assert.strictEqual(result.wasSplit, true); + assert.strictEqual(result.newFiles.length, 1); + assert.strictEqual(result.newFiles[0].path, '/workspace/project_analysis_2.deepnote'); + }); + }); +}); From b3278e94d74f3d3e19889728e2d2a52e3fb0b3b7 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 24 Apr 2026 13:43:13 +0000 Subject: [PATCH 47/47] Cleanup tests --- .../deepnoteExplorerView.unit.test.ts | 418 -------- .../deepnoteFileChangeWatcher.unit.test.ts | 916 +----------------- 2 files changed, 12 insertions(+), 1322 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts b/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts index 695ea2b146..3a832f5d2d 100644 --- a/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts @@ -11,7 +11,6 @@ import { DeepnoteTreeItem, DeepnoteTreeItemType, NOTEBOOK_FILE_CONTEXT_VALUE, - type DeepnoteTreeItemContext, type ProjectGroupData } from './deepnoteTreeItem'; import type { IExtensionContext } from '../../platform/common/types'; @@ -45,188 +44,6 @@ function createUuidMock(uuids: string[]): sinon.SinonStub { return stub; } -suite('DeepnoteExplorerView', () => { - let explorerView: DeepnoteExplorerView; - let mockExtensionContext: IExtensionContext; - let mockLogger: ILogger; - - setup(() => { - mockExtensionContext = { - subscriptions: [] - } as any; - - mockLogger = createMockLogger(); - explorerView = new DeepnoteExplorerView(mockExtensionContext, mockLogger); - }); - - teardown(() => { - explorerView.dispose(); - }); - - suite('constructor', () => { - test('should create instance with extension context', () => { - assert.isDefined(explorerView); - }); - - test('should initialize with proper dependencies', () => { - // Verify that internal components are accessible - assert.isDefined((explorerView as any).extensionContext); - assert.strictEqual((explorerView as any).extensionContext, mockExtensionContext); - }); - }); - - suite('activate', () => { - test('should attempt to activate without errors', () => { - // This test verifies the activate method can be called - try { - explorerView.activate(); - // If we get here, activation succeeded - assert.isTrue(true, 'activate() completed successfully'); - } catch (error) { - // Expected in test environment without full VS Code API - assert.isString(error.message, 'activate() method exists and attempts initialization'); - } - }); - }); - - suite('openNotebook', () => { - const mockContext: DeepnoteTreeItemContext = { - filePath: '/test/path/project.deepnote', - projectId: 'project-123', - notebookId: 'notebook-456' - }; - - test('should handle context without notebookId', async () => { - const contextWithoutId = { ...mockContext, notebookId: undefined }; - - // This should not throw an error - method should handle gracefully - try { - await (explorerView as any).openNotebook(contextWithoutId); - assert.isTrue(true, 'openNotebook handled undefined notebookId gracefully'); - } catch (error) { - // Expected in test environment - assert.isString(error.message, 'openNotebook method exists'); - } - }); - - test('should handle valid context', async () => { - try { - await (explorerView as any).openNotebook(mockContext); - assert.isTrue(true, 'openNotebook handled valid context'); - } catch (error) { - // Expected in test environment without VS Code APIs - assert.isString(error.message, 'openNotebook method exists and processes context'); - } - }); - - test('should use base file URI without fragments', async () => { - // This test verifies that we're using the simplified approach - // The actual URI creation is tested through integration, but we can verify - // that the method exists and processes the context correctly - try { - await (explorerView as any).openNotebook(mockContext); - assert.isTrue(true, 'openNotebook uses base file URI approach'); - } catch (error) { - // Expected in test environment - the method should exist and attempt to process - assert.isString(error.message, 'openNotebook method processes context'); - } - }); - }); - - suite('openFile', () => { - test('should handle non-project file items', async () => { - const mockTreeItem = { - type: 'notebook', // Not ProjectFile - context: { filePath: '/test/path' } - } as any; - - try { - await (explorerView as any).openFile(mockTreeItem); - assert.isTrue(true, 'openFile handled non-project file gracefully'); - } catch (error) { - // Expected in test environment - assert.isString(error.message, 'openFile method exists'); - } - }); - - test('should handle project file items', async () => { - const mockTreeItem = { - type: 'ProjectFile', - context: { filePath: '/test/path/project.deepnote' } - } as any; - - try { - await (explorerView as any).openFile(mockTreeItem); - assert.isTrue(true, 'openFile handled project file'); - } catch (error) { - // Expected in test environment - assert.isString(error.message, 'openFile method exists and processes files'); - } - }); - }); - - suite('revealActiveNotebook', () => { - test('should handle missing active notebook editor', async () => { - try { - await (explorerView as any).revealActiveNotebook(); - assert.isTrue(true, 'revealActiveNotebook handled missing editor gracefully'); - } catch (error) { - // Expected in test environment - assert.isString(error.message, 'revealActiveNotebook method exists'); - } - }); - }); - - suite('refreshExplorer', () => { - test('should call refresh method', () => { - try { - (explorerView as any).refreshExplorer(); - assert.isTrue(true, 'refreshExplorer method exists and can be called'); - } catch (error) { - // Expected in test environment - assert.isString(error.message, 'refreshExplorer method exists'); - } - }); - }); - - suite('integration scenarios', () => { - test('should handle multiple explorer view instances', () => { - const context1 = { subscriptions: [] } as any; - const context2 = { subscriptions: [] } as any; - - const logger1 = createMockLogger(); - const logger2 = createMockLogger(); - const view1 = new DeepnoteExplorerView(context1, logger1); - const view2 = new DeepnoteExplorerView(context2, logger2); - - try { - // Verify each view has its own context - assert.strictEqual((view1 as any).extensionContext, context1); - assert.strictEqual((view2 as any).extensionContext, context2); - assert.notStrictEqual((view1 as any).extensionContext, (view2 as any).extensionContext); - - // Verify views are independent instances - assert.notStrictEqual(view1, view2); - } finally { - view1.dispose(); - view2.dispose(); - } - }); - - test('should maintain component references', () => { - // Verify that internal components exist - assert.isDefined((explorerView as any).extensionContext); - - // After construction, some components should be initialized - const hasTreeDataProvider = (explorerView as any).treeDataProvider !== undefined; - const hasSerializer = (explorerView as any).serializer !== undefined; - - // At least one component should be defined after construction - assert.isTrue(hasTreeDataProvider || hasSerializer, 'Components are being initialized'); - }); - }); -}); - suite('DeepnoteExplorerView - Empty State Commands', () => { let explorerView: DeepnoteExplorerView; let mockContext: IExtensionContext; @@ -491,29 +308,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { expect(capturedUri!.path).to.include('test.deepnote'); }); - test('should import and convert jupyter files', async () => { - const workspaceFolder = { uri: Uri.file('/workspace') }; - const sourceUri = Uri.file('/external/my-notebook.ipynb'); - - when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder as any]); - when(mockedVSCodeNamespaces.window.showOpenDialog(anything())).thenReturn(Promise.resolve([sourceUri])); - - const mockFS = mock(); - when(mockFS.stat(anything())).thenReject(new Error('File not found')); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - - let infoMessageShown = false; - when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenCall(() => { - infoMessageShown = true; - return Promise.resolve(undefined); - }); - - await (explorerView as any).importNotebook(); - - // Verify success message was shown (indicating convert was called successfully) - expect(infoMessageShown).to.be.true; - }); - test('should import multiple files', async () => { const workspaceFolder = { uri: Uri.file('/workspace') }; const deepnoteUri = Uri.file('/external/test1.deepnote'); @@ -571,23 +365,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { expect(writeFileCalled).to.be.false; }); - test('should handle import errors', async () => { - const workspaceFolder = { uri: Uri.file('/workspace') }; - const sourceUri = Uri.file('/external/test.ipynb'); - - when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder as any]); - when(mockedVSCodeNamespaces.window.showOpenDialog(anything())).thenReturn(Promise.resolve([sourceUri])); - - const mockFS = mock(); - when(mockFS.stat(anything())).thenReject(new Error('File not found')); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - - // Test is simplified - the mock convert function succeeds by default - // To properly test error handling, we would need to modify the mock in vscode-mock.ts - // For now, we'll just verify the method completes without throwing - await (explorerView as any).importNotebook(); - }); - test('should return early if user cancels dialog', async () => { const workspaceFolder = { uri: Uri.file('/workspace') }; @@ -633,29 +410,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { }); suite('importJupyterNotebook', () => { - test('should import jupyter notebook with correct naming', async () => { - const workspaceFolder = { uri: Uri.file('/workspace') }; - const sourceUri = Uri.file('/external/my-analysis.ipynb'); - - when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder as any]); - when(mockedVSCodeNamespaces.window.showOpenDialog(anything())).thenReturn(Promise.resolve([sourceUri])); - - const mockFS = mock(); - when(mockFS.stat(anything())).thenReject(new Error('File not found')); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - - let infoMessageShown = false; - when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenCall(() => { - infoMessageShown = true; - return Promise.resolve(undefined); - }); - - await (explorerView as any).importJupyterNotebook(); - - // Verify success message was shown (indicating convert was called successfully) - expect(infoMessageShown).to.be.true; - }); - test('should import multiple jupyter notebooks', async () => { const workspaceFolder = { uri: Uri.file('/workspace') }; const sourceUris = [Uri.file('/external/notebook1.ipynb'), Uri.file('/external/notebook2.ipynb')]; @@ -701,23 +455,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { expect(errorShown).to.be.true; }); - test('should handle conversion errors', async () => { - const workspaceFolder = { uri: Uri.file('/workspace') }; - const sourceUri = Uri.file('/external/test.ipynb'); - - when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder as any]); - when(mockedVSCodeNamespaces.window.showOpenDialog(anything())).thenReturn(Promise.resolve([sourceUri])); - - const mockFS = mock(); - when(mockFS.stat(anything())).thenReject(new Error('File not found')); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - - // Test is simplified - the mock convert function succeeds by default - // To properly test error handling, we would need to modify the mock in vscode-mock.ts - // For now, we'll just verify the method completes without throwing - await (explorerView as any).importJupyterNotebook(); - }); - test('should return early if user cancels dialog', async () => { const workspaceFolder = { uri: Uri.file('/workspace') }; @@ -759,29 +496,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { expect(showInfoCalled).to.be.true; expect(executeCommandCalled).to.be.true; }); - - test('should remove .ipynb extension case-insensitively', async () => { - const workspaceFolder = { uri: Uri.file('/workspace') }; - const sourceUri = Uri.file('/external/notebook.IPYNB'); - - when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder as any]); - when(mockedVSCodeNamespaces.window.showOpenDialog(anything())).thenReturn(Promise.resolve([sourceUri])); - - const mockFS = mock(); - when(mockFS.stat(anything())).thenReject(new Error('File not found')); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - - let infoMessageShown = false; - when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenCall(() => { - infoMessageShown = true; - return Promise.resolve(undefined); - }); - - await (explorerView as any).importJupyterNotebook(); - - // Verify success message was shown (indicating convert was called successfully) - expect(infoMessageShown).to.be.true; - }); }); suite('createAndAddNotebookToProject', () => { @@ -1316,26 +1030,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); }); - test('should return early if tree item type is not Notebook', async () => { - const mockTreeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/workspace/test-project.deepnote', - projectId: 'test-project-id' - } - }; - - const mockFS = mock(); - when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(''))); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - - // Execute the method - await explorerView.renameNotebook(mockTreeItem as DeepnoteTreeItem); - - // Verify that readFile was not called (early return) - verify(mockFS.readFile(anything())).never(); - }); - test('should return early if user cancels input or provides same name', async () => { const projectId = 'test-project-id'; const notebookId = 'notebook-to-rename'; @@ -1503,28 +1197,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); }); - test('should return early if tree item type is not Notebook', async () => { - const mockTreeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/workspace/test-project.deepnote', - projectId: 'test-project-id' - } - }; - - const mockFS = mock(); - when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(''))); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - - // Execute the method - await explorerView.deleteNotebook(mockTreeItem as DeepnoteTreeItem); - - // Verify that readFile was not called (early return) - verify(mockFS.readFile(anything())).never(); - // Verify no warning message was shown - verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); - }); - test('should return early if user cancels confirmation', async () => { const projectId = 'test-project-id'; const notebookId = 'notebook-to-delete'; @@ -1706,26 +1378,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); }); - test('should return early if tree item type is not Notebook', async () => { - const mockTreeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/workspace/test-project.deepnote', - projectId: 'test-project-id' - } - }; - - const mockFS = mock(); - when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(''))); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - - // Execute the method - await explorerView.duplicateNotebook(mockTreeItem as DeepnoteTreeItem); - - // Verify that readFile was not called (early return) - verify(mockFS.readFile(anything())).never(); - }); - test('should show error if notebook is not found in project', async () => { const projectId = 'test-project-id'; const nonExistentNotebookId = 'non-existent-notebook-id'; @@ -2024,24 +1676,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { verify(mockFS.readFile(anything())).never(); }); - test('should return early if the tree item is a ProjectFile (not a group)', async () => { - const mockTreeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: '/workspace/test-project.deepnote', - projectId: 'test-project-id' - } - }; - - const mockFS = mock(); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - - await explorerView.renameProject(mockTreeItem as DeepnoteTreeItem); - - verify(mockedVSCodeNamespaces.window.showInputBox(anything())).never(); - verify(mockFS.readFile(anything())).never(); - }); - test('should return early if user cancels input or provides same name', async () => { const projectId = 'test-project-id'; const currentName = 'Current Project Name'; @@ -2266,58 +1900,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); }); - test('should export all notebooks when triggered from project', async () => { - resetVSCodeMocks(); - - const projectData: DeepnoteFile = { - version: '1.0.0', - metadata: { - createdAt: '2024-01-01T00:00:00.000Z', - modifiedAt: '2024-01-01T00:00:00.000Z' - }, - project: { - id: 'project-id', - name: 'Test Project', - notebooks: [ - { id: 'nb-1', name: 'Notebook 1', blocks: [], executionMode: 'block' }, - { id: 'nb-2', name: 'Notebook 2', blocks: [], executionMode: 'block' } - ] - } - }; - - const mockFS = mock(); - when(mockFS.readFile(anything())).thenReturn( - Promise.resolve(Buffer.from(serializeDeepnoteFile(projectData))) - ); - when(mockFS.stat(anything())).thenReject(new Error('File not found')); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - - when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( - Promise.resolve({ label: 'Jupyter Notebook (.ipynb)', value: 'jupyter' }) as any - ); - when(mockedVSCodeNamespaces.window.showOpenDialog(anything())).thenReturn( - Promise.resolve([Uri.file('/output/folder')]) - ); - when(mockedVSCodeNamespaces.window.showInformationMessage(anything())).thenReturn( - Promise.resolve(undefined) - ); - - let writeCount = 0; - when(mockFS.writeFile(anything(), anything())).thenCall(() => { - writeCount++; - return Promise.resolve(); - }); - - const treeItem = buildGroupTreeItem('project-id', [ - { filePath: '/test/project.deepnote', project: projectData } - ]); - - await (explorerView as any).exportProject(treeItem); - - assert.strictEqual(writeCount, 2); - verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); - }); - test('should write correct Jupyter notebook JSON format', async () => { resetVSCodeMocks(); diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index c122bbabc7..9ae6f366e7 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -3,7 +3,7 @@ import { assert } from 'chai'; import { mkdtempSync, rmSync } from 'fs'; import { tmpdir } from 'os'; import * as sinon from 'sinon'; -import { anything, instance, mock, reset, resetCalls, when } from 'ts-mockito'; +import { anything, instance, mock, resetCalls, when } from 'ts-mockito'; import { Disposable, EventEmitter, @@ -11,7 +11,6 @@ import { NotebookCellData, NotebookCellKind, NotebookDocument, - NotebookEdit, Uri } from 'vscode'; import { join } from '../../platform/vscode-path/path'; @@ -22,7 +21,6 @@ import type { DeepnoteOutput, DeepnoteProject } from '../../platform/deepnote/de import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; import type { IControllerRegistration } from '../controllers/types'; import { IDeepnoteNotebookManager } from '../types'; -import { DeepnoteDataConverter } from './deepnoteDataConverter'; import { DeepnoteFileChangeWatcher } from './deepnoteFileChangeWatcher'; import { SnapshotService } from './snapshots/snapshotService'; @@ -51,47 +49,6 @@ const validProject: DeepnoteFile = { } }; -const multiNotebookProject: DeepnoteFile = { - version: '1.0.0', - metadata: { createdAt: '2025-01-01T00:00:00Z' }, - project: { - id: 'project-1', - name: 'Multi Notebook Project', - notebooks: [ - { - id: 'notebook-1', - name: 'Notebook 1', - blocks: [ - { - id: 'block-nb1', - type: 'code', - sortingKey: 'a0', - blockGroup: '1', - content: 'print("nb1-new")', - metadata: {} - } - ] - }, - { - id: 'notebook-2', - name: 'Notebook 2', - blocks: [ - { - id: 'block-nb2', - type: 'code', - sortingKey: 'a0', - blockGroup: '1', - content: 'print("nb2-new")', - metadata: {} - } - ] - } - ] - } -}; - -const multiNotebookYaml = serializeDeepnoteFile(multiNotebookProject); - const waitForTimeoutMs = 5000; const waitForIntervalMs = 50; const debounceWaitMs = 800; @@ -99,17 +56,6 @@ const rapidChangeIntervalMs = 100; const autoSaveGraceMs = 200; const postSnapshotReadGraceMs = 100; -interface NotebookEditCapture { - uriKey: string; - cellSourceJoined: string; -} - -interface SnapshotInteractionCapture { - cellSourcesJoined: string; - outputPlainJoined: string; - uriKey: string; -} - /** * Polls until a condition is met or a timeout is reached. */ @@ -566,20 +512,20 @@ project: test('editor→external→editor→external: second external edit must reload (self-write leak regression)', async function () { this.timeout(15_000); const uri = testFileUri('self-write-leak.deepnote'); - - when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(multiNotebookProject); - - // Initial state: editor content matches disk — use the real converter - const converter = new DeepnoteDataConverter(); - const nb1 = multiNotebookProject.project.notebooks[0]; const notebook = createMockNotebook({ uri, - metadata: { deepnoteProjectId: multiNotebookProject.project.id, deepnoteNotebookId: nb1.id }, - cells: converter.convertBlocksToCells(nb1.blocks) + cells: [ + { + metadata: { id: 'block-1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")', languageId: 'python' } + } + ] }); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - setupMockFs(multiNotebookYaml); + setupMockFs(serializeDeepnoteFile(validProject)); // Real VS Code behavior: workspace.save() fires onDidSaveNotebookDocument. // executeMainFileSync calls markSelfWrite before workspace.save, AND the @@ -598,7 +544,7 @@ project: assert.strictEqual(applyEditCount, 0, 'Step 1: editor save FS event should be suppressed'); // Step 2: first external edit → triggers executeMainFileSync (reload works) - const externalProject1 = structuredClone(multiNotebookProject); + const externalProject1 = structuredClone(validProject); externalProject1.project.notebooks[0].blocks![0].content = 'print("external-1")'; setupMockFs(serializeDeepnoteFile(externalProject1)); @@ -619,7 +565,7 @@ project: await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); // Step 4: second external edit → should reload but phantom self-write blocks it - const externalProject2 = structuredClone(multiNotebookProject); + const externalProject2 = structuredClone(validProject); externalProject2.project.notebooks[0].blocks![0].content = 'print("external-2")'; setupMockFs(serializeDeepnoteFile(externalProject2)); @@ -1707,8 +1653,6 @@ project: return Promise.resolve(undefined); }); - when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(multiNotebookProject); - const uriNb1 = testFileUri('multi-scope-nb1.deepnote'); const uriNb2 = testFileUri('multi-scope-nb2.deepnote'); const notebook1 = createMockNotebook({ @@ -1758,841 +1702,5 @@ project: 'only the matching open notebook should receive snapshot edits' ); }); - - suite('snapshot and deserialization interaction', () => { - let interactionCaptures: SnapshotInteractionCapture[]; - let snapshotApplyEditStub: sinon.SinonStub; - - setup(() => { - interactionCaptures = []; - - reset(mockedNotebookManager); - when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(validProject); - when(mockedNotebookManager.updateOriginalProject(anything(), anything(), anything())).thenReturn(); - resetCalls(mockedNotebookManager); - - snapshotApplyEditStub = sinon.stub(snapshotWatcher, 'applyNotebookEdits').callsFake(async function ( - this: DeepnoteFileChangeWatcher, - ...args: unknown[] - ) { - const uri = args[0] as Uri; - const edits = args[1] as NotebookEdit[]; - - const replaceCellsEdit = edits.find((e) => (e as { newCells?: unknown[] }).newCells?.length); - if (replaceCellsEdit) { - const newCells = ( - replaceCellsEdit as { - newCells: Array<{ - outputs?: Array<{ items: Array<{ data?: Uint8Array }> }>; - value: string; - }>; - } - ).newCells; - const outputPlainJoined = newCells - .map((c) => { - const data = c.outputs?.[0]?.items?.[0]?.data; - - return data ? new TextDecoder().decode(data) : ''; - }) - .filter(Boolean) - .join(';'); - - interactionCaptures.push({ - uriKey: uri.toString(), - cellSourcesJoined: newCells.map((c) => c.value).join('\n'), - outputPlainJoined - }); - } - - return DeepnoteFileChangeWatcher.prototype.applyNotebookEdits.apply(this, [uri, edits]); - }); - }); - - teardown(() => { - snapshotApplyEditStub.restore(); - }); - - test('snapshot change with multi-notebook project applies only matching block outputs per notebook', async () => { - when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(multiNotebookProject); - - const multiOutputs = new Map([ - [ - 'block-nb1', - [ - { - output_type: 'execute_result', - data: { 'text/plain': 'OutputForNb1Only' }, - execution_count: 1 - } as DeepnoteOutput - ] - ], - [ - 'block-nb2', - [ - { - output_type: 'execute_result', - data: { 'text/plain': 'OutputForNb2Only' }, - execution_count: 1 - } as DeepnoteOutput - ] - ] - ]); - when(mockSnapshotService.readSnapshot(anything())).thenCall(() => { - readSnapshotCallCount++; - - return Promise.resolve(multiOutputs); - }); - - const basePath = testFileUri('multi-snap.deepnote'); - const uriNb1 = basePath.with({ query: 'view=1' }); - const uriNb2 = basePath.with({ query: 'view=2' }); - - const notebook1 = createMockNotebook({ - uri: uriNb1, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-1' - }, - cells: [ - { - metadata: { id: 'block-nb1', type: 'code' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb1-old")', languageId: 'python' } - } - ] - }); - - const notebook2 = createMockNotebook({ - uri: uriNb2, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-2' - }, - cells: [ - { - metadata: { id: 'block-nb2', type: 'code' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb2-old")', languageId: 'python' } - } - ] - }); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - - const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); - snapshotOnDidChange.fire(snapshotUri); - - await waitFor(() => snapshotApplyEditCount >= 4); - - const byUri = new Map(interactionCaptures.map((c) => [c.uriKey, c])); - - assert.include(byUri.get(uriNb1.toString())?.outputPlainJoined ?? '', 'OutputForNb1Only'); - assert.notInclude(byUri.get(uriNb1.toString())?.outputPlainJoined ?? '', 'OutputForNb2Only'); - - assert.include(byUri.get(uriNb2.toString())?.outputPlainJoined ?? '', 'OutputForNb2Only'); - assert.notInclude(byUri.get(uriNb2.toString())?.outputPlainJoined ?? '', 'OutputForNb1Only'); - }); - - test('main file change after snapshot update deserializes updated cell source', async function () { - this.timeout(10_000); - const notebookUri = testFileUri('after-snap.deepnote'); - const notebook = createMockNotebook({ - uri: notebookUri, - cells: [ - { - metadata: { id: 'block-1', type: 'code' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("hello")', languageId: 'python' } - } - ] - }); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - - const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); - when(mockSnapshotService.readSnapshot(anything())).thenCall(() => { - readSnapshotCallCount++; - - return Promise.resolve(snapshotOutputs); - }); - - snapshotOnDidChange.fire(snapshotUri); - await waitFor(() => snapshotApplyEditCount >= 2); - - const changedYaml = ` -version: '1.0.0' -metadata: - createdAt: '2025-01-01T00:00:00Z' -project: - id: project-1 - name: Test Project - notebooks: - - id: notebook-1 - name: Notebook 1 - blocks: - - id: block-1 - type: code - sortingKey: a0 - blockGroup: '1' - content: print("after-snapshot-main-sync") -`; - - const mockFs = mock(); - when(mockFs.readFile(anything())).thenCall(() => { - return Promise.resolve(new TextEncoder().encode(changedYaml)); - }); - when(mockFs.writeFile(anything(), anything())).thenReturn(Promise.resolve()); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); - - // Snapshot save marks self-write; first FS event consumes it without reloading. - snapshotOnDidChange.fire(notebookUri); - await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); - - snapshotOnDidChange.fire(notebookUri); - - await waitFor(() => - interactionCaptures.some((c) => c.cellSourcesJoined.includes('after-snapshot-main-sync')) - ); - - const mainSyncCapture = interactionCaptures.find((c) => - c.cellSourcesJoined.includes('after-snapshot-main-sync') - ); - - assert.isDefined(mainSyncCapture); - assert.include( - mainSyncCapture!.cellSourcesJoined, - 'after-snapshot-main-sync', - 'main-file sync should deserialize new source after snapshot outputs were applied' - ); - }); - - test('snapshot save self-write is consumed once then external main-file change applies', async function () { - this.timeout(10_000); - const baseUri = testFileUri('snap-self-write.deepnote'); - const notebook = createMockNotebook({ - uri: baseUri, - cells: [ - { - metadata: { id: 'block-1', type: 'code' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("hello")', languageId: 'python' } - } - ] - }); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - - const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); - snapshotOnDidChange.fire(snapshotUri); - await waitFor(() => snapshotSaveCount >= 1); - - const editsBefore = snapshotApplyEditCount; - - snapshotOnDidChange.fire(baseUri); - await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); - - assert.strictEqual( - snapshotApplyEditCount, - editsBefore, - 'first main-file FS event after snapshot save should be consumed as self-write' - ); - - const externalYaml = ` -version: '1.0.0' -metadata: - createdAt: '2025-01-01T00:00:00Z' -project: - id: project-1 - name: Test Project - notebooks: - - id: notebook-1 - name: Notebook 1 - blocks: - - id: block-1 - type: code - sortingKey: a0 - blockGroup: '1' - content: print("external-after-self-write") -`; - - const mockFs = mock(); - when(mockFs.readFile(anything())).thenCall(() => { - return Promise.resolve(new TextEncoder().encode(externalYaml)); - }); - when(mockFs.writeFile(anything(), anything())).thenReturn(Promise.resolve()); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); - - snapshotOnDidChange.fire(baseUri); - - await waitFor(() => - interactionCaptures.some((c) => c.cellSourcesJoined.includes('external-after-self-write')) - ); - - assert.isTrue( - interactionCaptures.some((c) => c.cellSourcesJoined.includes('external-after-self-write')), - 'second main-file change should deserialize external content' - ); - }); - - test('main-file sync runs after in-flight snapshot when both are triggered close together', async function () { - this.timeout(12_000); - let releaseSnapshot!: () => void; - const snapshotGate = new Promise((resolve) => { - releaseSnapshot = resolve; - }); - - when(mockSnapshotService.readSnapshot(anything())).thenCall(() => { - readSnapshotCallCount++; - - return snapshotGate.then(() => snapshotOutputs); - }); - - const baseUri = testFileUri('coalesce.deepnote'); - const notebook = createMockNotebook({ - uri: baseUri, - cells: [ - { - metadata: { id: 'block-1', type: 'code' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("hello")', languageId: 'python' } - } - ] - }); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - - const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); - snapshotOnDidChange.fire(snapshotUri); - - await waitFor(() => readSnapshotCallCount >= 1); - - const coalescedYaml = ` -version: '1.0.0' -metadata: - createdAt: '2025-01-01T00:00:00Z' -project: - id: project-1 - name: Test Project - notebooks: - - id: notebook-1 - name: Notebook 1 - blocks: - - id: block-1 - type: code - sortingKey: a0 - blockGroup: '1' - content: print("main-wins-after-snapshot") -`; - - const mockFs = mock(); - when(mockFs.readFile(anything())).thenCall(() => { - return Promise.resolve(new TextEncoder().encode(coalescedYaml)); - }); - when(mockFs.writeFile(anything(), anything())).thenReturn(Promise.resolve()); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); - - snapshotOnDidChange.fire(baseUri); - await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); - - releaseSnapshot(); - - await waitFor(() => - interactionCaptures.some((c) => c.cellSourcesJoined.includes('main-wins-after-snapshot')) - ); - - assert.isTrue( - interactionCaptures.some((c) => c.cellSourcesJoined.includes('main-wins-after-snapshot')), - 'main-file sync should deserialize disk YAML after snapshot operation completes' - ); - }); - - // Multi-notebook test removed — multi-notebook support has been replaced by auto-splitting into separate files - test.skip('multi-notebook: snapshot outputs then external YAML update keeps per-notebook sources', async function () { - this.timeout(12_000); - when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(multiNotebookProject); - - const multiOutputs = new Map([ - [ - 'block-nb1', - [ - { - output_type: 'execute_result', - data: { 'text/plain': 'SnapNb1' }, - execution_count: 1 - } as DeepnoteOutput - ] - ], - [ - 'block-nb2', - [ - { - output_type: 'execute_result', - data: { 'text/plain': 'SnapNb2' }, - execution_count: 1 - } as DeepnoteOutput - ] - ] - ]); - when(mockSnapshotService.readSnapshot(anything())).thenCall(() => { - readSnapshotCallCount++; - - return Promise.resolve(multiOutputs); - }); - - const basePath = testFileUri('multi-snap-then-yaml.deepnote'); - const uriNb1 = basePath.with({ query: 'view=1' }); - const uriNb2 = basePath.with({ query: 'view=2' }); - - const notebook1 = createMockNotebook({ - uri: uriNb1, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-1' - }, - cells: [ - { - metadata: { id: 'block-nb1', type: 'code' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb1-old")', languageId: 'python' } - } - ] - }); - - const notebook2 = createMockNotebook({ - uri: uriNb2, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-2' - }, - cells: [ - { - metadata: { id: 'block-nb2', type: 'code' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb2-old")', languageId: 'python' } - } - ] - }); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - - const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); - snapshotOnDidChange.fire(snapshotUri); - await waitFor(() => snapshotApplyEditCount >= 4); - - const round2Project = structuredClone(multiNotebookProject); - round2Project.project.notebooks[0].blocks![0].content = 'print("nb1-round2")'; - round2Project.project.notebooks[1].blocks![0].content = 'print("nb2-round2")'; - const yamlRound2 = serializeDeepnoteFile(round2Project); - - const mockFs = mock(); - when(mockFs.readFile(anything())).thenCall(() => { - return Promise.resolve(new TextEncoder().encode(yamlRound2)); - }); - when(mockFs.writeFile(anything(), anything())).thenReturn(Promise.resolve()); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); - - // Two snapshot saves increment self-write count to 2 for the shared base file URI. - snapshotOnDidChange.fire(basePath); - await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); - snapshotOnDidChange.fire(basePath); - await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); - - snapshotOnDidChange.fire(basePath); - - await waitFor(() => - interactionCaptures.some( - (c) => c.uriKey === uriNb1.toString() && c.cellSourcesJoined.includes('nb1-round2') - ) - ); - - assert.isTrue( - interactionCaptures.some( - (c) => c.uriKey === uriNb1.toString() && c.outputPlainJoined.includes('SnapNb1') - ), - 'snapshot phase should apply SnapNb1 output to notebook-1' - ); - assert.isTrue( - interactionCaptures.some( - (c) => c.uriKey === uriNb2.toString() && c.outputPlainJoined.includes('SnapNb2') - ), - 'snapshot phase should apply SnapNb2 output to notebook-2' - ); - - const nb1Main = interactionCaptures.find( - (c) => c.uriKey === uriNb1.toString() && c.cellSourcesJoined.includes('nb1-round2') - ); - const nb2Main = interactionCaptures.find( - (c) => c.uriKey === uriNb2.toString() && c.cellSourcesJoined.includes('nb2-round2') - ); - - assert.isDefined(nb1Main); - assert.include(nb1Main!.cellSourcesJoined, 'nb1-round2'); - assert.notInclude(nb1Main!.cellSourcesJoined, 'nb2-round2'); - - assert.isDefined(nb2Main); - assert.include(nb2Main!.cellSourcesJoined, 'nb2-round2'); - assert.notInclude(nb2Main!.cellSourcesJoined, 'nb1-round2'); - }); - - test('snapshot outputs for sibling notebook blocks do not leak into a single open notebook', async () => { - when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(multiNotebookProject); - - const multiOutputs = new Map([ - [ - 'block-nb1', - [ - { - output_type: 'execute_result', - data: { 'text/plain': 'OnlyNb1' }, - execution_count: 1 - } as DeepnoteOutput - ] - ], - [ - 'block-nb2', - [ - { - output_type: 'execute_result', - data: { 'text/plain': 'LeakIfApplied' }, - execution_count: 1 - } as DeepnoteOutput - ] - ] - ]); - when(mockSnapshotService.readSnapshot(anything())).thenCall(() => { - readSnapshotCallCount++; - - return Promise.resolve(multiOutputs); - }); - - const uriNb1 = testFileUri('only-nb1.deepnote'); - const notebook1Only = createMockNotebook({ - uri: uriNb1, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-1' - }, - cells: [ - { - metadata: { id: 'block-nb1', type: 'code' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb1-only")', languageId: 'python' } - } - ] - }); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1Only]); - - const snapshotUri = testFileUri('snapshots', 'my-project_project-1_latest.snapshot.deepnote'); - snapshotOnDidChange.fire(snapshotUri); - - await waitFor(() => snapshotApplyEditCount >= 2); - - const cap = interactionCaptures.find((c) => c.uriKey === uriNb1.toString()); - - assert.isDefined(cap); - assert.include(cap!.outputPlainJoined, 'OnlyNb1'); - assert.notInclude(cap!.outputPlainJoined, 'LeakIfApplied'); - }); - }); - }); - - // Multi-notebook file sync tests removed — multi-notebook support has been replaced by auto-splitting into separate files - suite.skip('multi-notebook file sync', () => { - let workspaceSetCaptures: NotebookEditCapture[] = []; - - setup(() => { - reset(mockedNotebookManager); - when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(multiNotebookProject); - when(mockedNotebookManager.updateOriginalProject(anything(), anything(), anything())).thenReturn(); - resetCalls(mockedNotebookManager); - workspaceSetCaptures = []; - sinon.stub(watcher, 'applyNotebookEdits' as any).callsFake(async (...args: unknown[]) => { - const uri = args[0] as Uri; - const edits = args[1] as NotebookEdit[]; - - applyEditCount++; - - 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 true; - }); - }); - - test('should reload each notebook with its own content when multiple notebooks are open', async () => { - const basePath = testFileUri('multi.deepnote'); - const uriNb1 = basePath.with({ query: 'view=1' }); - const uriNb2 = basePath.with({ query: 'view=2' }); - - const notebook1 = createMockNotebook({ - uri: uriNb1, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-1' - }, - cells: [ - { - metadata: { id: 'block-nb1', type: 'code' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb1-old")', languageId: 'python' } - } - ] - }); - - const notebook2 = createMockNotebook({ - uri: uriNb2, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-2' - }, - cells: [ - { - metadata: { id: 'block-nb2', type: 'code' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb2-old")', languageId: 'python' } - } - ] - }); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - setupMockFs(multiNotebookYaml); - - onDidChangeFile.fire(basePath); - - await waitFor(() => applyEditCount >= 2); - - assert.strictEqual(applyEditCount, 2, 'applyEdit should run once per open notebook'); - assert.strictEqual(workspaceSetCaptures.length, 2, 'each notebook should get a replaceCells edit'); - - const byUri = new Map(workspaceSetCaptures.map((c) => [c.uriKey, c.cellSourceJoined])); - - assert.include(byUri.get(uriNb1.toString()) ?? '', 'nb1-new'); - assert.notInclude(byUri.get(uriNb1.toString()) ?? '', 'nb2-new'); - assert.include(byUri.get(uriNb2.toString()) ?? '', 'nb2-new'); - assert.notInclude(byUri.get(uriNb2.toString()) ?? '', 'nb1-new'); - }); - - test('should not clear notebook selection before processing file change', async () => { - const basePath = testFileUri('multi.deepnote'); - const uriNb1 = basePath.with({ query: 'a=1' }); - const uriNb2 = basePath.with({ query: 'b=2' }); - - const notebook1 = createMockNotebook({ - uri: uriNb1, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-1' - }, - cells: [ - { - metadata: { id: 'block-nb1' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb1-old")' } - } - ] - }); - - const notebook2 = createMockNotebook({ - uri: uriNb2, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-2' - }, - cells: [ - { - metadata: { id: 'block-nb2' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb2-old")' } - } - ] - }); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - setupMockFs(multiNotebookYaml); - - onDidChangeFile.fire(basePath); - - await waitFor(() => applyEditCount >= 2); - }); - - test('should not corrupt other notebooks when one notebook triggers a file change', async () => { - const basePath = testFileUri('multi.deepnote'); - const uriNb1 = basePath.with({ query: 'n=1' }); - const uriNb2 = basePath.with({ query: 'n=2' }); - - const notebook1 = createMockNotebook({ - uri: uriNb1, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-1' - }, - cells: [ - { - metadata: { id: 'block-nb1' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb1-old")' } - } - ] - }); - - const notebook2 = createMockNotebook({ - uri: uriNb2, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-2' - }, - cells: [ - { - metadata: { id: 'block-nb2' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb2-old")' } - } - ] - }); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - setupMockFs(multiNotebookYaml); - - onDidChangeFile.fire(basePath); - - await waitFor(() => applyEditCount >= 2); - - const nb1Cells = workspaceSetCaptures.find((c) => c.uriKey === uriNb1.toString())?.cellSourceJoined; - const nb2Cells = workspaceSetCaptures.find((c) => c.uriKey === uriNb2.toString())?.cellSourceJoined; - - assert.isDefined(nb1Cells); - assert.isDefined(nb2Cells); - assert.notStrictEqual(nb1Cells, nb2Cells, 'each open notebook must receive distinct deserialized content'); - - assert.include(nb1Cells!, 'nb1-new'); - assert.include(nb2Cells!, 'nb2-new'); - assert.notInclude(nb1Cells!, 'nb2-new', 'notebook-1 must not receive notebook-2 block content'); - assert.notInclude(nb2Cells!, 'nb1-new', 'notebook-2 must not receive notebook-1 block content'); - }); - - test('external edit to disk should update each open notebook and not be suppressed as self-write', async () => { - const basePath = testFileUri('multi-external.deepnote'); - const uriNb1 = basePath.with({ query: 'view=1' }); - const uriNb2 = basePath.with({ query: 'view=2' }); - - const notebook1 = createMockNotebook({ - uri: uriNb1, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-1' - }, - cells: [ - { - metadata: { id: 'block-nb1', type: 'code' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb1-old")', languageId: 'python' } - } - ] - }); - - const notebook2 = createMockNotebook({ - uri: uriNb2, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-2' - }, - cells: [ - { - metadata: { id: 'block-nb2', type: 'code' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb2-old")', languageId: 'python' } - } - ] - }); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - setupMockFs(multiNotebookYaml); - - onDidChangeFile.fire(basePath); - - await waitFor(() => applyEditCount >= 2); - - const byUri = new Map(workspaceSetCaptures.map((c) => [c.uriKey, c.cellSourceJoined])); - - assert.include(byUri.get(uriNb1.toString()) ?? '', 'nb1-new'); - assert.include(byUri.get(uriNb2.toString()) ?? '', 'nb2-new'); - }); - - test('external edit after a user save should still be processed', async function () { - this.timeout(8000); - const basePath = testFileUri('multi-save-then-external.deepnote'); - const uriNb1 = basePath.with({ query: 'view=1' }); - const uriNb2 = basePath.with({ query: 'view=2' }); - - const notebook1 = createMockNotebook({ - uri: uriNb1, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-1' - }, - cells: [ - { - metadata: { id: 'block-nb1' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb1-old")' } - } - ] - }); - - const notebook2 = createMockNotebook({ - uri: uriNb2, - metadata: { - deepnoteProjectId: 'project-1', - deepnoteNotebookId: 'notebook-2' - }, - cells: [ - { - metadata: { id: 'block-nb2' }, - outputs: [], - kind: NotebookCellKind.Code, - document: { getText: () => 'print("nb2-old")' } - } - ] - }); - - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook1, notebook2]); - setupMockFs(multiNotebookYaml); - - onDidSaveNotebook.fire(notebook1); - onDidChangeFile.fire(basePath); - - await new Promise((resolve) => setTimeout(resolve, debounceWaitMs)); - - assert.strictEqual(applyEditCount, 0, 'first FS event after save should be suppressed as self-write'); - - onDidChangeFile.fire(basePath); - - await waitFor(() => applyEditCount >= 1); - - assert.isAtLeast(applyEditCount, 1); - }); }); });