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/package.json b/package.json index dfba8b252b..05af26b156 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%", @@ -1594,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/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", diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 112f7f0e09..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,12 +277,12 @@ 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, - serverInfo?.jupyterPort ?? 0, + 0, 'unknown', capturedOutput?.stdout || '', capturedOutput?.stderr || '', @@ -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,13 +341,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); - this.serverOutputByFile.delete(fileKey); + this.serverOutputByFile.delete(projectKey); - const disposables = this.disposablesByFile.get(fileKey); - if (disposables) { - disposables.forEach((d) => d.dispose()); - this.disposablesByFile.delete(fileKey); - } + this.disposeOutputListeners(projectKey); } /** @@ -377,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); @@ -400,19 +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(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); } @@ -429,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); } @@ -446,6 +441,22 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } + /** + * 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(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 project ${projectKey}`, ex); + } + } + this.disposablesByFile.delete(projectKey); + } + public async dispose(): Promise { logger.info('Disposing DeepnoteServerStarter - stopping all servers...'); @@ -460,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); }) ); } @@ -485,12 +496,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 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/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/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..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, @@ -14,6 +15,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 +29,7 @@ import { DeepnoteKernelConnectionMetadata, IDeepnoteEnvironmentManager, IDeepnoteKernelAutoSelector, - IDeepnoteNotebookEnvironmentMapper + IDeepnoteProjectEnvironmentMapper } from '../types'; import { CreateDeepnoteEnvironmentOptions, DeepnoteEnvironment } from './deepnoteEnvironment'; import { DeepnoteEnvironmentTreeDataProvider } from './deepnoteEnvironmentTreeDataProvider.node'; @@ -49,8 +51,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 +302,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 +347,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,13 +382,19 @@ 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 currentEnvironment = currentEnvironmentId - ? this.environmentManager.getEnvironment(currentEnvironmentId) + const previousEnvironmentId = this.projectEnvironmentMapper.getEnvironmentForProject(projectId); + const currentEnvironment = previousEnvironmentId + ? this.environmentManager.getEnvironment(previousEnvironmentId) : undefined; // Get all environments @@ -429,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) { @@ -437,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); @@ -461,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( @@ -471,15 +513,71 @@ export class DeepnoteEnvironmentsView implements Disposable { cancellable: true }, async (progress, token) => { - // Update the notebook-to-environment mapping - await this.notebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, 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 0407420162..e1f9a79b3e 100644 --- a/src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts +++ b/src/kernels/deepnote/environments/deepnoteEnvironmentsView.unit.test.ts @@ -3,10 +3,10 @@ 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'; +import { IKernel, IKernelProvider, INotebookKernelExecution } from '../../../kernels/types'; import { DeepnoteEnvironment } from './deepnoteEnvironment'; import { PythonEnvironment } from '../../../platform/pythonEnvironments/info'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; @@ -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,42 @@ suite('DeepnoteEnvironmentsView', () => { lastUsedAt: new Date() }; + 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); - 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 +851,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 +861,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 +885,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 +895,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 +919,8 @@ suite('DeepnoteEnvironmentsView', () => { const mockNotebook = { uri: notebookUri, notebookType: 'deepnote', - cellCount: 5 + cellCount: 5, + metadata: { deepnoteProjectId: testProjectId } }; const mockNotebookEditor = { notebook: mockNotebook @@ -902,8 +929,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 +968,7 @@ suite('DeepnoteEnvironmentsView', () => { // Mock environment mapping update when( - mockNotebookEnvironmentMapper.setEnvironmentForNotebook(baseFileUri, newExternalEnvironment.id) + mockProjectEnvironmentMapper.setEnvironmentForProject(testProjectId, newExternalEnvironment.id) ).thenResolve(); // Mock controller rebuild @@ -952,12 +978,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(); @@ -965,13 +991,122 @@ 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'); const mockNotebook = { uri: notebookUri, notebookType: 'deepnote', - cellCount: 5 + cellCount: 5, + metadata: { deepnoteProjectId: testProjectId } }; const mockNotebookEditor = { notebook: mockNotebook @@ -980,8 +1115,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 +1150,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 +1159,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..6049cdfe75 --- /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, unparsable) + * 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/deepnoteActivationService.ts b/src/notebooks/deepnote/deepnoteActivationService.ts index 96c35288cd..2eb47f6a4f 100644 --- a/src/notebooks/deepnote/deepnoteActivationService.ts +++ b/src/notebooks/deepnote/deepnoteActivationService.ts @@ -1,14 +1,16 @@ +import { deserializeDeepnoteFile } from '@deepnote/blocks'; 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'; 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,13 +48,25 @@ 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(); this.registerSerializer(); this.extensionContext.subscriptions.push(this.editProtection); + this.extensionContext.subscriptions.push( + workspace.onDidOpenNotebookDocument((doc) => { + 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')) { @@ -70,6 +86,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/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 new file mode 100644 index 0000000000..6bb5487c6c --- /dev/null +++ b/src/notebooks/deepnote/deepnoteAutoSplitter.ts @@ -0,0 +1,241 @@ +import { serializeDeepnoteFile, type DeepnoteFile } from '@deepnote/blocks'; +import { l10n, RelativePattern, Uri, window, workspace } from 'vscode'; + +import { logger } from '../../platform/logging'; +import { + buildSiblingNotebookFileUri, + buildSingleNotebookFile, + computeSnapshotHash +} from './deepnoteNotebookFileFactory'; +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 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 newFileUri = await buildSiblingNotebookFileUri(fileUri, notebook.name, exists); + + reservedPaths.add(newFileUri.path); + + const newProject = await buildSingleNotebookFile(deepnoteFile, notebook); + + const yaml = serializeDeepnoteFile(newProject); + await workspace.fs.writeFile(newFileUri, new TextEncoder().encode(yaml)); + + newFiles.push(newFileUri); + + const newFileName = newFileUri.path.split('/').pop(); + 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 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 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); + } +} 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'); + }); + }); +}); diff --git a/src/notebooks/deepnote/deepnoteExplorerView.ts b/src/notebooks/deepnote/deepnoteExplorerView.ts index 33599da2f0..a2fb193002 100644 --- a/src/notebooks/deepnote/deepnoteExplorerView.ts +++ b/src/notebooks/deepnote/deepnoteExplorerView.ts @@ -1,16 +1,25 @@ 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'; import { IExtensionContext } from '../../platform/common/types'; -import { IDeepnoteNotebookManager } from '../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'; +import { SNAPSHOT_FILE_SUFFIX } from './snapshots/snapshotFiles'; import { ILogger } from '../../platform/logging/types'; /** @@ -18,18 +27,24 @@ import { ILogger } from '../../platform/logging/types'; */ @injectable() export class DeepnoteExplorerView { + private readonly logger: ILogger; private readonly treeDataProvider: DeepnoteTreeDataProvider; private treeView: TreeView; constructor( @inject(IExtensionContext) private readonly extensionContext: IExtensionContext, - @inject(IDeepnoteNotebookManager) private readonly manager: IDeepnoteNotebookManager, @inject(ILogger) logger: ILogger ) { + this.logger = logger; 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, @@ -43,25 +58,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; @@ -70,46 +87,58 @@ 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, newNotebook.id); + 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 }; } 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 @@ -143,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 confirmation = await window.showWarningMessage( - l10n.t('Are you sure you want to delete notebook "{0}"?', notebookName), - { modal: true }, - l10n.t('Delete') - ); + const target = this.resolveNotebookTarget(treeItem); - 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) { @@ -169,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) { @@ -191,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) { @@ -207,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)) { @@ -225,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); @@ -259,10 +322,8 @@ export class DeepnoteExplorerView { await this.treeDataProvider.refreshNotebook(treeItem.context.projectId); - // Optionally open the duplicated notebook - this.manager.selectNotebookForProject(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 @@ -276,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'), @@ -299,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'; @@ -412,16 +476,75 @@ export class DeepnoteExplorerView { } /** - * Generates a suggested unique notebook name based on existing notebooks - * @param projectData The project data containing existing notebooks - * @returns A unique suggested notebook name + * 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 generateSuggestedNotebookName(projectData: DeepnoteFile): string { - const notebookCount = projectData.project.notebooks?.length || 0; - const existingNames = new Set(projectData.project.notebooks?.map((nb: DeepnoteNotebook) => nb.name) || []); + 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; + }) + }; + } - let nextNumber = notebookCount + 1; + /** + * 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). + * @param existingNames The set of already-taken notebook names + * @returns A unique suggested notebook name + */ + private generateSuggestedNotebookName(existingNames: Set): string { + let nextNumber = existingNames.size + 1; let suggestedName = `Notebook ${nextNumber}`; + while (existingNames.has(suggestedName)) { nextNumber++; suggestedName = `Notebook ${nextNumber}`; @@ -483,38 +606,46 @@ 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, - notebookId: string - ): 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 - use granular refresh for notebooks - await this.treeDataProvider.refreshNotebook(projectData.project.id); + for (const file of projectFiles) { + try { + const project = await readDeepnoteProjectFile(file); - // Open the new notebook - this.manager.selectNotebookForProject(projectData.project.id, notebookId); - const notebookUri = fileUri.with({ query: `notebook=${notebookId}` }); - const document = await workspace.openNotebookDocument(notebookUri); - 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 { + this.treeDataProvider.refresh(); } private refreshExplorer(): void { @@ -522,29 +653,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.manager.selectNotebookForProject(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 @@ -590,7 +704,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 }); @@ -701,10 +815,7 @@ export class DeepnoteExplorerView { this.treeDataProvider.refresh(); - this.manager.selectNotebookForProject(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, @@ -725,12 +836,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 @@ -930,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), @@ -948,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) { @@ -959,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) { @@ -979,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') @@ -996,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, @@ -1017,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[] = []; @@ -1065,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') @@ -1082,7 +1215,6 @@ export class DeepnoteExplorerView { return; } - const fileUri = Uri.file(treeItem.context.filePath); const projectData = await readDeepnoteProjectFile(fileUri); if (!projectData?.project) { @@ -1103,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 5e5e2de826..3a832f5d2d 100644 --- a/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteExplorerView.unit.test.ts @@ -6,8 +6,13 @@ 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 { + DeepnoteTreeItem, + DeepnoteTreeItemType, + NOTEBOOK_FILE_CONTEXT_VALUE, + 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'; @@ -39,187 +44,9 @@ function createUuidMock(uuids: string[]): sinon.SinonStub { return stub; } -suite('DeepnoteExplorerView', () => { - let explorerView: DeepnoteExplorerView; - let mockExtensionContext: IExtensionContext; - let manager: DeepnoteNotebookManager; - let mockLogger: ILogger; - - setup(() => { - mockExtensionContext = { - subscriptions: [] - } as any; - - manager = new DeepnoteNotebookManager(); - mockLogger = createMockLogger(); - explorerView = new DeepnoteExplorerView(mockExtensionContext, manager, mockLogger); - }); - - 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 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); - - // 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); - }); - - 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; - let mockManager: DeepnoteNotebookManager; let sandbox: sinon.SinonSandbox; let uuidStubs: sinon.SinonStub[] = []; @@ -232,12 +59,12 @@ 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(() => { + explorerView.dispose(); sandbox.restore(); uuidStubs.forEach((stub) => stub.restore()); uuidStubs = []; @@ -481,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'); @@ -561,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') }; @@ -623,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')]; @@ -691,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') }; @@ -749,33 +496,10 @@ 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', () => { - 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'; @@ -783,6 +507,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 = { @@ -807,13 +532,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)); @@ -821,20 +547,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 @@ -842,25 +572,43 @@ 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 return null if user cancels notebook name input', async () => { + 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'; - // Mock existing project data const existingProjectData: DeepnoteFile = { version: '1.0.0', metadata: { @@ -870,34 +618,95 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { project: { id: projectId, name: 'Test Project', - notebooks: [] + 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); - // Mock file system const mockFS = mock(); when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(yamlContent))); - when(mockFS.writeFile(anything(), anything())).thenReturn(Promise.resolve()); + 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)); - // Mock user cancelling input - when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(notebookName)); - // Execute the method - const result = await explorerView.createAndAddNotebookToProject(fileUri); + const uuidStub = createUuidMock([newNotebookId, blockGroupId, blockId]); + uuidStubs.push(uuidStub); - // Verify result is null and file was not written - expect(result).to.be.null; - verify(mockFS.writeFile(anything(), anything())).never(); + 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 generate unique notebook name suggestions', async () => { + 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'; - // Mock existing project data with multiple notebooks const existingProjectData: DeepnoteFile = { version: '1.0.0', metadata: { @@ -907,57 +716,61 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { project: { id: projectId, name: 'Test Project', - notebooks: [ - { id: 'nb1', name: 'Notebook 1', blocks: [], executionMode: 'block' }, - { id: 'nb2', name: 'Notebook 2', blocks: [], executionMode: 'block' } - ] + notebooks: [{ id: 'existing', name: 'Existing', blocks: [], executionMode: 'block' }] } }; const yamlContent = serializeDeepnoteFile(existingProjectData); - // Mock file system const mockFS = mock(); when(mockFS.readFile(anything())).thenReturn(Promise.resolve(Buffer.from(yamlContent))); - when(mockFS.writeFile(anything(), anything())).thenReturn(Promise.resolve()); - when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); - let capturedInputBoxOptions: any; - when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenCall((options: any) => { - capturedInputBoxOptions = options; - return Promise.resolve('Test Notebook'); + // 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 uuidStub = createUuidMock(['test-id', 'test-id', 'test-id']); + 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({} as any) - ); + 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) ); - // Execute the method await explorerView.createAndAddNotebookToProject(fileUri); - // Verify suggested name is 'Notebook 3' (next in sequence) - expect(capturedInputBoxOptions).to.exist; - expect(capturedInputBoxOptions.value).to.equal('Notebook 3'); + // 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); }); - }); - suite('renameNotebook', () => { - test('should successfully rename a notebook with valid input', async () => { + test('should validate name uniqueness across sibling project files', async () => { const projectId = 'test-project-id'; - const notebookId = 'notebook-to-rename'; - const otherNotebookId = 'other-notebook-id'; - const oldName = 'Old Notebook Name'; - const newName = 'New Notebook Name'; const fileUri = Uri.file('/workspace/test-project.deepnote'); + const siblingUri = Uri.file('/workspace/test-project_other.deepnote'); - // Mock existing project data - const existingProjectData: DeepnoteFile = { + const sourceData: DeepnoteFile = { version: '1.0.0', metadata: { createdAt: '2024-01-01T00:00:00.000Z', @@ -966,13 +779,185 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { project: { id: projectId, name: 'Test Project', - notebooks: [ - { - id: otherNotebookId, - name: 'Other Notebook', - blocks: [], - executionMode: 'block' - }, + 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'); + + 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: [] + } + }; + + 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')); + when(mockFS.writeFile(anything(), anything())).thenReturn(Promise.resolve()); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFS)); + + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const result = await explorerView.createAndAddNotebookToProject(fileUri); + + // Null result and no file writes occurred (neither original nor sibling) + expect(result).to.be.null; + verify(mockFS.writeFile(anything(), anything())).never(); + }); + + test('should generate unique notebook name suggestions', async () => { + const projectId = 'test-project-id'; + const fileUri = Uri.file('/workspace/test-project.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: 'nb1', name: 'Notebook 1', blocks: [], executionMode: 'block' }, + { id: 'nb2', name: 'Notebook 2', blocks: [], executionMode: 'block' } + ] + } + }; + + const yamlContent = serializeDeepnoteFile(existingProjectData); + + // 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)); + + let capturedInputBoxOptions: any; + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenCall((options: any) => { + capturedInputBoxOptions = options; + return Promise.resolve('Test Notebook'); + }); + + const uuidStub = createUuidMock(['test-id', 'test-id', 'test-id']); + uuidStubs.push(uuidStub); + + when(mockedVSCodeNamespaces.workspace.openNotebookDocument(anything())).thenReturn( + Promise.resolve({} as any) + ); + when(mockedVSCodeNamespaces.window.showNotebookDocument(anything(), anything())).thenReturn( + Promise.resolve(undefined as any) + ); + + await explorerView.createAndAddNotebookToProject(fileUri); + + // With two existing notebooks, suggestion is `Notebook ${size + 1}` = 'Notebook 3' + expect(capturedInputBoxOptions).to.exist; + expect(capturedInputBoxOptions.value).to.equal('Notebook 3'); + }); + }); + + suite('renameNotebook', () => { + test('should successfully rename a notebook with valid input', async () => { + const projectId = 'test-project-id'; + const notebookId = 'notebook-to-rename'; + const otherNotebookId = 'other-notebook-id'; + const oldName = 'Old Notebook Name'; + const newName = 'New Notebook Name'; + const fileUri = Uri.file('/workspace/test-project.deepnote'); + + // Mock existing project data + 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: otherNotebookId, + name: 'Other Notebook', + blocks: [], + executionMode: 'block' + }, { id: notebookId, name: oldName, @@ -1045,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'; @@ -1232,36 +1197,29 @@ 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'; 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)); @@ -1270,7 +1228,6 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { Promise.resolve(undefined) ); - // Create mock tree item const mockTreeItem: Partial = { type: DeepnoteTreeItemType.Notebook, context: { @@ -1286,12 +1243,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(); }); }); @@ -1422,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'; @@ -1632,14 +1568,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', @@ -1648,71 +1584,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 mockTreeItem: Partial = { - type: DeepnoteTreeItemType.ProjectFile, - context: { - filePath: fileUri.fsPath, - projectId: projectId - }, - data: existingProjectData as unknown as DeepnoteFile + const groupData = { + projectId, + projectName: oldProjectName, + files: [ + { filePath: fileA, project: projectDataA }, + { filePath: fileB, project: projectDataB } + ] }; - // 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); + const mockTreeItem: Partial = { + type: DeepnoteTreeItemType.ProjectGroup, + context: { filePath: fileA, projectId }, + data: groupData as any + }; - // Verify project was renamed - expect(updatedProjectData.project.name).to.equal(newProjectName); + await explorerView.renameProject(mockTreeItem as DeepnoteTreeItem); - // Verify notebooks were not affected - expect(updatedProjectData.project.notebooks).to.have.lengthOf(1); - expect(updatedProjectData.project.notebooks[0].name).to.equal('Notebook 1'); + // 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 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: { @@ -1738,9 +1679,7 @@ suite('DeepnoteExplorerView - Empty State Commands', () => { 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: { @@ -1754,30 +1693,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 @@ -1785,34 +1722,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(); }); @@ -1822,10 +1780,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', @@ -1839,30 +1794,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' } @@ -1875,45 +1823,52 @@ 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 error message was shown verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); }); - test('should export all notebooks when triggered from project', async () => { + test('should concatenate Jupyter outputs across every file in the group', async () => { resetVSCodeMocks(); - const projectData: DeepnoteFile = { + const projectA: 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', - notebooks: [ - { id: 'nb-1', name: 'Notebook 1', blocks: [], executionMode: 'block' }, - { id: 'nb-2', name: 'Notebook 2', blocks: [], executionMode: 'block' } - ] + 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())).thenReturn( - Promise.resolve(Buffer.from(serializeDeepnoteFile(projectData))) - ); + + 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)); @@ -1933,17 +1888,14 @@ 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: projectA }, + { filePath: '/test/project_b.deepnote', project: projectB } + ]); await (explorerView as any).exportProject(treeItem); - // Verify both notebooks were exported + // One notebook per file -> two writes assert.strictEqual(writeCount, 2); verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); }); @@ -2003,17 +1955,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')); @@ -2062,17 +2009,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')); @@ -2109,20 +2051,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(); }); @@ -2149,7 +2085,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)); @@ -2159,24 +2094,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(); }); @@ -2200,7 +2129,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)); @@ -2210,7 +2138,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 ); @@ -2224,18 +2151,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(); }); }); @@ -2617,3 +2557,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/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index b7da2080f3..38a9c0f046 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,17 +12,19 @@ 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'; -import { extractProjectIdFromSnapshotUri, isSnapshotFile } from './snapshots/snapshotFiles'; +import { + extractNotebookIdFromSnapshotUri, + extractProjectIdFromSnapshotUri, + isSnapshotFile +} from './snapshots/snapshotFiles'; import { SnapshotService } from './snapshots/snapshotService'; const debounceTimeInMilliseconds = 500; @@ -38,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; } /** @@ -99,6 +108,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) => { @@ -122,6 +140,13 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } } + public 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); @@ -145,7 +170,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 +212,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic if (liveCells.length !== newCells.length) { return true; } + return liveCells.some( (live, i) => live.kind !== newCells[i].kind || @@ -213,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); @@ -243,12 +269,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) { @@ -258,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); } } @@ -332,22 +361,53 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } } - const wsEdit = new WorkspaceEdit(); - wsEdit.set(notebook.uri, edits); - const applied = await workspace.applyEdit(wsEdit); + // Apply the edit to update in-memory cells immediately (responsive UX). + const applied = await this.applyNotebookEdits(notebook.uri, edits); if (!applied) { logger.warn(`[FileChangeWatcher] Failed to apply edit: ${notebook.uri.path}`); 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); + + // 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; + } + + // Save to clear dirty state. VS Code serializes (same bytes) and sees the + // mtime from our recent write, so no "content is newer" conflict. + // NOTE: onDidSaveNotebookDocument handles the self-write mark for this save. + try { + const saved = await workspace.save(notebook.uri); + if (!saved) { + logger.warn(`[FileChangeWatcher] Save after sync write returned undefined: ${notebook.uri.path}`); + return; + } + } catch (saveError) { + 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}`); @@ -358,18 +418,30 @@ 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; } + 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) { @@ -378,8 +450,7 @@ 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 originalBlocks = metadataNotebookId ? notebookBlocksMap.get(metadataNotebookId) : undefined; // Collect cells that need output updates const cellUpdates: Array<{ @@ -452,9 +523,13 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } } if (metadataEdits.length > 0) { - const wsEdit = new WorkspaceEdit(); - wsEdit.set(notebook.uri, metadataEdits); - await workspace.applyEdit(wsEdit); + 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}`); @@ -482,9 +557,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}`); @@ -503,19 +576,16 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic }) ); } - const metaEdit = new WorkspaceEdit(); - metaEdit.set(notebook.uri, metadataEdits); - await workspace.applyEdit(metaEdit); - - // 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; + 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. + // 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}`); } @@ -523,6 +593,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)) { @@ -561,7 +640,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); @@ -571,17 +652,18 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic key, setTimeout(() => { this.debounceTimers.delete(key); - this.enqueueSnapshotOutputUpdate(projectId); + this.enqueueSnapshotOutputUpdate(projectId, notebookId); }, debounceTimeInMilliseconds) ); } /** - * 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 = uri.toString(); + const key = this.normalizeFileUri(uri); const count = this.selfWriteCounts.get(key) ?? 0; this.selfWriteCounts.set(key, count + 1); diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index c9229820e2..9ae6f366e7 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -1,21 +1,60 @@ +import { DeepnoteFile, serializeDeepnoteFile } from '@deepnote/blocks'; import { assert } from 'chai'; +import { mkdtempSync, rmSync } from 'fs'; +import { tmpdir } from 'os'; 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, resetCalls, when } from 'ts-mockito'; +import { + Disposable, + EventEmitter, + FileSystemWatcher, + NotebookCellData, + NotebookCellKind, + NotebookDocument, + Uri +} from 'vscode'; +import { join } from '../../platform/vscode-path/path'; +import { logger } from '../../platform/logging'; -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'; +const validProject: DeepnoteFile = { + 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")', + metadata: {} + } + ] + } + ] + } +}; + const waitForTimeoutMs = 5000; 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. @@ -35,11 +74,27 @@ 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.joinPath(Uri.file(testFixturesDir), ...pathSegments); + } + let watcher: DeepnoteFileChangeWatcher; let mockDisposables: IDisposableRegistry; + let mockedNotebookManager: IDeepnoteNotebookManager; let mockNotebookManager: IDeepnoteNotebookManager; let onDidChangeFile: EventEmitter; let onDidCreateFile: EventEmitter; + let onDidSaveNotebook: EventEmitter; let readFileCalls: number; let applyEditCount: number; let saveCount: number; @@ -51,7 +106,11 @@ suite('DeepnoteFileChangeWatcher', () => { saveCount = 0; mockDisposables = []; - mockNotebookManager = instance(mock()); + + mockedNotebookManager = mock(); + when(mockedNotebookManager.getOriginalProject(anything(), anything())).thenReturn(validProject); + when(mockedNotebookManager.updateOriginalProject(anything(), anything(), anything())).thenReturn(); + mockNotebookManager = instance(mockedNotebookManager); // Set up FileSystemWatcher mock onDidChangeFile = new EventEmitter(); @@ -63,13 +122,16 @@ 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); }); - 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); @@ -83,6 +145,7 @@ suite('DeepnoteFileChangeWatcher', () => { } onDidChangeFile.dispose(); onDidCreateFile.dispose(); + onDidSaveNotebook.dispose(); }); function createMockNotebook(opts: { @@ -91,20 +154,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 { @@ -126,6 +195,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; @@ -149,8 +219,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, @@ -178,7 +318,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]); @@ -195,7 +335,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([]); @@ -210,7 +350,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]); @@ -228,7 +368,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]); @@ -249,7 +389,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]); @@ -265,7 +405,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, @@ -289,7 +429,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, @@ -314,20 +454,35 @@ project: }); test('should not suppress real changes after auto-save', async function () { - this.timeout(5000); - const uri = Uri.file('/workspace/test.deepnote'); + // 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 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); - await waitFor(() => saveCount >= 1); + await waitFor(() => saveCount >= 1 && applyEditCount >= 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); + + // 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 = ` @@ -349,13 +504,83 @@ 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'); }); + 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'); + const notebook = createMockNotebook({ + uri, + cells: [ + { + metadata: { id: 'block-1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(serializeDeepnoteFile(validProject)); + + // 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(validProject); + 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(validProject); + 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 = Uri.file('/workspace/test.deepnote'); + const uri = testFileUri('test.deepnote'); const notebook = createMockNotebook({ uri, cellCount: 0 }); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); @@ -363,14 +588,14 @@ 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)'); }); 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({ @@ -458,9 +683,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( @@ -480,9 +705,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' }, @@ -509,7 +734,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([]); @@ -526,9 +751,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: [] }] }); @@ -549,7 +774,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); @@ -559,12 +784,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' }, @@ -587,9 +810,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' }, @@ -611,9 +834,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' }, @@ -638,7 +861,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' }, @@ -657,10 +880,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' }, @@ -682,8 +905,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: [ @@ -726,9 +949,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' }, @@ -754,9 +977,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' }, @@ -778,10 +1001,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' }, @@ -813,7 +1036,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: { @@ -853,10 +1076,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: [ { @@ -900,7 +1123,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 = { @@ -908,7 +1131,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' }, @@ -1020,9 +1243,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' }, @@ -1053,9 +1276,45 @@ project: }); test('should not apply updates when cells have no block IDs and no fallback', async () => { - const snapshotUri = Uri.file('/workspace/snapshots/my-project_project-1_latest.snapshot.deepnote'); + 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(), anything())).thenReturn(undefined); + + const nfWatcher = new DeepnoteFileChangeWatcher( + noFallbackDisposables, + instance(nfManager), + instance(nfSnapshotService) + ); + nfWatcher.activate(); + + 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: {}, @@ -1068,16 +1327,19 @@ project: when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); - snapshotOnDidChange.fire(snapshotUri); + noFallbackOnDidChange.fire(snapshotUri); - await waitFor(() => readSnapshotCallCount >= 1); + await waitFor(() => nfReadSnapshotCount >= 1); + await new Promise((resolve) => setTimeout(resolve, postSnapshotReadGraceMs)); - assert.isAtLeast(readSnapshotCallCount, 1, 'readSnapshot should be called'); - assert.strictEqual( - snapshotApplyEditCount, - 0, - 'applyEdit should NOT be called when no block IDs can be resolved' - ); + assert.isAtLeast(nfReadSnapshotCount, 1, 'readSnapshot should be called'); + assert.strictEqual(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 () => { @@ -1109,9 +1371,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' }, @@ -1136,5 +1398,309 @@ project: fbOnDidChange.dispose(); 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.getOriginalProject(anything(), anything())).thenReturn(validProject); + when(mockedManagerEx.updateOriginalProject(anything(), 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(); + } + }); + + 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); + }); + + 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' + ); + }); }); }); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 63e7129200..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,10 +688,10 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, this.notebookControllers.set(notebookKey, controller); // Prepare init notebook execution - const projectId = notebook.metadata?.deepnoteProjectId; - const project = projectId - ? (this.notebookManager.getOriginalProject(projectId) as DeepnoteFile | undefined) - : undefined; + const notebookIdForProject = notebook.metadata?.deepnoteNotebookId; + 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 @@ -688,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}`); } } @@ -793,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); @@ -807,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); @@ -830,7 +857,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, environment, baseFileUri, notebookKey, - projectKey, + projectId, token ); } @@ -847,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); } /** @@ -857,7 +884,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, notebook: NotebookDocument, baseFileUri: Uri, notebookKey: string, - projectKey: string, + projectId: string, token: CancellationToken ): Promise { Cancellation.throwIfCanceled(token); @@ -873,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 ); @@ -899,7 +926,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, environment: DeepnoteEnvironment, baseFileUri: Uri, notebookKey: string, - projectKey: string, + projectId: string, token: CancellationToken ): Promise { try { @@ -915,7 +942,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, environment, baseFileUri, notebookKey, - projectKey, + projectId, progress, progressToken ); @@ -947,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/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..407235e71e --- /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 '!!!' (name that produces no slug)", () => { + 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); + }); + }); +}); 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/deepnote/deepnoteNotebookManager.ts b/src/notebooks/deepnote/deepnoteNotebookManager.ts index 069f3a570c..d515fd8bf2 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.ts @@ -4,126 +4,121 @@ import { IDeepnoteNotebookManager, ProjectIntegration } from '../types'; import type { DeepnoteProject } from '../../platform/deepnote/deepnoteTypes'; /** - * 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. + * + * 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 currentNotebookId = new Map(); - private readonly originalProjects = new Map(); + private readonly originalProjects = new Map>(); private readonly projectsWithInitNotebookRun = new Set(); - private readonly selectedNotebookByProject = new Map(); /** - * 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. + * 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; } /** - * Gets the selected notebook ID for a specific project. + * Checks if the init notebook has already been run for a project. * @param projectId Project identifier - * @returns Selected notebook ID or undefined if not set + * @returns True if init notebook has been run, false otherwise */ - getTheSelectedNotebookForAProject(projectId: string): string | undefined { - return this.selectedNotebookByProject.get(projectId); + hasInitNotebookBeenRun(projectId: string): boolean { + return this.projectsWithInitNotebookRun.has(projectId); } /** - * 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. - * - * @param projectId - The project ID that identifies the Deepnote project - * @param notebookId - The ID of the selected notebook within the project + * Marks the init notebook as having been run for a project. + * @param projectId Project identifier */ - selectNotebookForProject(projectId: string, notebookId: string): void { - this.selectedNotebookByProject.set(projectId, notebookId); + 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 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 - * @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, notebookId: string, project: DeepnoteProject): void { const clonedProject = structuredClone(project); + let notebookMap = this.originalProjects.get(projectId); - this.originalProjects.set(projectId, clonedProject); - this.currentNotebookId.set(projectId, notebookId); + if (!notebookMap) { + notebookMap = new Map(); + this.originalProjects.set(projectId, notebookMap); + } + + notebookMap.set(notebookId, clonedProject); } /** - * Updates the current notebook ID for a project. - * Used when switching notebooks within the same project. + * 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 New current notebook ID + * @param notebookId Notebook identifier within the project + * @param project Updated project data to store */ - updateCurrentNotebookId(projectId: string, notebookId: string): void { - this.currentNotebookId.set(projectId, notebookId); + updateOriginalProject(projectId: string, notebookId: string, project: DeepnoteProject): void { + const clonedProject = structuredClone(project); + 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; - - const currentNotebookId = this.currentNotebookId.get(projectId); - - if (currentNotebookId) { - this.storeOriginalProject(projectId, updatedProject, currentNotebookId); - } else { - 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; } - - /** - * 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); - } } diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts index feb2842848..2ebd2c7c8a 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.unit.test.ts @@ -25,31 +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'); - }); - - 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', () => { test('should return undefined for unknown project', () => { const result = manager.getOriginalProject('unknown-project'); @@ -58,7 +33,7 @@ suite('DeepnoteNotebookManager', () => { }); test('should return original project after storing', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); const result = manager.getOriginalProject('project-123'); @@ -66,127 +41,119 @@ 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'); + suite('storeOriginalProject', () => { + test('should store project data', () => { + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); - const result1 = manager.getTheSelectedNotebookForAProject('project-1'); - const result2 = manager.getTheSelectedNotebookForAProject('project-2'); + const storedProject = manager.getOriginalProject('project-123'); - assert.strictEqual(result1, 'notebook-1'); - assert.strictEqual(result2, 'notebook-2'); + assert.deepStrictEqual(storedProject, mockProject); }); - }); - - 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 project data', () => { + const updatedProject: DeepnoteProject = { + ...mockProject, + project: { + ...mockProject.project, + name: 'Updated Project' + } + }; - test('should overwrite existing selection', () => { - manager.selectNotebookForProject('project-123', 'notebook-456'); - manager.selectNotebookForProject('project-123', 'notebook-789'); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); + manager.storeOriginalProject('project-123', 'notebook-1', updatedProject); - const result = manager.getTheSelectedNotebookForAProject('project-123'); + const storedProject = manager.getOriginalProject('project-123'); - assert.strictEqual(result, 'notebook-789'); + assert.deepStrictEqual(storedProject, updatedProject); }); - test('should handle multiple projects independently', () => { - manager.selectNotebookForProject('project-1', 'notebook-1'); - manager.selectNotebookForProject('project-2', 'notebook-2'); + 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' } + }; - const result1 = manager.getTheSelectedNotebookForAProject('project-1'); - const result2 = manager.getTheSelectedNotebookForAProject('project-2'); + manager.storeOriginalProject('project-123', 'notebook-a', projectA); + manager.storeOriginalProject('project-123', 'notebook-b', projectB); - assert.strictEqual(result1, 'notebook-1'); - assert.strictEqual(result2, 'notebook-2'); + 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('storeOriginalProject', () => { - test('should store both project and current notebook ID', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); - - 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', () => { + suite('updateOriginalProject', () => { + test('should update project data', () => { const updatedProject: DeepnoteProject = { ...mockProject, project: { ...mockProject.project, - name: 'Updated Project' + name: 'Updated Name Only' } }; - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); - manager.storeOriginalProject('project-123', updatedProject, 'notebook-789'); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); + manager.updateOriginalProject('project-123', 'notebook-1', updatedProject); const storedProject = manager.getOriginalProject('project-123'); - const currentNotebookId = manager.getCurrentNotebookId('project-123'); assert.deepStrictEqual(storedProject, updatedProject); - assert.strictEqual(currentNotebookId, 'notebook-789'); }); - }); - suite('updateCurrentNotebookId', () => { - test('should update notebook ID for existing project', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); - manager.updateCurrentNotebookId('project-123', 'notebook-789'); + 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', 'notebook-1', mockProject); + manager.updateOriginalProject('project-123', 'notebook-1', updatedProject); - const result = manager.getCurrentNotebookId('project-123'); + updatedProject.project.name = 'After Mutation'; - assert.strictEqual(result, 'notebook-789'); + const storedProject = manager.getOriginalProject('project-123'); + + assert.strictEqual(storedProject?.project.name, 'Before Mutation'); }); - test('should set notebook ID for new project', () => { - manager.updateCurrentNotebookId('new-project', 'notebook-123'); + test('should overwrite existing project data on successive updates', () => { + const firstUpdate: DeepnoteProject = { + ...mockProject, + project: { ...mockProject.project, name: 'First Update' } + }; + const secondUpdate: DeepnoteProject = { + ...mockProject, + project: { ...mockProject.project, name: 'Second Update' } + }; - const result = manager.getCurrentNotebookId('new-project'); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); + manager.updateOriginalProject('project-123', 'notebook-1', firstUpdate); + manager.updateOriginalProject('project-123', 'notebook-1', secondUpdate); - assert.strictEqual(result, 'notebook-123'); + assert.strictEqual(manager.getOriginalProject('project-123')?.project.name, 'Second Update'); }); - test('should handle multiple projects independently', () => { - manager.updateCurrentNotebookId('project-1', 'notebook-1'); - manager.updateCurrentNotebookId('project-2', 'notebook-2'); + test('should store project when no prior data exists', () => { + const projectOnly: DeepnoteProject = { + ...mockProject, + project: { ...mockProject.project, name: 'No Notebook Id Yet' } + }; - const result1 = manager.getCurrentNotebookId('project-1'); - const result2 = manager.getCurrentNotebookId('project-2'); + manager.updateOriginalProject('project-123', 'notebook-1', projectOnly); - assert.strictEqual(result1, 'notebook-1'); - assert.strictEqual(result2, 'notebook-2'); + 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', 'notebook-1', mockProject); const integrations: ProjectIntegration[] = [ { id: 'int-1', name: 'PostgreSQL', type: 'pgsql' }, @@ -210,7 +177,7 @@ suite('DeepnoteNotebookManager', () => { } }; - manager.storeOriginalProject('project-123', projectWithIntegrations, 'notebook-456'); + manager.storeOriginalProject('project-123', 'notebook-1', projectWithIntegrations); const newIntegrations: ProjectIntegration[] = [ { id: 'new-int-1', name: 'New Integration 1', type: 'pgsql' }, @@ -234,7 +201,7 @@ suite('DeepnoteNotebookManager', () => { } }; - manager.storeOriginalProject('project-123', projectWithIntegrations, 'notebook-456'); + manager.storeOriginalProject('project-123', 'notebook-1', projectWithIntegrations); const result = manager.updateProjectIntegrations('project-123', []); @@ -256,7 +223,7 @@ suite('DeepnoteNotebookManager', () => { }); test('should preserve other project properties and return true', () => { - manager.storeOriginalProject('project-123', mockProject, 'notebook-456'); + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); const integrations: ProjectIntegration[] = [{ id: 'int-1', name: 'PostgreSQL', type: 'pgsql' }]; @@ -271,10 +238,25 @@ suite('DeepnoteNotebookManager', () => { assert.deepStrictEqual(updatedProject?.metadata, mockProject.metadata); }); - 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); + 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', () => { + manager.updateOriginalProject('project-123', 'notebook-1', mockProject); const integrations: ProjectIntegration[] = [ { id: 'int-1', name: 'PostgreSQL', type: 'pgsql' }, @@ -295,41 +277,33 @@ 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'); + suite('hasInitNotebookBeenRun', () => { + test('should return false for unknown project', () => { + assert.strictEqual(manager.hasInitNotebookBeenRun('unknown-project'), false); + }); - manager.storeOriginalProject('project-2', mockProject, 'notebook-2'); - manager.selectNotebookForProject('project-2', 'notebook-2'); + 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.getTheSelectedNotebookForAProject('project-1'), 'notebook-1'); - assert.strictEqual(manager.getTheSelectedNotebookForAProject('project-2'), 'notebook-2'); + assert.strictEqual(manager.hasInitNotebookBeenRun('project-123'), true); }); + }); - test('should handle notebook switching within same project', () => { - 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'); + suite('markInitNotebookAsRun', () => { + test('should mark init notebook as run for a project', () => { + manager.markInitNotebookAsRun('project-123'); - assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-2'); - assert.strictEqual(manager.getTheSelectedNotebookForAProject('project-123'), 'notebook-2'); + assert.strictEqual(manager.hasInitNotebookBeenRun('project-123'), true); }); + }); - 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'); + suite('integration scenarios', () => { + test('should handle complete workflow for multiple projects', () => { + manager.storeOriginalProject('project-1', 'notebook-1', mockProject); + manager.storeOriginalProject('project-2', 'notebook-1', mockProject); - // Both should be maintained independently - assert.strictEqual(manager.getCurrentNotebookId('project-123'), 'notebook-original'); - assert.strictEqual(manager.getTheSelectedNotebookForAProject('project-123'), 'notebook-selected'); + assert.deepStrictEqual(manager.getOriginalProject('project-1'), mockProject); + assert.deepStrictEqual(manager.getOriginalProject('project-2'), mockProject); }); }); }); 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/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 17443e93b9..fe3207a2e1 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -1,14 +1,14 @@ 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 { workspace, type CancellationToken, type NotebookData, type NotebookSerializer } from 'vscode'; +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'; @@ -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 @@ -74,13 +75,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); @@ -90,29 +84,28 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } const projectId = deepnoteFile.project.id; - const notebookId = this.findCurrentNotebookId(projectId); + const selectedNotebook = this.findDefaultNotebook(deepnoteFile); - logger.debug(`DeepnoteSerializer: Project ID: ${projectId}, Selected notebook ID: ${notebookId}`); + logger.debug(`DeepnoteSerializer: Project ID: ${projectId}, Selected notebook ID: ${selectedNotebook?.id}`); - if (deepnoteFile.project.notebooks.length === 0) { + if (!selectedNotebook) { throw new Error('Deepnote project contains no notebooks.'); } - const selectedNotebook = notebookId - ? deepnoteFile.project.notebooks.find((nb) => nb.id === notebookId) - : this.findDefaultNotebook(deepnoteFile); + // Initialize vega-lite for output conversion (lazy-loaded ESM module) + await this.converter.initialize(); - if (!selectedNotebook) { - throw new Error(l10n.t('No notebook selected or found')); + 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]; + 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`); @@ -125,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`); @@ -157,7 +150,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { ); } - this.notebookManager.storeOriginalProject(deepnoteFile.project.id, deepnoteFile, selectedNotebook.id); + this.notebookManager.storeOriginalProject(deepnoteFile.project.id, selectedNotebook.id, deepnoteFile); logger.debug(`DeepnoteSerializer: Stored project ${projectId} in notebook manager`); return { @@ -181,39 +174,6 @@ 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. - * @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 { - // 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; - } - - // 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 - const openNotebook = workspace.notebookDocuments.find( - (doc) => doc.notebookType === 'deepnote' && doc.metadata?.deepnoteProjectId === projectId - ); - - return openNotebook?.metadata?.deepnoteNotebookId; - } - /** * Gets the data converter instance for cell/block conversion. * @returns DeepnoteDataConverter instance @@ -245,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.'); @@ -258,15 +228,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { logger.debug('SerializeNotebook: Got and cloned original project'); - const notebookId = - data.metadata?.deepnoteNotebookId || this.notebookManager.getTheSelectedNotebookForAProject(projectId); - - 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) { @@ -346,8 +307,8 @@ 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. + this.notebookManager.updateOriginalProject(projectId, notebookId, originalProject); logger.debug('SerializeNotebook: Serializing to YAML'); @@ -551,25 +512,23 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } /** - * 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 c3332f974d..c865dcbdf5 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,10 +71,7 @@ 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', async () => { const yamlContent = ` version: '1.0.0' metadata: @@ -107,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 () => { @@ -144,9 +135,19 @@ project: await assert.isRejected( serializer.deserializeNotebook(contentWithoutNotebooks, {} as any), - /no notebooks|notebooks.*must contain at least 1/i + /Failed to parse Deepnote file/ ); }); + + 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); + }); }); suite('serializeNotebook', () => { @@ -179,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', 'notebook-1', mockProject); const mockNotebookData = { cells: [ @@ -205,163 +206,28 @@ project: assert.include(yamlString, 'project-123'); assert.include(yamlString, 'notebook-1'); }); - }); - - suite('findCurrentNotebookId', () => { - teardown(() => { - // Reset only the specific mocks used in this suite - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(undefined); - when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); - }); - - test('should return stored notebook ID when available', () => { - manager.selectNotebookForProject('project-123', 'notebook-456'); - - const result = serializer.findCurrentNotebookId('project-123'); - - assert.strictEqual(result, 'notebook-456'); - }); - - test('should fall back to active notebook document when no stored selection', () => { - // Create a mock notebook document - const mockNotebookDoc = { - then: undefined, // Prevent mock from being treated as a Promise-like thenable - 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; - - // 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'); - }); - - 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 active notebook editor over stored selection', () => { - // Store a selection for the project - manager.selectNotebookForProject('project-123', 'stored-notebook'); - - // Mock the active notebook editor to return a different notebook - 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'); - - // Should return the active editor's notebook, not the stored one - 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: { - deepnoteProjectId: 'different-project', - deepnoteNotebookId: 'active-editor-notebook' - } - }; - - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ - notebook: mockActiveNotebook - } as any); - - 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: { - deepnoteProjectId: 'project-123', - deepnoteNotebookId: 'active-editor-notebook' - } - }; - - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ - notebook: mockActiveNotebook - } as any); - - 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'); + test('should throw error when metadata notebook ID is missing', async () => { + manager.storeOriginalProject('project-123', 'notebook-1', mockProject); - // Mock active editor without notebook ID in metadata - const mockActiveNotebook = { - notebookType: 'deepnote', + const mockNotebookData = { + cells: [ + { + kind: 1, // NotebookCellKind.Markup + value: '# Updated second notebook', + languageId: 'markdown', + metadata: {} + } + ], metadata: { deepnoteProjectId: 'project-123' - // Missing deepnoteNotebookId } }; - when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn({ - notebook: mockActiveNotebook - } as any); - - const result = serializer.findCurrentNotebookId('project-123'); - - // Should fall back to stored selection since active editor has no notebook ID - assert.strictEqual(result, 'stored-notebook'); + await assert.isRejected( + serializer.serializeNotebook(mockNotebookData as any, {} as any), + /Cannot determine which notebook to save/ + ); }); }); @@ -388,19 +254,9 @@ project: }); test('should handle manager state operations', () => { - assert.isFunction(manager.getCurrentNotebookId, 'has getCurrentNotebookId method'); assert.isFunction(manager.getOriginalProject, 'has getOriginalProject method'); - assert.isFunction( - manager.getTheSelectedNotebookForAProject, - 'has getTheSelectedNotebookForAProject method' - ); - assert.isFunction(manager.selectNotebookForProject, 'has selectNotebookForProject method'); assert.isFunction(manager.storeOriginalProject, 'has storeOriginalProject method'); }); - - test('should have findCurrentNotebookId method', () => { - assert.isFunction(serializer.findCurrentNotebookId, 'has findCurrentNotebookId method'); - }); }); suite('data structure handling', () => { @@ -472,7 +328,7 @@ project: } }; - manager.storeOriginalProject('project-circular', projectWithCircularRef, 'notebook-1'); + manager.storeOriginalProject('project-circular', 'notebook-1', projectWithCircularRef); const notebookData = { cells: [ @@ -540,7 +396,7 @@ project: }; // Store the project - manager.storeOriginalProject('project-id-test', projectData, 'notebook-1'); + 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 @@ -626,7 +482,7 @@ project: } }; - manager.storeOriginalProject('project-recover-ids', projectData, 'notebook-1'); + 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 @@ -693,7 +549,7 @@ project: } }; - manager.storeOriginalProject('project-new-content', projectData, 'notebook-1'); + manager.storeOriginalProject('project-new-content', 'notebook-1', projectData); // Cell with different content than any original block const notebookData = { @@ -757,251 +613,6 @@ project: }); }); - suite('default notebook selection', () => { - test('should not select Init notebook when other notebooks are available', async () => { - 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 = await serializer.deserializeNotebook(content, {} as any); - - // Should select the Main notebook, not the Init notebook - assert.strictEqual(result.metadata?.deepnoteNotebookId, 'main-notebook'); - assert.strictEqual(result.metadata?.deepnoteNotebookName, 'Main'); - }); - - test('should select Init notebook when it is the only notebook', async () => { - 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 = await serializer.deserializeNotebook(content, {} as any); - - // Should select the Init notebook since it's the only one - assert.strictEqual(result.metadata?.deepnoteNotebookId, 'init-notebook'); - assert.strictEqual(result.metadata?.deepnoteNotebookName, 'Init'); - }); - - test('should select alphabetically first notebook when no initNotebookId', async () => { - 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 = await serializer.deserializeNotebook(content, {} as any); - - // Should select the alphabetically first notebook - assert.strictEqual(result.metadata?.deepnoteNotebookId, 'alpha-notebook'); - assert.strictEqual(result.metadata?.deepnoteNotebookName, 'Alpha Notebook'); - }); - - test('should sort Init notebook last when multiple notebooks exist', async () => { - 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 = await serializer.deserializeNotebook(content, {} as any); - - // 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'); - }); - }); - suite('detectContentChanges', () => { test('should detect no changes when content is identical', () => { const project: DeepnoteFile = { @@ -1491,7 +1102,7 @@ project: } }; - manager.storeOriginalProject('project-snapshot-hash', projectData, 'notebook-1'); + manager.storeOriginalProject('project-snapshot-hash', 'notebook-1', projectData); const notebookData = { cells: [ @@ -1566,13 +1177,13 @@ project: }; // Serialize twice - manager.storeOriginalProject('project-deterministic', structuredClone(projectData), 'notebook-1'); + 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), 'notebook-1'); + 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 }; @@ -1667,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', '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 }; @@ -1716,7 +1327,7 @@ project: } }; - manager.storeOriginalProject('project-content-change', projectData1, 'notebook-1'); + manager.storeOriginalProject('project-content-change', 'notebook-1', projectData1); const notebookData1 = { cells: [ @@ -1794,7 +1405,7 @@ project: } }; - manager.storeOriginalProject('project-version-change', projectData1, 'notebook-1'); + manager.storeOriginalProject('project-version-change', 'notebook-1', projectData1); const notebookData = { cells: [ @@ -1818,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', 'notebook-1', projectData2); const result2 = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed2 = parseYaml(new TextDecoder().decode(result2)) as DeepnoteFile & { @@ -1860,7 +1471,7 @@ project: } }; - manager.storeOriginalProject('project-integrations-change', projectData1, 'notebook-1'); + manager.storeOriginalProject('project-integrations-change', 'notebook-1', projectData1); const notebookData = { cells: [ @@ -1885,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', 'notebook-1', projectData2); const result2 = await serializer.serializeNotebook(notebookData as any, {} as any); const parsed2 = parseYaml(new TextDecoder().decode(result2)) as DeepnoteFile & { @@ -1927,7 +1538,7 @@ project: } }; - manager.storeOriginalProject('project-env-hash', projectData1, 'notebook-1'); + manager.storeOriginalProject('project-env-hash', 'notebook-1', projectData1); const notebookData = { cells: [ @@ -1952,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', '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/deepnoteTreeDataProvider.ts b/src/notebooks/deepnote/deepnoteTreeDataProvider.ts index be9e471e27..5587ccc759 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,108 @@ export class DeepnoteTreeDataProvider implements TreeDataProvider 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; + if (project) { + results.push({ filePath: file.path, project }); } - - deepnoteFiles.push(treeItem); } catch (error) { this.logger?.error(`Failed to load Deepnote project from ${file.path}`, error); } } } - // Sort projects alphabetically by name (case-insensitive) - deepnoteFiles.sort(compareTreeItemsByLabel); + return results; + } + + /** + * Get top-level project groups (grouped by project ID) + */ + private async getProjectGroups(): Promise { + 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; + + const groupData: ProjectGroupData = { + projectId, + projectName, + files: files.map((f) => ({ filePath: f.filePath, project: f.project })) + }; + + const context: DeepnoteTreeItemContext = { + filePath: files[0].filePath, + projectId + }; + + // Expand single-file groups by default so the lone notebook stays visible + const collapsibleState = + files.length === 1 ? TreeItemCollapsibleState.Expanded : TreeItemCollapsibleState.Collapsed; + + const groupItem = new DeepnoteTreeItem( + DeepnoteTreeItemType.ProjectGroup, + context, + groupData, + collapsibleState + ); + groups.push(groupItem); + } + + groups.sort(compareTreeItemsByLabel); - return deepnoteFiles; + return groups; } - private async getNotebooksForProject(projectItem: DeepnoteTreeItem): Promise { + /** + * 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); + } + + fileItems.sort(compareTreeItemsByLabel); + + return fileItems; + } + + /** + * 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 +317,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..e391bcfd26 100644 --- a/src/notebooks/deepnote/deepnoteTreeDataProvider.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteTreeDataProvider.unit.test.ts @@ -1,9 +1,17 @@ +import { serializeDeepnoteFile } from '@deepnote/blocks'; import { assert } from 'chai'; -import { l10n } from 'vscode'; +import { instance, mock, when, anything } from 'ts-mockito'; +import { l10n, TreeItemCollapsibleState, Uri, workspace } from 'vscode'; import { DeepnoteTreeDataProvider, compareTreeItemsByLabel } from './deepnoteTreeDataProvider'; -import { DeepnoteTreeItem, DeepnoteTreeItemType } from './deepnoteTreeItem'; +import { + DeepnoteTreeItem, + DeepnoteTreeItemType, + NOTEBOOK_FILE_CONTEXT_VALUE, + type ProjectGroupData +} from './deepnoteTreeItem'; import type { DeepnoteProject } from '../../platform/deepnote/deepnoteTypes'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; suite('DeepnoteTreeDataProvider', () => { 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, 'Original Name'); - assert.strictEqual(mockTreeItem.description, '1 notebook'); + // 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,24 +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 - mockTreeItem.label = 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' : ''}`; - } - - // 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' - ); + // 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); + + // 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', @@ -522,16 +521,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 +583,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 +640,160 @@ 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'); + }); + }); +}); + +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 0763f0e1bb..9211098793 100644 --- a/src/notebooks/deepnote/deepnoteTreeItem.ts +++ b/src/notebooks/deepnote/deepnoteTreeItem.ts @@ -1,10 +1,13 @@ -import { TreeItem, TreeItemCollapsibleState, Uri, ThemeIcon } from 'vscode'; +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 */ export enum DeepnoteTreeItemType { + ProjectGroup = 'projectGroup', ProjectFile = 'projectFile', Notebook = 'notebook', Loading = 'loading' @@ -20,25 +23,45 @@ 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 }>; +} + +/** + * 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 */ 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 +74,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 +82,15 @@ 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 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'; @@ -68,8 +99,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 +110,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 +142,32 @@ 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; + const singleNonInitNotebook = getSingleNonInitNotebook(project); - this.label = project.project.name || 'Untitled 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 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; @@ -124,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 bbadeceee4..40676a2b7e 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,15 +93,15 @@ suite('DeepnoteTreeItem', () => { const item = new DeepnoteTreeItem( DeepnoteTreeItemType.ProjectFile, context, - mockProject, + multiNotebookProject, TreeItemCollapsibleState.Collapsed ); 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', () => { @@ -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,30 +281,31 @@ suite('DeepnoteTreeItem', () => { const item = new DeepnoteTreeItem( DeepnoteTreeItemType.ProjectFile, context, - mockProject, + multiNotebookProject, 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', () => { + test('should handle project with three notebooks', () => { const projectWithMultipleNotebooks: DeepnoteProject = { - ...mockProject, + ...multiNotebookProject, project: { - ...mockProject.project, + ...multiNotebookProject.project, notebooks: [ { id: 'notebook-1', @@ -175,14 +344,16 @@ suite('DeepnoteTreeItem', () => { TreeItemCollapsibleState.Collapsed ); - assert.strictEqual(item.description, '3 notebooks'); + 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: [] } }; @@ -199,14 +370,16 @@ suite('DeepnoteTreeItem', () => { TreeItemCollapsibleState.Collapsed ); - assert.strictEqual(item.description, '0 notebooks'); + 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 } }; @@ -223,7 +396,7 @@ suite('DeepnoteTreeItem', () => { TreeItemCollapsibleState.Collapsed ); - assert.strictEqual(item.label, 'Untitled Project'); + assert.strictEqual(item.label, 'project.deepnote'); }); }); @@ -259,12 +432,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', () => { @@ -398,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' }, @@ -413,12 +591,13 @@ suite('DeepnoteTreeItem', () => { ); assert.strictEqual(projectItem.contextValue, 'projectFile'); + assert.strictEqual(notebookFileItem.contextValue, NOTEBOOK_FILE_CONTEXT_VALUE); assert.strictEqual(notebookItem.contextValue, 'notebook'); }); }); 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' @@ -427,11 +606,14 @@ suite('DeepnoteTreeItem', () => { const item = new DeepnoteTreeItem( DeepnoteTreeItemType.ProjectFile, context, - mockProject, + multiNotebookProject, 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 +639,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' @@ -466,12 +648,12 @@ suite('DeepnoteTreeItem', () => { const item = new DeepnoteTreeItem( DeepnoteTreeItemType.ProjectFile, context, - mockProject, + multiNotebookProject, TreeItemCollapsibleState.Collapsed ); 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', () => { @@ -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 behavior, 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 ) ); @@ -746,7 +1066,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/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/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..f0302f61da 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, @@ -202,9 +254,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 +274,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 +450,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 +464,13 @@ 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`; + + 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}`); @@ -413,24 +478,43 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync 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 = `**/snapshots/*_${projectId}_*.snapshot.deepnote`; + const timestampedGlob = notebookId + ? `**/snapshots/*_${projectId}_${notebookId}_*.snapshot.deepnote` + : `**/snapshots/*_${projectId}_*.snapshot.deepnote`; let allTimestampedFiles: Uri[] = []; for (const folder of workspaceFolders) { @@ -453,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)}`); + } - logger.debug(`[Snapshot] Using timestamped snapshot: ${Utils.basename(newestFile)}`); + 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. + } + } + + // 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; } @@ -494,11 +606,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); } @@ -618,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; @@ -629,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(); @@ -712,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`); @@ -739,26 +972,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; } @@ -778,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; @@ -800,7 +1046,8 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync projectId, originalProject.project.name, snapshotProject, - notebookUri + notebookUri, + notebookId ); if (snapshotUri) { @@ -814,7 +1061,8 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync projectId, originalProject.project.name, snapshotProject, - notebookUri + notebookUri, + notebookId ); if (snapshotUri) { @@ -866,12 +1114,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 +1260,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/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', () => { diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index cbd8b860fe..f671334941 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'; @@ -66,7 +67,7 @@ import { IDeepnoteKernelAutoSelector, IDeepnoteServerProvider, IDeepnoteEnvironmentManager, - IDeepnoteNotebookEnvironmentMapper, + IDeepnoteProjectEnvironmentMapper, IDeepnoteLspClientManager } from '../kernels/deepnote/types'; import { DeepnoteAgentSkillsManager } from '../kernels/deepnote/deepnoteAgentSkillsManager.node'; @@ -82,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'; @@ -169,6 +170,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteActivationService ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + DeepnoteNotebookInfoStatusBar + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNotebookCommandListener @@ -250,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/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/notebooks/types.ts b/src/notebooks/types.ts index a12dacf52a..a8823edb95 100644 --- a/src/notebooks/types.ts +++ b/src/notebooks/types.ts @@ -37,12 +37,11 @@ export interface ProjectIntegration { export const IDeepnoteNotebookManager = Symbol('IDeepnoteNotebookManager'); export interface IDeepnoteNotebookManager { - getCurrentNotebookId(projectId: string): string | undefined; - getOriginalProject(projectId: string): DeepnoteProject | undefined; - getTheSelectedNotebookForAProject(projectId: string): string | undefined; - selectNotebookForProject(projectId: string, notebookId: string): void; - storeOriginalProject(projectId: string, project: DeepnoteProject, notebookId: string): void; - updateCurrentNotebookId(projectId: string, notebookId: string): void; + getOriginalProject(projectId: string, notebookId?: string): DeepnoteProject | undefined; + hasInitNotebookBeenRun(projectId: string): boolean; + markInitNotebookAsRun(projectId: string): 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. @@ -53,7 +52,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; } 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'; 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 new file mode 100644 index 0000000000..b39d3c56c8 --- /dev/null +++ b/src/platform/deepnote/deepnoteProjectIdResolver.ts @@ -0,0 +1,36 @@ +import { NotebookDocument, Uri } from 'vscode'; + +import { logger } from '../logging'; +import { readDeepnoteProjectFile } from './deepnoteProjectFileReader'; + +/** + * 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.toString()}`, 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}`; } 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 b6b0f5371f..d413837dc7 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 @@ -300,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; }