diff --git a/README.md b/README.md index 04935a5956..7e97d7a529 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,6 @@ A powerful [Visual Studio Code](https://marketplace.visualstudio.com/items?itemN ![Deepnote Projects](./assets/deepnote-projects.png) - Run Deepnote locally inside your IDE and unlock the next generation of data workflows: - **Rich block types:** Combine Python, Markdown, data visualizations, tables, and more — all in one place @@ -66,6 +65,7 @@ Open the Command Palette (`Cmd+Shift+P` or `Ctrl+Shift+P`) and type `Deepnote` t | `Deepnote: Import Notebook` | Import an existing notebook into your project | | `Notebook: Select Notebook Kernel` | Select or switch kernels within your notebook | | `Notebook: Change Cell Language` | Change the language of the cell currently in focus | +| `Deepnote: Enable Snapshots` | Enable snapshot mode for the current workspace | ### Database integrations @@ -89,6 +89,24 @@ SELECT * FROM users WHERE created_at > '2024-01-01' Results are displayed as interactive tables that you can explore and export. +### Snapshot mode + +Snapshot mode gives you a historical, portable record of all notebook executions without polluting your main project files. This makes it easier to work with Git since outputs are stored separately from your source code. + +**How it works:** + +- Execution outputs are saved to a `snapshots/` folder alongside your project +- Your main `.deepnote` file stays clean (no outputs), making diffs readable +- Each "Run All" execution creates a timestamped snapshot for historical tracking +- Running individual cells updates only the latest snapshot + +**To enable:** + +1. Open Command Palette (`Cmd+Shift+P` or `Ctrl+Shift+P`) +2. Run `Deepnote: Enable Snapshots` + +Once enabled, snapshots are automatically created when you execute notebooks. You can add the `snapshots/` folder to `.gitignore` to keep outputs local, or commit them to share execution history with your team. + ## Need help? - Join our [Community](https://github.com/deepnote/deepnote/discussions)! diff --git a/cspell.json b/cspell.json index 944d2c7eb1..7181203fad 100644 --- a/cspell.json +++ b/cspell.json @@ -78,9 +78,14 @@ "scikit", "scipy", "sklearn", + "slugification", + "slugified", + "slugifies", + "slugify", "sqlalchemy", "taskkill", "testdb", + "testproject", "toolsai", "trino", "Trino", diff --git a/package.json b/package.json index 631461563d..7d7791a9ec 100644 --- a/package.json +++ b/package.json @@ -93,6 +93,16 @@ "category": "Deepnote", "icon": "$(reveal)" }, + { + "command": "deepnote.enableSnapshots", + "title": "%deepnote.commands.enableSnapshots.title%", + "category": "Deepnote" + }, + { + "command": "deepnote.disableSnapshots", + "title": "%deepnote.commands.disableSnapshots.title%", + "category": "Deepnote" + }, { "command": "deepnote.environments.create", "title": "%deepnote.commands.environments.create.title%", @@ -1640,6 +1650,12 @@ "description": "Disable SSL certificate verification (for development only)", "scope": "application" }, + "deepnote.snapshots.enabled": { + "type": "boolean", + "default": false, + "description": "When enabled, outputs are saved to separate snapshot files in a 'snapshots' folder instead of the main .deepnote file.", + "scope": "resource" + }, "deepnote.experiments.enabled": { "type": "boolean", "default": true, diff --git a/package.nls.json b/package.nls.json index cbc6704968..01f16b7e7e 100644 --- a/package.nls.json +++ b/package.nls.json @@ -250,6 +250,8 @@ "deepnote.commands.openNotebook.title": "Open Notebook", "deepnote.commands.openFile.title": "Open File", "deepnote.commands.revealInExplorer.title": "Reveal in Explorer", + "deepnote.commands.enableSnapshots.title": "Enable Snapshots", + "deepnote.commands.disableSnapshots.title": "Disable Snapshots", "deepnote.commands.manageIntegrations.title": "Manage Integrations", "deepnote.commands.newProject.title": "New Project", "deepnote.commands.importNotebook.title": "Import Notebook", diff --git a/src/kernels/execution/cellExecutionQueue.ts b/src/kernels/execution/cellExecutionQueue.ts index 1973fb50f2..66d78822e4 100644 --- a/src/kernels/execution/cellExecutionQueue.ts +++ b/src/kernels/execution/cellExecutionQueue.ts @@ -14,7 +14,8 @@ import { CodeExecution } from './codeExecution'; import { once } from '../../platform/common/utils/events'; import { getCellMetadata } from '../../platform/common/utils'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; -import { ISnapshotMetadataService } from '../../platform/notebooks/deepnote/types'; +// eslint-disable-next-line import/no-restricted-paths +import { ISnapshotMetadataService } from '../../notebooks/deepnote/snapshots/snapshotService'; /** * A queue responsible for execution of cells. @@ -324,5 +325,10 @@ export class CellExecutionQueue implements Disposable { break; } } + + // Notify listeners that execution queue is complete + if (this.notebook) { + notebookCellExecutions.notifyQueueComplete(this.notebook.uri.toString()); + } } } diff --git a/src/kernels/kernelExecution.ts b/src/kernels/kernelExecution.ts index 3fac55d16b..9158d7a2cf 100644 --- a/src/kernels/kernelExecution.ts +++ b/src/kernels/kernelExecution.ts @@ -45,7 +45,7 @@ import { import { CodeExecution } from './execution/codeExecution'; import type { ICodeExecution } from './execution/types'; import { NotebookCellExecutionState, notebookCellExecutions } from '../platform/notebooks/cellExecutionStateService'; -import { ISnapshotMetadataService } from '../platform/notebooks/deepnote/types'; +import { ISnapshotMetadataService } from '../notebooks/deepnote/snapshots/snapshotService'; /** * Everything in this classes gets disposed via the `onWillCancel` hook. diff --git a/src/kernels/kernelProvider.node.ts b/src/kernels/kernelProvider.node.ts index c91575f1f0..caf3c30028 100644 --- a/src/kernels/kernelProvider.node.ts +++ b/src/kernels/kernelProvider.node.ts @@ -32,7 +32,8 @@ import { IReplNotebookTrackerService } from '../platform/notebooks/replNotebookT import { logger } from '../platform/logging'; import { getDisplayPath } from '../platform/common/platform/fs-paths.node'; import { IRawNotebookSupportedService } from './raw/types'; -import { ISnapshotMetadataService } from '../platform/notebooks/deepnote/types'; +// eslint-disable-next-line import/no-restricted-paths +import { ISnapshotMetadataService } from '../notebooks/deepnote/snapshots/snapshotService'; /** * Node version of a kernel provider. Needed in order to create the node version of a kernel. diff --git a/src/notebooks/deepnote/deepnoteActivationService.ts b/src/notebooks/deepnote/deepnoteActivationService.ts index 49a600c5ca..4a33418a26 100644 --- a/src/notebooks/deepnote/deepnoteActivationService.ts +++ b/src/notebooks/deepnote/deepnoteActivationService.ts @@ -9,7 +9,7 @@ import { DeepnoteNotebookSerializer } from './deepnoteSerializer'; import { DeepnoteExplorerView } from './deepnoteExplorerView'; import { IIntegrationManager } from './integrations/types'; import { DeepnoteInputBlockEditProtection } from './deepnoteInputBlockEditProtection'; -import { ISnapshotMetadataService, ISnapshotMetadataServiceFull } from './snapshotMetadataService'; +import { SnapshotService } from './snapshots/snapshotService'; /** * Service responsible for activating and configuring Deepnote notebook support in VS Code. @@ -30,7 +30,7 @@ export class DeepnoteActivationService implements IExtensionSyncActivationServic @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, @inject(IIntegrationManager) integrationManager: IIntegrationManager, @inject(ILogger) private readonly logger: ILogger, - @inject(ISnapshotMetadataService) @optional() private readonly snapshotService?: ISnapshotMetadataServiceFull + @inject(SnapshotService) @optional() private readonly snapshotService?: SnapshotService ) { this.integrationManager = integrationManager; } diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts index 75c06ec2e3..3439ee83a8 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts @@ -1,6 +1,7 @@ import { injectable, inject } from 'inversify'; import { commands, + ConfigurationTarget, window, NotebookCellData, NotebookCellKind, @@ -16,7 +17,7 @@ import z from 'zod'; import { logger } from '../../platform/logging'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; -import { IDisposableRegistry } from '../../platform/common/types'; +import { IConfigurationService, IDisposableRegistry } from '../../platform/common/types'; import { Commands } from '../../platform/common/constants'; import { notebookUpdaterUtils } from '../../kernels/execution/notebookUpdater'; import { WrappedError } from '../../platform/errors/types'; @@ -149,7 +150,10 @@ export function getNextDeepnoteVariableName(cells: NotebookCell[], prefix: 'df' */ @injectable() export class DeepnoteNotebookCommandListener implements IExtensionSyncActivationService { - constructor(@inject(IDisposableRegistry) private readonly disposableRegistry: IDisposableRegistry) {} + constructor( + @inject(IConfigurationService) private readonly configurationService: IConfigurationService, + @inject(IDisposableRegistry) private readonly disposableRegistry: IDisposableRegistry + ) {} /** * Activates the service by registering Deepnote-specific commands. @@ -217,6 +221,10 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation this.addTextBlockCommandHandler({ textBlockType: 'text-cell-p' }) ) ); + this.disposableRegistry.push(commands.registerCommand(Commands.EnableSnapshots, () => this.enableSnapshots())); + this.disposableRegistry.push( + commands.registerCommand(Commands.DisableSnapshots, () => this.disableSnapshots()) + ); } public async addSqlBlock(): Promise { @@ -537,4 +545,34 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation // Enter edit mode on the new cell await commands.executeCommand('notebook.cell.edit'); } + + private async disableSnapshots(): Promise { + try { + await this.configurationService.updateSetting( + 'snapshots.enabled', + false, + undefined, + ConfigurationTarget.Workspace + ); + void window.showInformationMessage(l10n.t('Snapshots disabled for this workspace.')); + } catch (error) { + logger.error('Failed to disable snapshots', error); + void window.showErrorMessage(l10n.t('Failed to disable snapshots.')); + } + } + + private async enableSnapshots(): Promise { + try { + await this.configurationService.updateSetting( + 'snapshots.enabled', + true, + undefined, + ConfigurationTarget.Workspace + ); + void window.showInformationMessage(l10n.t('Snapshots enabled for this workspace.')); + } catch (error) { + logger.error('Failed to enable snapshots', error); + void window.showErrorMessage(l10n.t('Failed to enable snapshots.')); + } + } } diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts index 8124e48307..4cc7dc1d0e 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts @@ -18,7 +18,7 @@ import { InputBlockType } from './deepnoteNotebookCommandListener'; import { formatInputBlockCellContent, getInputBlockLanguage } from './inputBlockContentFormatter'; -import { IDisposable } from '../../platform/common/types'; +import { IConfigurationService, IDisposable } from '../../platform/common/types'; import * as notebookUpdater from '../../kernels/execution/notebookUpdater'; import { createMockedNotebookDocument } from '../../test/datascience/editor-integration/helpers'; import { WrappedError } from '../../platform/errors/types'; @@ -29,11 +29,21 @@ suite('DeepnoteNotebookCommandListener', () => { let commandListener: DeepnoteNotebookCommandListener; let disposables: IDisposable[]; let sandbox: sinon.SinonSandbox; + let mockConfigService: IConfigurationService; + + function createMockConfigService(): IConfigurationService { + return { + getSettings: sinon.stub().returns({}), + updateSetting: sinon.stub().resolves(), + updateSectionSetting: sinon.stub().resolves() + } as unknown as IConfigurationService; + } setup(() => { sandbox = sinon.createSandbox(); disposables = []; - commandListener = new DeepnoteNotebookCommandListener(disposables); + mockConfigService = createMockConfigService(); + commandListener = new DeepnoteNotebookCommandListener(mockConfigService, disposables); }); teardown(() => { @@ -78,7 +88,7 @@ suite('DeepnoteNotebookCommandListener', () => { // Create new instance and activate again const disposables2: IDisposable[] = []; - const commandListener2 = new DeepnoteNotebookCommandListener(disposables2); + const commandListener2 = new DeepnoteNotebookCommandListener(createMockConfigService(), disposables2); commandListener2.activate(); // Both should register the same number of commands diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.ts b/src/notebooks/deepnote/deepnoteNotebookManager.ts index b6875d5048..8e43380466 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.ts @@ -11,8 +11,8 @@ import type { DeepnoteProject } from '../../platform/deepnote/deepnoteTypes'; export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { private readonly currentNotebookId = new Map(); private readonly originalProjects = new Map(); - private readonly selectedNotebookByProject = new Map(); private readonly projectsWithInitNotebookRun = new Set(); + private readonly selectedNotebookByProject = new Map(); /** * Gets the currently selected notebook ID for a project. diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index ea2440d11c..de5989d29b 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -7,7 +7,7 @@ import { logger } from '../../platform/logging'; import { IDeepnoteNotebookManager } from '../types'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; -import { ISnapshotMetadataService, ISnapshotMetadataServiceFull } from './snapshotMetadataService'; +import { SnapshotService } from './snapshots/snapshotService'; import { computeHash } from '../../platform/common/crypto'; export type { DeepnoteBlock, DeepnoteFile } from '@deepnote/blocks'; @@ -56,7 +56,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { constructor( @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, - @inject(ISnapshotMetadataService) @optional() private readonly snapshotService?: ISnapshotMetadataServiceFull + @inject(SnapshotService) @optional() private readonly snapshotService?: SnapshotService ) {} /** @@ -114,10 +114,33 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { throw new Error(l10n.t('No notebook selected or found')); } - const cells = this.converter.convertBlocksToCells(selectedNotebook.blocks ?? []); + let cells = this.converter.convertBlocksToCells(selectedNotebook.blocks ?? []); logger.debug(`DeepnoteSerializer: Converted ${cells.length} cells from notebook blocks`); + // Merge outputs from snapshot if snapshots are enabled + if (this.snapshotService?.isSnapshotsEnabled()) { + try { + const snapshotOutputs = await this.snapshotService.readSnapshot(projectId); + + if (snapshotOutputs) { + logger.debug(`DeepnoteSerializer: Merging ${snapshotOutputs.size} outputs from snapshot`); + const blocksWithOutputs = this.snapshotService.mergeOutputsIntoBlocks( + selectedNotebook.blocks ?? [], + snapshotOutputs + ); + + cells = this.converter.convertBlocksToCells(blocksWithOutputs); + } + } catch (error) { + logger.warn( + `DeepnoteSerializer: Failed to merge snapshot outputs for project ${projectId}, using baseline cells`, + error + ); + // Fall back to baseline cells (already set above) + } + } + this.notebookManager.storeOriginalProject(deepnoteFile.project.id, deepnoteFile, selectedNotebook.id); logger.debug(`DeepnoteSerializer: Stored project ${projectId} in notebook manager`); @@ -199,7 +222,17 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { // Add snapshot metadata to blocks (contentHash and execution timing) await this.addSnapshotMetadataToBlocks(blocks, data); - notebook.blocks = cloneWithoutCircularRefs(blocks); + // Handle snapshot mode: strip outputs from main file (snapshots are created on execution, not save) + if (this.snapshotService?.isSnapshotsEnabled()) { + // Strip outputs from main file blocks - snapshots are created during cell execution + // Also clone to remove circular references that may cause yaml.dump to fail + const strippedBlocks = this.snapshotService.stripOutputsFromBlocks(blocks); + notebook.blocks = cloneWithoutCircularRefs(strippedBlocks); + logger.debug('SerializeNotebook: Stripped outputs from main file (snapshot mode)'); + } else { + // Default behavior: outputs in main file + notebook.blocks = cloneWithoutCircularRefs(blocks); + } logger.debug('SerializeNotebook: Cloned blocks, updating modifiedAt'); diff --git a/src/notebooks/deepnote/deepnoteTreeDataProvider.ts b/src/notebooks/deepnote/deepnoteTreeDataProvider.ts index 1dec6ad59e..026def0a24 100644 --- a/src/notebooks/deepnote/deepnoteTreeDataProvider.ts +++ b/src/notebooks/deepnote/deepnoteTreeDataProvider.ts @@ -18,6 +18,16 @@ import type { DeepnoteProject, DeepnoteNotebook } from '../../platform/deepnote/ import { readDeepnoteProjectFile } from './deepnoteProjectUtils'; import { ILogger } from '../../platform/logging/types'; +/** File suffix for snapshot files, used to filter them from the explorer */ +const SNAPSHOT_FILE_SUFFIX = '.snapshot.deepnote'; + +/** + * Checks if a URI represents a snapshot file + */ +function isSnapshotFile(uri: Uri): boolean { + return uri.path.endsWith(SNAPSHOT_FILE_SUFFIX); +} + /** * Comparator function for sorting tree items alphabetically by label (case-insensitive) */ @@ -203,8 +213,9 @@ export class DeepnoteTreeDataProvider implements TreeDataProvider !file.path.endsWith(SNAPSHOT_FILE_SUFFIX)); - for (const file of files) { + for (const file of projectFiles) { try { const project = await this.loadDeepnoteProject(file); if (!project) { @@ -316,16 +327,25 @@ export class DeepnoteTreeDataProvider implements TreeDataProvider { + if (isSnapshotFile(uri)) { + return; + } // Use granular refresh for file changes void this.refreshProject(uri.path); }); - this.fileWatcher.onDidCreate(() => { + this.fileWatcher.onDidCreate((uri) => { + if (isSnapshotFile(uri)) { + return; + } // New file created, do full refresh this._onDidChangeTreeData.fire(); }); this.fileWatcher.onDidDelete((uri) => { + if (isSnapshotFile(uri)) { + return; + } // File deleted, clear both caches and do full refresh this.cachedProjects.delete(uri.path); this.treeItemCache.delete(`project:${uri.path}`); diff --git a/src/notebooks/deepnote/snapshotMetadataService.ts b/src/notebooks/deepnote/snapshotMetadataService.ts deleted file mode 100644 index f1803bce5e..0000000000 --- a/src/notebooks/deepnote/snapshotMetadataService.ts +++ /dev/null @@ -1,428 +0,0 @@ -import { inject, injectable } from 'inversify'; -import { NotebookCell, workspace } from 'vscode'; - -import type { Environment, Execution, ExecutionError } from '@deepnote/blocks'; - -import { IEnvironmentCapture } from './environmentCapture.node'; -import { IDisposableRegistry } from '../../platform/common/types'; -import { logger } from '../../platform/logging'; -import { IExtensionSyncActivationService } from '../../platform/activation/types'; -import { notebookCellExecutions, NotebookCellExecutionState } from '../../platform/notebooks/cellExecutionStateService'; -import { ISnapshotMetadataService as IPlatformSnapshotMetadataService } from '../../platform/notebooks/deepnote/types'; - -class TimeoutError extends Error { - constructor(message: string) { - super(message); - this.name = 'TimeoutError'; - } -} - -/** - * Block-level execution metadata. - */ -export interface BlockExecutionMetadata { - /** SHA-256 hash of block source code (prefixed with "sha256:") */ - contentHash: string; - - /** ISO 8601 timestamp when block execution started */ - executionStartedAt?: string; - - /** ISO 8601 timestamp when block execution completed */ - executionFinishedAt?: string; -} - -/** - * Internal state tracking for a notebook execution session. - * Used by SnapshotMetadataService to aggregate execution data. - */ -interface NotebookExecutionState { - /** Number of blocks executed so far */ - blocksExecuted: number; - - /** Number of blocks that failed */ - blocksFailed: number; - - /** Number of blocks that succeeded */ - blocksSucceeded: number; - - /** Promise that resolves when environment capture completes */ - capturePromise?: Promise; - - /** Per-cell execution metadata, keyed by cell ID */ - cellMetadata: Map; - - /** Cached environment metadata */ - environment?: Environment; - - /** Whether environment has been captured for this session */ - environmentCaptured: boolean; - - /** Top-level error if any */ - error?: { name?: string; message?: string; traceback?: string[] }; - - /** ISO 8601 timestamp when last cell finished executing */ - finishedAt?: string; - - /** ISO 8601 timestamp when first cell started executing */ - startedAt: string; - - /** Total duration in milliseconds */ - totalDurationMs: number; -} - -/** - * Service interface for tracking and retrieving snapshot metadata. - * Extends the platform interface with additional methods for the notebooks layer. - */ -export { ISnapshotMetadataService } from '../../platform/notebooks/deepnote/types'; - -export interface ISnapshotMetadataServiceFull extends IPlatformSnapshotMetadataService { - /** - * Get block-level execution metadata for a specific cell. - */ - getBlockExecutionMetadata(notebookUri: string, cellId: string): BlockExecutionMetadata | undefined; - - /** - * Get environment metadata for a notebook. - * If capture is in progress, waits for it to complete before returning. - */ - getEnvironmentMetadata(notebookUri: string): Promise; - - /** - * Get execution metadata for a notebook. - */ - getExecutionMetadata(notebookUri: string): Execution | undefined; - - /** - * Record the end of a cell execution. - */ - recordCellExecutionEnd( - notebookUri: string, - cellId: string, - endTime: number, - success: boolean, - error?: ExecutionError - ): void; - - /** - * Record the start of a cell execution. - */ - recordCellExecutionStart(notebookUri: string, cellId: string, startTime: number): void; -} - -/** - * Service that tracks and aggregates execution metadata for notebooks. - * This service captures: - * - Per-cell execution timing (startedAt, finishedAt) - * - Notebook-level execution summary (blocksExecuted, blocksSucceeded, etc.) - * - Environment metadata (captured once per execution session) - */ -@injectable() -export class SnapshotMetadataService implements ISnapshotMetadataServiceFull, IExtensionSyncActivationService { - private readonly executionStates = new Map(); - - constructor( - @inject(IEnvironmentCapture) private readonly environmentCapture: IEnvironmentCapture, - @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry - ) {} - - activate(): void { - logger.info('[Snapshot] SnapshotMetadataService activated'); - - workspace.onDidCloseNotebookDocument( - (notebook) => { - this.clearExecutionState(notebook.uri.toString()); - }, - this, - this.disposables - ); - - notebookCellExecutions.onDidChangeNotebookCellExecutionState( - (e) => { - logger.debug(`[Snapshot] Cell execution state changed: ${e.state} for cell ${e.cell.metadata?.id}`); - this.handleCellExecutionStateChange(e.cell, e.state); - }, - this, - this.disposables - ); - } - - async captureEnvironmentBeforeExecution(notebookUri: string): Promise { - logger.info(`[Snapshot] captureEnvironmentBeforeExecution called for ${notebookUri}`); - - const state = this.getOrCreateExecutionState(notebookUri, Date.now()); - - // If capture is already in progress, wait for it - if (state.capturePromise) { - logger.info(`[Snapshot] Capture already in progress, waiting...`); - await state.capturePromise; - - return; - } - - // Start capture and store the promise so other callers can wait for it - state.capturePromise = this.captureEnvironmentForNotebook(notebookUri); - } - - clearExecutionState(notebookUri: string): void { - this.executionStates.delete(notebookUri); - - logger.trace(`[Snapshot] Cleared execution state for ${notebookUri}`); - } - - getBlockExecutionMetadata(notebookUri: string, cellId: string): BlockExecutionMetadata | undefined { - const state = this.executionStates.get(notebookUri); - - if (!state) { - return; - } - - return state.cellMetadata.get(cellId); - } - - async getEnvironmentMetadata(notebookUri: string): Promise { - const state = this.executionStates.get(notebookUri); - - logger.info(`[Snapshot] getEnvironmentMetadata for ${notebookUri}`); - logger.info(Boolean(state) ? '[Snapshot] State exists.' : '[Snapshot] No state found.'); - - if (!state) { - logger.info(`[Snapshot] Available URIs: ${Array.from(this.executionStates.keys()).join(', ')}`); - - return; - } - - // If capture is in progress, wait for it to complete - if (state.capturePromise && !state.environmentCaptured) { - logger.info(`[Snapshot] Waiting for capture to complete before returning metadata.`); - - const timeoutPromise = new Promise((_, reject) => { - setTimeout(() => reject(new TimeoutError('Timeout waiting for environment capture.')), 10_000); - }); - - await Promise.race([state.capturePromise, timeoutPromise]).catch((error) => { - if (error instanceof TimeoutError) { - logger.warn('[Snapshot] Timed out waiting for environment capture'); - } else { - throw error; - } - }); - } - - logger.info(`[Snapshot] state.environment exists: ${!!state.environment}`); - logger.info(`[Snapshot] state.environmentCaptured: ${state.environmentCaptured}`); - - // Environment is captured before execution starts via captureEnvironmentBeforeExecution() - return state.environment; - } - - getExecutionMetadata(notebookUri: string): Execution | undefined { - const state = this.executionStates.get(notebookUri); - - if (!state) { - return; - } - - // Don't return execution metadata if no cells have been executed - if (state.blocksExecuted === 0) { - return; - } - - const execution: Execution = { - finishedAt: state.finishedAt || state.startedAt, - startedAt: state.startedAt, - summary: { - blocksExecuted: state.blocksExecuted, - blocksFailed: state.blocksFailed, - blocksSucceeded: state.blocksSucceeded, - totalDurationMs: state.totalDurationMs - }, - triggeredBy: 'user' - }; - - if (state.error) { - execution.error = state.error; - } - - return execution; - } - - recordCellExecutionEnd( - notebookUri: string, - cellId: string, - endTime: number, - success: boolean, - error?: ExecutionError - ): void { - const state = this.executionStates.get(notebookUri); - - if (!state) { - logger.warn(`[Snapshot] No execution state found for notebook ${notebookUri}`); - return; - } - - const isoTimestamp = new Date(endTime).toISOString(); - - const cellMetadata = state.cellMetadata.get(cellId); - - if (cellMetadata) { - cellMetadata.executionFinishedAt = isoTimestamp; - } - - state.blocksExecuted++; - - if (success) { - state.blocksSucceeded++; - } else { - state.blocksFailed++; - - if (error) { - state.error = error; - } - } - - state.finishedAt = isoTimestamp; - - const startMs = new Date(state.startedAt).getTime(); - - state.totalDurationMs = endTime - startMs; - - logger.trace(`[Snapshot] Cell ${cellId} execution ended at ${isoTimestamp} (success: ${success})`); - } - - recordCellExecutionStart(notebookUri: string, cellId: string, startTime: number): void { - const state = this.getOrCreateExecutionState(notebookUri, startTime); - const isoTimestamp = new Date(startTime).toISOString(); - - // Create or update cell metadata - const cellMetadata = state.cellMetadata.get(cellId) || { contentHash: '' }; - - cellMetadata.executionStartedAt = isoTimestamp; - - delete cellMetadata.executionFinishedAt; - - state.cellMetadata.set(cellId, cellMetadata); - - logger.trace(`[Snapshot] Cell ${cellId} execution started at ${isoTimestamp}`); - } - - /** - * Updates the content hash for a cell. - * Called during serialization when we compute the hash. - */ - updateContentHash(notebookUri: string, cellId: string, contentHash: string): void { - const state = this.executionStates.get(notebookUri); - - if (!state) { - return; - } - - const cellMetadata = state.cellMetadata.get(cellId) || { - contentHash: '', - executionFinishedAt: undefined, - executionStartedAt: undefined - }; - cellMetadata.contentHash = contentHash; - state.cellMetadata.set(cellId, cellMetadata); - } - - private async captureEnvironmentForNotebook(notebookUri: string): Promise { - logger.info(`[Snapshot] captureEnvironmentForNotebook called for ${notebookUri}`); - - const state = this.executionStates.get(notebookUri); - - if (!state) { - logger.info(`[Snapshot] Skipping capture: no state found`); - - return; - } - - if (state.environmentCaptured) { - logger.info(`[Snapshot] Skipping capture: already captured`); - - return; - } - - try { - // Find the notebook document to get its resource for interpreter resolution - const notebook = workspace.notebookDocuments.find((n) => n.uri.toString() === notebookUri); - - if (!notebook) { - logger.info(`[Snapshot] Could not find notebook document for ${notebookUri}`); - logger.info( - `[Snapshot] Available notebooks: ${workspace.notebookDocuments - .map((d) => d.uri.toString()) - .join(', ')}` - ); - state.environmentCaptured = true; // Mark as captured to prevent retries - - return; - } - - logger.info(`[Snapshot] Found notebook, getting interpreter...`); - - const environment = await this.environmentCapture.captureEnvironment(notebook.uri); - - if (environment) { - state.environment = environment; - logger.info(`[Snapshot] Captured environment successfully for ${notebookUri}`); - } else { - logger.info(`[Snapshot] environmentCapture returned undefined`); - } - - // Mark as captured only after completion (success or failure) - state.environmentCaptured = true; - } catch (error) { - logger.error('[Snapshot] Failed to capture environment', error); - state.environmentCaptured = true; // Mark as captured to prevent retries - } - } - - private getOrCreateExecutionState(notebookUri: string, startTime: number): NotebookExecutionState { - let state = this.executionStates.get(notebookUri); - - if (!state) { - state = { - blocksFailed: 0, - blocksExecuted: 0, - blocksSucceeded: 0, - cellMetadata: new Map(), - environmentCaptured: false, - startedAt: new Date(startTime).toISOString(), - totalDurationMs: 0 - }; - this.executionStates.set(notebookUri, state); - logger.trace(`[Snapshot] Created new execution state for ${notebookUri}`); - } - - return state; - } - - private handleCellExecutionStateChange(cell: NotebookCell, state: NotebookCellExecutionState): void { - const notebookUri = cell.notebook.uri.toString(); - const cellId = cell.metadata?.id as string | undefined; - - if (!cellId) { - return; - } - - if (state === NotebookCellExecutionState.Executing) { - // Cell started executing - record start time - const startTime = Date.now(); - this.recordCellExecutionStart(notebookUri, cellId, startTime); - } else if (state === NotebookCellExecutionState.Idle) { - // Cell finished executing - record end time and success/failure - const endTime = Date.now(); - const executionSummary = cell.executionSummary; - - // Use VSCode's timing if available (more accurate) - const actualEndTime = executionSummary?.timing?.endTime || endTime; - - // Determine success based on execution summary - // If success is undefined, treat as success (cell may have had no output) - const success = executionSummary?.success !== false; - - this.recordCellExecutionEnd(notebookUri, cellId, actualEndTime, success); - } - } -} diff --git a/src/notebooks/deepnote/snapshotMetadataService.unit.test.ts b/src/notebooks/deepnote/snapshotMetadataService.unit.test.ts deleted file mode 100644 index 80990100bb..0000000000 --- a/src/notebooks/deepnote/snapshotMetadataService.unit.test.ts +++ /dev/null @@ -1,278 +0,0 @@ -import { assert } from 'chai'; -import { instance, mock } from 'ts-mockito'; - -import { SnapshotMetadataService } from './snapshotMetadataService'; -import { IEnvironmentCapture } from './environmentCapture.node'; -import { IDisposableRegistry } from '../../platform/common/types'; - -suite('SnapshotMetadataService', () => { - let service: SnapshotMetadataService; - let mockEnvironmentCapture: IEnvironmentCapture; - let mockDisposables: IDisposableRegistry; - - const notebookUri = 'file:///path/to/notebook.deepnote'; - const cellId = 'cell-123'; - - setup(() => { - mockEnvironmentCapture = mock(); - mockDisposables = []; - - service = new SnapshotMetadataService(instance(mockEnvironmentCapture), mockDisposables); - }); - - suite('recordCellExecutionStart', () => { - test('should record cell execution start time', () => { - const startTime = Date.now(); - - service.recordCellExecutionStart(notebookUri, cellId, startTime); - - const metadata = service.getBlockExecutionMetadata(notebookUri, cellId); - assert.isDefined(metadata); - assert.isDefined(metadata!.executionStartedAt); - assert.isUndefined(metadata!.executionFinishedAt); - }); - - test('should initialize notebook execution state', () => { - const startTime = Date.now(); - - service.recordCellExecutionStart(notebookUri, cellId, startTime); - - const executionMetadata = service.getExecutionMetadata(notebookUri); - // Should not have execution metadata yet since no cells have completed - assert.isUndefined(executionMetadata); - }); - - test('should handle multiple cells in same notebook', () => { - const startTime = Date.now(); - - service.recordCellExecutionStart(notebookUri, 'cell-1', startTime); - service.recordCellExecutionStart(notebookUri, 'cell-2', startTime + 1000); - - const metadata1 = service.getBlockExecutionMetadata(notebookUri, 'cell-1'); - const metadata2 = service.getBlockExecutionMetadata(notebookUri, 'cell-2'); - - assert.isDefined(metadata1); - assert.isDefined(metadata2); - assert.notStrictEqual(metadata1!.executionStartedAt, metadata2!.executionStartedAt); - }); - }); - - suite('recordCellExecutionEnd', () => { - test('should record successful cell execution end', () => { - const startTime = Date.now(); - const endTime = startTime + 1000; - - service.recordCellExecutionStart(notebookUri, cellId, startTime); - service.recordCellExecutionEnd(notebookUri, cellId, endTime, true); - - const metadata = service.getBlockExecutionMetadata(notebookUri, cellId); - assert.isDefined(metadata); - assert.isDefined(metadata!.executionStartedAt); - assert.isDefined(metadata!.executionFinishedAt); - }); - - test('should update execution summary on success', () => { - const startTime = Date.now(); - const endTime = startTime + 1000; - - service.recordCellExecutionStart(notebookUri, cellId, startTime); - service.recordCellExecutionEnd(notebookUri, cellId, endTime, true); - - const executionMetadata = service.getExecutionMetadata(notebookUri); - assert.isDefined(executionMetadata); - assert.isDefined(executionMetadata!.summary); - assert.strictEqual(executionMetadata!.summary!.blocksExecuted, 1); - assert.strictEqual(executionMetadata!.summary!.blocksSucceeded, 1); - assert.strictEqual(executionMetadata!.summary!.blocksFailed, 0); - }); - - test('should update execution summary on failure', () => { - const startTime = Date.now(); - const endTime = startTime + 1000; - - service.recordCellExecutionStart(notebookUri, cellId, startTime); - service.recordCellExecutionEnd(notebookUri, cellId, endTime, false); - - const executionMetadata = service.getExecutionMetadata(notebookUri); - assert.isDefined(executionMetadata); - assert.isDefined(executionMetadata!.summary); - assert.strictEqual(executionMetadata!.summary!.blocksExecuted, 1); - assert.strictEqual(executionMetadata!.summary!.blocksSucceeded, 0); - assert.strictEqual(executionMetadata!.summary!.blocksFailed, 1); - }); - - test('should record error details on failure', () => { - const startTime = Date.now(); - const endTime = startTime + 1000; - const error = { name: 'TypeError', message: 'undefined is not a function' }; - - service.recordCellExecutionStart(notebookUri, cellId, startTime); - service.recordCellExecutionEnd(notebookUri, cellId, endTime, false, error); - - const executionMetadata = service.getExecutionMetadata(notebookUri); - assert.isDefined(executionMetadata); - assert.isDefined(executionMetadata!.error); - assert.strictEqual(executionMetadata!.error!.name, 'TypeError'); - assert.strictEqual(executionMetadata!.error!.message, 'undefined is not a function'); - }); - - test('should accumulate multiple cell executions', () => { - const startTime = Date.now(); - - // Execute 3 cells: 2 successful, 1 failed - service.recordCellExecutionStart(notebookUri, 'cell-1', startTime); - service.recordCellExecutionEnd(notebookUri, 'cell-1', startTime + 100, true); - - service.recordCellExecutionStart(notebookUri, 'cell-2', startTime + 200); - service.recordCellExecutionEnd(notebookUri, 'cell-2', startTime + 300, true); - - service.recordCellExecutionStart(notebookUri, 'cell-3', startTime + 400); - service.recordCellExecutionEnd(notebookUri, 'cell-3', startTime + 500, false); - - const executionMetadata = service.getExecutionMetadata(notebookUri); - assert.isDefined(executionMetadata); - assert.isDefined(executionMetadata!.summary); - assert.strictEqual(executionMetadata!.summary!.blocksExecuted, 3); - assert.strictEqual(executionMetadata!.summary!.blocksSucceeded, 2); - assert.strictEqual(executionMetadata!.summary!.blocksFailed, 1); - }); - - test('should calculate total duration', () => { - const startTime = Date.now(); - const endTime = startTime + 5000; - - service.recordCellExecutionStart(notebookUri, cellId, startTime); - service.recordCellExecutionEnd(notebookUri, cellId, endTime, true); - - const executionMetadata = service.getExecutionMetadata(notebookUri); - assert.isDefined(executionMetadata); - assert.isDefined(executionMetadata!.summary); - assert.strictEqual(executionMetadata!.summary!.totalDurationMs, 5000); - }); - }); - - suite('getExecutionMetadata', () => { - test('should return undefined for unknown notebook', () => { - const metadata = service.getExecutionMetadata('unknown-notebook'); - assert.isUndefined(metadata); - }); - - test('should return undefined if no cells have been executed', () => { - const startTime = Date.now(); - service.recordCellExecutionStart(notebookUri, cellId, startTime); - - const metadata = service.getExecutionMetadata(notebookUri); - assert.isUndefined(metadata); - }); - - test('should include ISO timestamps', () => { - const startTime = Date.now(); - service.recordCellExecutionStart(notebookUri, cellId, startTime); - service.recordCellExecutionEnd(notebookUri, cellId, startTime + 1000, true); - - const metadata = service.getExecutionMetadata(notebookUri); - assert.isDefined(metadata); - assert.isDefined(metadata!.startedAt); - assert.isDefined(metadata!.finishedAt); - // Should be valid ISO date strings - assert.doesNotThrow(() => new Date(metadata!.startedAt!)); - assert.doesNotThrow(() => new Date(metadata!.finishedAt!)); - }); - }); - - suite('getBlockExecutionMetadata', () => { - test('should return undefined for unknown notebook', () => { - const metadata = service.getBlockExecutionMetadata('unknown-notebook', cellId); - assert.isUndefined(metadata); - }); - - test('should return undefined for unknown cell', () => { - const startTime = Date.now(); - service.recordCellExecutionStart(notebookUri, cellId, startTime); - - const metadata = service.getBlockExecutionMetadata(notebookUri, 'unknown-cell'); - assert.isUndefined(metadata); - }); - }); - - suite('updateContentHash', () => { - test('should update content hash for existing cell', () => { - const startTime = Date.now(); - service.recordCellExecutionStart(notebookUri, cellId, startTime); - - service.updateContentHash(notebookUri, cellId, 'sha256:abc123'); - - const metadata = service.getBlockExecutionMetadata(notebookUri, cellId); - assert.isDefined(metadata); - assert.strictEqual(metadata!.contentHash, 'sha256:abc123'); - }); - - test('should not fail for unknown notebook', () => { - // Should not throw - service.updateContentHash('unknown-notebook', cellId, 'md5:abc123'); - }); - }); - - suite('clearExecutionState', () => { - test('should clear all state for a notebook', () => { - const startTime = Date.now(); - service.recordCellExecutionStart(notebookUri, cellId, startTime); - service.recordCellExecutionEnd(notebookUri, cellId, startTime + 1000, true); - - service.clearExecutionState(notebookUri); - - const executionMetadata = service.getExecutionMetadata(notebookUri); - const blockMetadata = service.getBlockExecutionMetadata(notebookUri, cellId); - - assert.isUndefined(executionMetadata); - assert.isUndefined(blockMetadata); - }); - - test('should only clear state for specified notebook', () => { - const startTime = Date.now(); - const otherNotebookUri = 'file:///other/notebook.deepnote'; - - service.recordCellExecutionStart(notebookUri, cellId, startTime); - service.recordCellExecutionEnd(notebookUri, cellId, startTime + 1000, true); - - service.recordCellExecutionStart(otherNotebookUri, 'other-cell', startTime); - service.recordCellExecutionEnd(otherNotebookUri, 'other-cell', startTime + 1000, true); - - service.clearExecutionState(notebookUri); - - // First notebook should be cleared - assert.isUndefined(service.getExecutionMetadata(notebookUri)); - - // Second notebook should still have state - assert.isDefined(service.getExecutionMetadata(otherNotebookUri)); - }); - }); - - suite('multiple notebooks', () => { - test('should track state independently for different notebooks', () => { - const notebook1 = 'file:///notebook1.deepnote'; - const notebook2 = 'file:///notebook2.deepnote'; - const startTime = Date.now(); - - // Execute cells in different notebooks - service.recordCellExecutionStart(notebook1, 'cell-1', startTime); - service.recordCellExecutionEnd(notebook1, 'cell-1', startTime + 100, true); - - service.recordCellExecutionStart(notebook2, 'cell-2', startTime); - service.recordCellExecutionEnd(notebook2, 'cell-2', startTime + 200, false); - - const metadata1 = service.getExecutionMetadata(notebook1); - const metadata2 = service.getExecutionMetadata(notebook2); - - assert.isDefined(metadata1); - assert.isDefined(metadata1!.summary); - assert.strictEqual(metadata1!.summary!.blocksSucceeded, 1); - assert.strictEqual(metadata1!.summary!.blocksFailed, 0); - - assert.isDefined(metadata2); - assert.isDefined(metadata2!.summary); - assert.strictEqual(metadata2!.summary!.blocksSucceeded, 0); - assert.strictEqual(metadata2!.summary!.blocksFailed, 1); - }); - }); -}); diff --git a/src/notebooks/deepnote/environmentCapture.node.ts b/src/notebooks/deepnote/snapshots/environmentCapture.node.ts similarity index 94% rename from src/notebooks/deepnote/environmentCapture.node.ts rename to src/notebooks/deepnote/snapshots/environmentCapture.node.ts index 8b12fcae8c..dcd3c923cb 100644 --- a/src/notebooks/deepnote/environmentCapture.node.ts +++ b/src/notebooks/deepnote/snapshots/environmentCapture.node.ts @@ -5,15 +5,15 @@ import { promisify } from 'node:util'; import type { Environment } from '@deepnote/blocks'; -import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; -import { computeHash } from '../../platform/common/crypto'; -import { raceTimeout } from '../../platform/common/utils/async'; -import { logger } from '../../platform/logging'; +import { PythonEnvironment } from '../../../platform/pythonEnvironments/info'; +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, IDeepnoteNotebookEnvironmentMapper } from '../../../kernels/deepnote/types'; import { Uri } from 'vscode'; -import { DeepnoteEnvironment } from '../../kernels/deepnote/environments/deepnoteEnvironment'; -import * as path from '../../platform/vscode-path/path'; +import { DeepnoteEnvironment } from '../../../kernels/deepnote/environments/deepnoteEnvironment'; +import * as path from '../../../platform/vscode-path/path'; const captureTimeoutInMilliseconds = 5_000; diff --git a/src/notebooks/deepnote/environmentCapture.unit.test.ts b/src/notebooks/deepnote/snapshots/environmentCapture.unit.test.ts similarity index 100% rename from src/notebooks/deepnote/environmentCapture.unit.test.ts rename to src/notebooks/deepnote/snapshots/environmentCapture.unit.test.ts diff --git a/src/notebooks/deepnote/pipFileParser.ts b/src/notebooks/deepnote/snapshots/pipFileParser.ts similarity index 100% rename from src/notebooks/deepnote/pipFileParser.ts rename to src/notebooks/deepnote/snapshots/pipFileParser.ts diff --git a/src/notebooks/deepnote/pipFileParser.unit.test.ts b/src/notebooks/deepnote/snapshots/pipFileParser.unit.test.ts similarity index 100% rename from src/notebooks/deepnote/pipFileParser.unit.test.ts rename to src/notebooks/deepnote/snapshots/pipFileParser.unit.test.ts diff --git a/src/notebooks/deepnote/snapshots/snapshotService.ts b/src/notebooks/deepnote/snapshots/snapshotService.ts new file mode 100644 index 0000000000..966d008523 --- /dev/null +++ b/src/notebooks/deepnote/snapshots/snapshotService.ts @@ -0,0 +1,951 @@ +import type { DeepnoteBlock, DeepnoteFile, Environment, Execution, ExecutionError } from '@deepnote/blocks'; +import fastDeepEqual from 'fast-deep-equal'; +import { inject, injectable, optional } from 'inversify'; +import * as yaml from 'js-yaml'; +import { FileType, NotebookCell, RelativePattern, Uri, window, workspace } from 'vscode'; +import { Utils } from 'vscode-uri'; + +import { IExtensionSyncActivationService } from '../../../platform/activation/types'; +import { IDisposableRegistry } from '../../../platform/common/types'; +import type { DeepnoteOutput } from '../../../platform/deepnote/deepnoteTypes'; +import { InvalidProjectNameError } from '../../../platform/errors/invalidProjectNameError'; +import { logger } from '../../../platform/logging'; +import { + notebookCellExecutions, + NotebookCellExecutionState +} from '../../../platform/notebooks/cellExecutionStateService'; +import { IDeepnoteNotebookManager } from '../../types'; +import { DeepnoteDataConverter } from '../deepnoteDataConverter'; +import { IEnvironmentCapture } from './environmentCapture.node'; + +/** + * Platform-layer interface for snapshot metadata service. + * Used by the kernel execution layer to capture environment before cell execution. + */ +export const ISnapshotMetadataService = Symbol('ISnapshotMetadataService'); +export interface ISnapshotMetadataService { + /** + * Capture environment before execution starts. + * Called at the start of a cell execution batch. + * This is blocking and should complete before cells execute. + */ + captureEnvironmentBeforeExecution(notebookUri: string): Promise; + + /** + * Clear execution state for a notebook (e.g., when kernel restarts). + */ + clearExecutionState(notebookUri: string): void; + + /** + * Set "Run All" mode for a notebook. + * When execution completes, a full snapshot (timestamped + latest) will be created. + */ + setRunAllMode(notebookUri: string): void; + + /** + * Check if "Run All" mode is set for a notebook. + */ + isRunAllMode(notebookUri: string): boolean; +} + +/** + * Block-level execution metadata. + */ +export interface BlockExecutionMetadata { + /** SHA-256 hash of block source code (prefixed with "sha256:") */ + contentHash: string; + + /** ISO 8601 timestamp when block execution started */ + executionStartedAt?: string; + + /** ISO 8601 timestamp when block execution completed */ + executionFinishedAt?: string; +} + +/** + * Internal state tracking for a notebook execution session. + */ +interface NotebookExecutionState { + /** Number of blocks executed so far */ + blocksExecuted: number; + + /** Number of blocks that failed */ + blocksFailed: number; + + /** Number of blocks that succeeded */ + blocksSucceeded: number; + + /** Promise that resolves when environment capture completes */ + capturePromise?: Promise; + + /** Per-cell execution metadata, keyed by cell ID */ + cellMetadata: Map; + + /** Cached environment metadata */ + environment?: Environment; + + /** Whether environment has been captured for this session */ + environmentCaptured: boolean; + + /** Top-level error if any */ + error?: { name?: string; message?: string; traceback?: string[] }; + + /** ISO 8601 timestamp when last cell finished executing */ + finishedAt?: string; + + /** ISO 8601 timestamp when first cell started executing */ + startedAt: string; + + /** Total duration in milliseconds */ + totalDurationMs: number; +} + +class TimeoutError extends Error { + constructor(message: string) { + super(message); + this.name = 'TimeoutError'; + } +} + +/** + * Slugifies a project name for use in filenames. + * Converts to lowercase, replaces spaces with hyphens, removes non-alphanumeric chars. + * @throws Error if the result is empty after transformation + */ +function slugifyProjectName(name: string): string { + if (typeof name !== 'string' || !name.trim()) { + throw new InvalidProjectNameError(); + } + + const slug = name + .toLowerCase() + .replace(/\s+/g, '-') + .replace(/[^a-z0-9-]/g, '') + .replace(/-+/g, '-') + .replace(/^-|-$/g, ''); + + if (!slug) { + throw new InvalidProjectNameError(); + } + + return slug; +} + +/** + * Generates a timestamp string for snapshot filenames. + * Format: 2025-12-11T10-31-48 (ISO 8601 with colons replaced by hyphens) + */ +function generateTimestamp(): string { + return new Date().toISOString().replace(/:/g, '-').slice(0, 19); +} + +/** + * Unified service for managing Deepnote notebook snapshots. + * Handles both file I/O operations (reading/writing snapshot files) and + * execution metadata tracking (timing, environment capture). + */ +@injectable() +export class SnapshotService implements ISnapshotMetadataService, IExtensionSyncActivationService { + private readonly converter = new DeepnoteDataConverter(); + private readonly executionStates = new Map(); + private readonly runAllModeNotebooks = new Set(); + + constructor( + @inject(IEnvironmentCapture) private readonly environmentCapture: IEnvironmentCapture, + @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, + @inject(IDeepnoteNotebookManager) @optional() private readonly notebookManager?: IDeepnoteNotebookManager + ) {} + + activate(): void { + logger.info('[Snapshot] SnapshotService activated'); + + workspace.onDidCloseNotebookDocument( + (notebook) => { + this.clearExecutionState(notebook.uri.toString()); + }, + this, + this.disposables + ); + + notebookCellExecutions.onDidChangeNotebookCellExecutionState( + (e) => { + logger.debug(`[Snapshot] Cell execution state changed: ${e.state} for cell ${e.cell.metadata?.id}`); + this.handleCellExecutionStateChange(e.cell, e.state); + }, + this, + this.disposables + ); + + notebookCellExecutions.onDidCompleteQueueExecution( + async (e) => { + logger.debug(`[Snapshot] Queue execution complete for ${e.notebookUri}`); + await this.onExecutionComplete(e.notebookUri); + }, + this, + this.disposables + ); + } + + async captureEnvironmentBeforeExecution(notebookUri: string): Promise { + logger.info(`[Snapshot] captureEnvironmentBeforeExecution called for ${notebookUri}`); + + const state = this.getOrCreateExecutionState(notebookUri, Date.now()); + + // If capture is already in progress, wait for it + if (state.capturePromise) { + logger.info(`[Snapshot] Capture already in progress, waiting...`); + await state.capturePromise; + + return; + } + + // Start capture and store the promise so other callers can wait for it + state.capturePromise = this.captureEnvironmentForNotebook(notebookUri); + } + + clearExecutionState(notebookUri: string): void { + this.executionStates.delete(notebookUri); + this.runAllModeNotebooks.delete(notebookUri); + + logger.trace(`[Snapshot] Cleared execution state for ${notebookUri}`); + } + + async createSnapshot( + projectUri: Uri, + projectId: string, + projectName: string, + projectData: DeepnoteFile + ): Promise { + const prepared = await this.prepareSnapshotData(projectUri, projectId, projectName, projectData); + + if (!prepared) { + logger.debug(`[Snapshot] No changes detected, skipping snapshot creation`); + + return; + } + + const { latestPath, content } = prepared; + const timestamp = generateTimestamp(); + const timestampedPath = this.buildSnapshotPath(projectUri, projectId, projectName, timestamp); + + // Write to timestamped file first (safe - doesn't touch existing files) + try { + await workspace.fs.writeFile(timestampedPath, content); + logger.debug(`[Snapshot] Wrote timestamped snapshot: ${Utils.basename(timestampedPath)}`); + } catch (error) { + logger.error(`[Snapshot] Failed to write timestamped snapshot: ${Utils.basename(timestampedPath)}`, error); + + const message = error instanceof Error ? error.message : String(error); + + await window.showErrorMessage(`Failed to create snapshot: ${message}`); + + return; + } + + // Copy timestamped file to 'latest' pointer + try { + await workspace.fs.copy(timestampedPath, latestPath, { overwrite: true }); + + logger.debug(`[Snapshot] Updated latest snapshot: ${Utils.basename(latestPath)}`); + } catch (error) { + logger.warn( + `[Snapshot] Wrote timestamped snapshot but failed to update latest pointer: ${Utils.basename( + latestPath + )}. ` + `Timestamped snapshot available at: ${Utils.basename(timestampedPath)}`, + error + ); + + const message = error instanceof Error ? error.message : String(error); + + await window.showErrorMessage(`Failed to update latest snapshot pointer: ${message}`); + + return; + } + + return timestampedPath; + } + + extractOutputsFromBlocks(blocks: DeepnoteBlock[]): Map { + const outputsMap = new Map(); + + for (const block of blocks) { + if (block.id && block.outputs) { + outputsMap.set(block.id, block.outputs as DeepnoteOutput[]); + } + } + + return outputsMap; + } + + getBlockExecutionMetadata(notebookUri: string, cellId: string): BlockExecutionMetadata | undefined { + const state = this.executionStates.get(notebookUri); + + if (!state) { + return; + } + + return state.cellMetadata.get(cellId); + } + + async getEnvironmentMetadata(notebookUri: string): Promise { + const state = this.executionStates.get(notebookUri); + + logger.info(`[Snapshot] getEnvironmentMetadata for ${notebookUri}`); + logger.info(Boolean(state) ? '[Snapshot] State exists.' : '[Snapshot] No state found.'); + + if (!state) { + logger.info(`[Snapshot] Available URIs: ${Array.from(this.executionStates.keys()).join(', ')}`); + + return; + } + + // If capture is in progress, wait for it to complete + if (state.capturePromise && !state.environmentCaptured) { + logger.info(`[Snapshot] Waiting for capture to complete before returning metadata.`); + + const timeoutPromise = new Promise((_, reject) => { + setTimeout(() => reject(new TimeoutError('Timeout waiting for environment capture.')), 10_000); + }); + + await Promise.race([state.capturePromise, timeoutPromise]).catch((error) => { + if (error instanceof TimeoutError) { + logger.warn('[Snapshot] Timed out waiting for environment capture'); + } else { + throw error; + } + }); + } + + logger.info(`[Snapshot] state.environment exists: ${!!state.environment}`); + logger.info(`[Snapshot] state.environmentCaptured: ${state.environmentCaptured}`); + + return state.environment; + } + + getExecutionMetadata(notebookUri: string): Execution | undefined { + const state = this.executionStates.get(notebookUri); + + if (!state) { + return; + } + + // Don't return execution metadata if no cells have been executed + if (state.blocksExecuted === 0) { + return; + } + + const execution: Execution = { + finishedAt: state.finishedAt || state.startedAt, + startedAt: state.startedAt, + summary: { + blocksExecuted: state.blocksExecuted, + blocksFailed: state.blocksFailed, + blocksSucceeded: state.blocksSucceeded, + totalDurationMs: state.totalDurationMs + }, + triggeredBy: 'user' + }; + + if (state.error) { + execution.error = state.error; + } + + return execution; + } + + isRunAllMode(notebookUri: string): boolean { + return this.runAllModeNotebooks.has(notebookUri); + } + + isSnapshotsEnabled(): boolean { + const config = workspace.getConfiguration('deepnote'); + + return config.get('snapshots.enabled', false); + } + + mergeOutputsIntoBlocks(blocks: DeepnoteBlock[], outputs: Map): DeepnoteBlock[] { + let mergedCount = 0; + + const mergedBlocks = blocks.map((block) => { + if (!block.id) { + return block; + } + + const blockOutputs = outputs.get(block.id); + + if (blockOutputs !== undefined) { + mergedCount++; + + return { ...block, outputs: blockOutputs }; + } + + return block; + }); + + logger.debug(`[Snapshot] Merged outputs into ${mergedCount}/${blocks.length} blocks`); + + return mergedBlocks; + } + + async readSnapshot(projectId: string): Promise | undefined> { + const workspaceFolders = workspace.workspaceFolders; + + if (!workspaceFolders || workspaceFolders.length === 0) { + logger.debug(`[Snapshot] No workspace folders found`); + + await window.showWarningMessage('Cannot read snapshot: No workspace folders found.'); + + return; + } + + // 1. Try to find a 'latest' snapshot file + const latestGlob = `**/snapshots/*_${projectId}_latest.snapshot.deepnote`; + + for (const folder of workspaceFolders) { + const latestPattern = new RelativePattern(folder, latestGlob); + const latestFiles = await workspace.findFiles(latestPattern, null, 1); + + if (latestFiles.length > 0) { + logger.debug(`[Snapshot] Found latest snapshot: ${Utils.basename(latestFiles[0])}`); + + try { + return await this.parseSnapshotFile(latestFiles[0]); + } catch (error) { + logger.error(`[Snapshot] Failed to parse snapshot: ${Utils.basename(latestFiles[0])}`, error); + + await window.showErrorMessage(`Failed to read latest snapshot: ${Utils.basename(latestFiles[0])}`); + + return; + } + } + } + + 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`; + let allTimestampedFiles: Uri[] = []; + + for (const folder of workspaceFolders) { + const timestampedPattern = new RelativePattern(folder, timestampedGlob); + const files = await workspace.findFiles(timestampedPattern, null, 100); + + allTimestampedFiles = allTimestampedFiles.concat(files); + } + + // Filter out 'latest' files and sort by filename descending + const sortedFiles = allTimestampedFiles + .filter((uri) => !Utils.basename(uri).endsWith('_latest.snapshot.deepnote')) + .sort((a, b) => { + const nameA = Utils.basename(a); + const nameB = Utils.basename(b); + + return nameB.localeCompare(nameA); + }); + + if (sortedFiles.length === 0) { + logger.debug(`[Snapshot] No timestamped snapshots found`); + + return; + } + + const newestFile = sortedFiles[0]; + + logger.debug(`[Snapshot] Using timestamped snapshot: ${Utils.basename(newestFile)}`); + + try { + return await this.parseSnapshotFile(newestFile); + } catch (error) { + logger.error(`[Snapshot] Failed to parse snapshot: ${Utils.basename(newestFile)}`, error); + + return; + } + } + + setRunAllMode(notebookUri: string): void { + this.runAllModeNotebooks.add(notebookUri); + + logger.debug(`[Snapshot] Set Run All mode for ${notebookUri}`); + } + + stripOutputsFromBlocks(blocks: DeepnoteBlock[]): DeepnoteBlock[] { + return blocks.map((block) => { + // eslint-disable-next-line @typescript-eslint/no-unused-vars + const { outputs, ...strippedBlock } = block; + + return strippedBlock as DeepnoteBlock; + }); + } + + private buildSnapshotPath( + projectUri: Uri, + projectId: string, + projectName: string, + variant: 'latest' | string + ): Uri { + const parentDir = Uri.joinPath(projectUri, '..'); + const slug = slugifyProjectName(projectName); + const filename = `${slug}_${projectId}_${variant}.snapshot.deepnote`; + + return Uri.joinPath(parentDir, 'snapshots', filename); + } + + private async captureEnvironmentForNotebook(notebookUri: string): Promise { + logger.info(`[Snapshot] captureEnvironmentForNotebook called for ${notebookUri}`); + + const state = this.executionStates.get(notebookUri); + + if (!state) { + logger.info(`[Snapshot] Skipping capture: no state found`); + + return; + } + + if (state.environmentCaptured) { + logger.info(`[Snapshot] Skipping capture: already captured`); + + return; + } + + try { + const notebook = workspace.notebookDocuments.find((n) => n.uri.toString() === notebookUri); + + if (!notebook) { + logger.info(`[Snapshot] Could not find notebook document for ${notebookUri}`); + logger.info( + `[Snapshot] Available notebooks: ${workspace.notebookDocuments + .map((d) => d.uri.toString()) + .join(', ')}` + ); + state.environmentCaptured = true; + + return; + } + + logger.info(`[Snapshot] Found notebook, getting interpreter...`); + + const environment = await this.environmentCapture.captureEnvironment(notebook.uri); + + if (environment) { + state.environment = environment; + logger.info(`[Snapshot] Captured environment successfully for ${notebookUri}`); + } else { + logger.info(`[Snapshot] environmentCapture returned undefined`); + } + + state.environmentCaptured = true; + } catch (error) { + logger.error('[Snapshot] Failed to capture environment', error); + state.environmentCaptured = true; + } + } + + private async ensureSnapshotsDirectory(snapshotsDir: Uri): Promise { + try { + const stat = await workspace.fs.stat(snapshotsDir); + + if (stat.type !== FileType.Directory) { + logger.error( + `[Snapshot] Snapshots path exists but is not a directory: ${Utils.basename(snapshotsDir)}` + ); + + return false; + } + + return true; + } catch { + logger.debug(`[Snapshot] Creating snapshots directory: ${Utils.basename(snapshotsDir)}`); + + try { + await workspace.fs.createDirectory(snapshotsDir); + + return true; + } catch (error) { + logger.error(`[Snapshot] Failed to create snapshots directory: ${Utils.basename(snapshotsDir)}`, error); + + return false; + } + } + } + + private findProjectUriFromId(projectId: string): Uri | undefined { + const notebookDoc = workspace.notebookDocuments.find( + (doc) => doc.notebookType === 'deepnote' && doc.metadata?.deepnoteProjectId === projectId + ); + + return notebookDoc?.uri.with({ query: '' }); + } + + private getComparableProjectContent(data: DeepnoteFile): object { + return { + version: data.version, + project: data.project, + environment: data.environment, + execution: data.execution + }; + } + + private getOrCreateExecutionState(notebookUri: string, startTime: number): NotebookExecutionState { + let state = this.executionStates.get(notebookUri); + + if (!state) { + state = { + blocksFailed: 0, + blocksExecuted: 0, + blocksSucceeded: 0, + cellMetadata: new Map(), + environmentCaptured: false, + startedAt: new Date(startTime).toISOString(), + totalDurationMs: 0 + }; + this.executionStates.set(notebookUri, state); + logger.trace(`[Snapshot] Created new execution state for ${notebookUri}`); + } + + return state; + } + + private handleCellExecutionStateChange(cell: NotebookCell, state: NotebookCellExecutionState): void { + const notebookUri = cell.notebook.uri.toString(); + const cellId = cell.metadata?.id as string | undefined; + + if (!cellId) { + return; + } + + if (state === NotebookCellExecutionState.Executing) { + const startTime = Date.now(); + + this.recordCellExecutionStart(notebookUri, cellId, startTime); + } else if (state === NotebookCellExecutionState.Idle) { + const endTime = Date.now(); + const executionSummary = cell.executionSummary; + const actualEndTime = executionSummary?.timing?.endTime || endTime; + const success = executionSummary?.success !== false; + + this.recordCellExecutionEnd(notebookUri, cellId, actualEndTime, success); + } + } + + private async hasSnapshotChanges(latestPath: Uri, projectData: DeepnoteFile): Promise { + try { + const existingContent = await workspace.fs.readFile(latestPath); + const existingString = new TextDecoder('utf-8').decode(existingContent); + const existingData = yaml.load(existingString) as DeepnoteFile; + + const existingProject = this.getComparableProjectContent(existingData); + const newProject = this.getComparableProjectContent(projectData); + + return !fastDeepEqual(existingProject, newProject); + } catch { + logger.debug(`[Snapshot] No existing snapshot found, treating as changed`); + + return true; + } + } + + private async onExecutionComplete(notebookUri: string): Promise { + logger.debug(`[Snapshot] onExecutionComplete called for ${notebookUri}`); + + if (!this.isSnapshotsEnabled()) { + logger.debug(`[Snapshot] Snapshots not enabled, skipping`); + this.runAllModeNotebooks.delete(notebookUri); + + return; + } + + const notebook = workspace.notebookDocuments.find((n) => n.uri.toString() === notebookUri); + + if (!notebook) { + logger.warn(`[Snapshot] Could not find notebook document for ${notebookUri}`); + this.runAllModeNotebooks.delete(notebookUri); + + return; + } + + const projectId = notebook.metadata?.deepnoteProjectId as string | undefined; + + if (!projectId) { + logger.warn(`[Snapshot] No project ID in notebook metadata`); + this.runAllModeNotebooks.delete(notebookUri); + + return; + } + + const originalProject = this.notebookManager?.getOriginalProject(projectId); + + if (!originalProject) { + logger.warn(`[Snapshot] No original project found for ${projectId}`); + this.runAllModeNotebooks.delete(notebookUri); + + return; + } + + const projectUri = this.findProjectUriFromId(projectId); + + if (!projectUri) { + logger.warn(`[Snapshot] Could not find project URI for ${projectId}`); + this.runAllModeNotebooks.delete(notebookUri); + + return; + } + + const notebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; + + if (!notebookId) { + logger.warn(`[Snapshot] No notebook ID in notebook metadata`); + this.runAllModeNotebooks.delete(notebookUri); + + return; + } + + const deepnoteNotebook = originalProject.project.notebooks?.find((nb) => nb.id === notebookId); + + if (!deepnoteNotebook) { + logger.warn(`[Snapshot] Notebook ${notebookId} not found in project`); + this.runAllModeNotebooks.delete(notebookUri); + + return; + } + + const cellData = notebook.getCells().map((cell) => ({ + kind: cell.kind, + value: cell.document.getText(), + languageId: cell.document.languageId, + metadata: cell.metadata, + outputs: [...cell.outputs] + })); + const blocks = this.converter.convertCellsToBlocks(cellData); + + const snapshotProject = structuredClone(originalProject) as DeepnoteFile; + const snapshotNotebook = snapshotProject.project.notebooks?.find((nb) => nb.id === notebookId); + + if (snapshotNotebook) { + snapshotNotebook.blocks = blocks as DeepnoteBlock[]; + } + + const isRunAll = this.runAllModeNotebooks.has(notebookUri); + + if (isRunAll) { + logger.debug(`[Snapshot] Creating full snapshot (Run All mode)`); + + const snapshotUri = await this.createSnapshot( + projectUri, + projectId, + originalProject.project.name, + snapshotProject + ); + + if (snapshotUri) { + logger.info(`[Snapshot] Created full snapshot: ${snapshotUri.toString()}`); + } + } else { + logger.debug(`[Snapshot] Updating latest snapshot only (partial run)`); + + const snapshotUri = await this.updateLatestSnapshot( + projectUri, + projectId, + originalProject.project.name, + snapshotProject + ); + + if (snapshotUri) { + logger.info(`[Snapshot] Updated latest snapshot: ${snapshotUri.toString()}`); + } + } + + this.runAllModeNotebooks.delete(notebookUri); + } + + private async parseSnapshotFile(path: Uri): Promise> { + const outputsMap = new Map(); + + let snapshotData: unknown; + + try { + const content = await workspace.fs.readFile(path); + const contentString = new TextDecoder('utf-8').decode(content); + + snapshotData = yaml.load(contentString); + } catch (error) { + logger.error(`[Snapshot] Failed to read or parse snapshot file: ${Utils.basename(path)}`, error); + + return outputsMap; + } + + if (typeof snapshotData !== 'object' || snapshotData === null) { + logger.error(`[Snapshot] Invalid snapshot structure (not an object): ${Utils.basename(path)}`); + + return outputsMap; + } + + const data = snapshotData as DeepnoteFile; + const notebooks = data.project?.notebooks; + + if (!Array.isArray(notebooks)) { + logger.debug(`[Snapshot] No notebooks array in snapshot: ${Utils.basename(path)}`); + + return outputsMap; + } + + for (const notebook of notebooks) { + const blocks = notebook?.blocks; + + if (!Array.isArray(blocks)) { + continue; + } + + for (const block of blocks) { + if ( + typeof block === 'object' && + block !== null && + typeof block.id === 'string' && + Array.isArray(block.outputs) + ) { + outputsMap.set(block.id, block.outputs as DeepnoteOutput[]); + } + } + } + + logger.debug(`[Snapshot] Read ${outputsMap.size} block outputs from snapshot`); + + return outputsMap; + } + + private async prepareSnapshotData( + projectUri: Uri, + projectId: string, + projectName: string, + projectData: DeepnoteFile + ): Promise<{ latestPath: Uri; content: Uint8Array } | undefined> { + let latestPath: Uri; + + try { + latestPath = this.buildSnapshotPath(projectUri, projectId, projectName, 'latest'); + } catch (error) { + if (error instanceof InvalidProjectNameError) { + logger.warn('[Snapshot] Skipping snapshots due to invalid project name', error); + + return; + } + throw error; + } + + const snapshotsDir = Uri.joinPath(latestPath, '..'); + const dirExists = await this.ensureSnapshotsDirectory(snapshotsDir); + + if (!dirExists) { + return; + } + + const hasChanges = await this.hasSnapshotChanges(latestPath, projectData); + + if (!hasChanges) { + return; + } + + const snapshotData = structuredClone(projectData); + + snapshotData.metadata = snapshotData.metadata || {}; + if (!snapshotData.metadata.createdAt) { + snapshotData.metadata.createdAt = new Date().toISOString(); + } + snapshotData.metadata.modifiedAt = new Date().toISOString(); + + const yamlString = yaml.dump(snapshotData, { + indent: 2, + lineWidth: -1, + noRefs: true, + sortKeys: false + }); + + const content = new TextEncoder().encode(yamlString); + + return { latestPath, content }; + } + + private recordCellExecutionEnd( + notebookUri: string, + cellId: string, + endTime: number, + success: boolean, + error?: ExecutionError + ): void { + const state = this.executionStates.get(notebookUri); + + if (!state) { + logger.warn(`[Snapshot] No execution state found for notebook ${notebookUri}`); + + return; + } + + const isoTimestamp = new Date(endTime).toISOString(); + const cellMetadata = state.cellMetadata.get(cellId); + + if (cellMetadata) { + cellMetadata.executionFinishedAt = isoTimestamp; + } + + state.blocksExecuted++; + + if (success) { + state.blocksSucceeded++; + } else { + state.blocksFailed++; + + if (error) { + state.error = error; + } + } + + state.finishedAt = isoTimestamp; + + const startMs = new Date(state.startedAt).getTime(); + + state.totalDurationMs = endTime - startMs; + + logger.trace(`[Snapshot] Cell ${cellId} execution ended at ${isoTimestamp} (success: ${success})`); + } + + private recordCellExecutionStart(notebookUri: string, cellId: string, startTime: number): void { + const state = this.getOrCreateExecutionState(notebookUri, startTime); + const isoTimestamp = new Date(startTime).toISOString(); + const cellMetadata = state.cellMetadata.get(cellId) || { contentHash: '' }; + + cellMetadata.executionStartedAt = isoTimestamp; + + delete cellMetadata.executionFinishedAt; + + state.cellMetadata.set(cellId, cellMetadata); + + logger.trace(`[Snapshot] Cell ${cellId} execution started at ${isoTimestamp}`); + } + + private async updateLatestSnapshot( + projectUri: Uri, + projectId: string, + projectName: string, + projectData: DeepnoteFile + ): Promise { + const prepared = await this.prepareSnapshotData(projectUri, projectId, projectName, projectData); + + if (!prepared) { + logger.debug(`[Snapshot] No changes detected, skipping latest snapshot update`); + + return; + } + + const { latestPath, content } = prepared; + + try { + await workspace.fs.writeFile(latestPath, content); + logger.debug(`[Snapshot] Updated latest snapshot: ${Utils.basename(latestPath)}`); + + return latestPath; + } catch (error) { + logger.error(`[Snapshot] Failed to update latest snapshot: ${Utils.basename(latestPath)}`, error); + + return; + } + } +} diff --git a/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts b/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts new file mode 100644 index 0000000000..7820e01c61 --- /dev/null +++ b/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts @@ -0,0 +1,988 @@ +import { assert } from 'chai'; +import { anything, instance, mock, when } from 'ts-mockito'; +import { FileType, Uri, WorkspaceConfiguration, WorkspaceFolder } from 'vscode'; + +import type { DeepnoteBlock, DeepnoteFile } from '@deepnote/blocks'; + +import { IEnvironmentCapture } from './environmentCapture.node'; +import { SnapshotService } from './snapshotService'; +import type { DeepnoteOutput } from '../../../platform/deepnote/deepnoteTypes'; +import { IDisposableRegistry } from '../../../platform/common/types'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../../test/vscode-mock'; + +suite('SnapshotService', () => { + let service: SnapshotService; + // eslint-disable-next-line @typescript-eslint/no-explicit-any + let serviceAny: any; + let mockEnvironmentCapture: IEnvironmentCapture; + let mockDisposables: IDisposableRegistry; + + setup(() => { + resetVSCodeMocks(); + mockEnvironmentCapture = mock(); + mockDisposables = []; + service = new SnapshotService(instance(mockEnvironmentCapture), mockDisposables); + serviceAny = service; + }); + + function createProjectData(projectId = 'test-project-id-123', projectName = 'My Project'): DeepnoteFile { + return { + metadata: { + createdAt: '2025-01-01T00:00:00Z' + }, + version: '1.0', + project: { + id: projectId, + name: projectName, + notebooks: [ + { + id: 'notebook-1', + name: 'Notebook 1', + blocks: [ + { + id: 'block-1', + type: 'code', + sortingKey: 'a0', + content: 'print(1)', + outputs: [{ output_type: 'stream', text: '1' }] + } + ] + } + ] + } + }; + } + + suite('buildSnapshotPath', () => { + test('should build correct path for latest variant', () => { + const projectUri = Uri.file('/path/to/my-project.deepnote'); + const projectId = 'e132b172-b114-410e-8331-011517db664f'; + const projectName = 'My Project'; + + const result = serviceAny.buildSnapshotPath(projectUri, projectId, projectName, 'latest'); + + assert.include(result.fsPath, 'snapshots'); + assert.include(result.fsPath, 'my-project'); + assert.include(result.fsPath, projectId); + assert.include(result.fsPath, 'latest'); + assert.include(result.fsPath, '.snapshot.deepnote'); + }); + + test('should build correct path for timestamped variant', () => { + const projectUri = Uri.file('/path/to/my-project.deepnote'); + const projectId = 'e132b172-b114-410e-8331-011517db664f'; + const projectName = 'My Project'; + const timestamp = '2025-12-11T10-31-48'; + + const result = serviceAny.buildSnapshotPath(projectUri, projectId, projectName, timestamp); + + assert.include(result.fsPath, 'snapshots'); + assert.include(result.fsPath, 'my-project'); + assert.include(result.fsPath, projectId); + assert.include(result.fsPath, timestamp); + assert.include(result.fsPath, '.snapshot.deepnote'); + }); + + test('should slugify project name correctly', () => { + const projectUri = Uri.file('/path/to/file.deepnote'); + const projectId = 'abc-123'; + const projectName = 'Customer Churn ML Playbook!'; + + const result = serviceAny.buildSnapshotPath(projectUri, projectId, projectName, 'latest'); + + assert.include(result.fsPath, 'customer-churn-ml-playbook'); + assert.notInclude(result.fsPath, '!'); + assert.notInclude(result.fsPath, ' '); + }); + + test('should handle project names with special characters', () => { + const projectUri = Uri.file('/path/to/file.deepnote'); + const projectId = 'abc-123'; + const projectName = 'Test@#$%Project'; + + const result = serviceAny.buildSnapshotPath(projectUri, projectId, projectName, 'latest'); + + assert.include(result.fsPath, 'testproject'); + }); + + test('should handle project names with multiple spaces', () => { + const projectUri = Uri.file('/path/to/file.deepnote'); + const projectId = 'abc-123'; + const projectName = 'My Project Name'; + + const result = serviceAny.buildSnapshotPath(projectUri, projectId, projectName, 'latest'); + + assert.include(result.fsPath, 'my-project-name'); + assert.notInclude(result.fsPath, '--'); + }); + + test('should throw error for empty project name', () => { + const projectUri = Uri.file('/path/to/file.deepnote'); + const projectId = 'abc-123'; + const projectName = ''; + + assert.throws( + () => serviceAny.buildSnapshotPath(projectUri, projectId, projectName, 'latest'), + 'Project name cannot be empty or contain only special characters' + ); + }); + + test('should throw error for project name with only special characters', () => { + const projectUri = Uri.file('/path/to/file.deepnote'); + const projectId = 'abc-123'; + const projectName = '@#$%^&*()'; + + assert.throws( + () => serviceAny.buildSnapshotPath(projectUri, projectId, projectName, 'latest'), + 'Project name cannot be empty or contain only special characters' + ); + }); + + test('should throw error for project name with only whitespace', () => { + const projectUri = Uri.file('/path/to/file.deepnote'); + const projectId = 'abc-123'; + const projectName = ' '; + + assert.throws( + () => serviceAny.buildSnapshotPath(projectUri, projectId, projectName, 'latest'), + 'Project name cannot be empty or contain only special characters' + ); + }); + }); + + suite('mergeOutputsIntoBlocks', () => { + test('should merge outputs into blocks by ID', () => { + const blocks: DeepnoteBlock[] = [ + { id: 'block-1', type: 'code', sortingKey: 'a0', content: 'print(1)' }, + { id: 'block-2', type: 'code', sortingKey: 'a1', content: 'print(2)' }, + { id: 'block-3', type: 'markdown', sortingKey: 'a2', content: '# Hello' } + ]; + + const outputs = new Map(); + + outputs.set('block-1', [{ output_type: 'stream', name: 'stdout', text: '1\n' }]); + outputs.set('block-2', [{ output_type: 'stream', name: 'stdout', text: '2\n' }]); + + const result = service.mergeOutputsIntoBlocks(blocks, outputs); + + assert.deepEqual(result[0].outputs, [{ output_type: 'stream', name: 'stdout', text: '1\n' }]); + assert.deepEqual(result[1].outputs, [{ output_type: 'stream', name: 'stdout', text: '2\n' }]); + assert.isUndefined(result[2].outputs); + }); + + test('should not modify original blocks', () => { + const blocks: DeepnoteBlock[] = [{ id: 'block-1', type: 'code', sortingKey: 'a0', content: 'print(1)' }]; + + const outputs = new Map(); + + outputs.set('block-1', [{ output_type: 'stream', text: 'new' }]); + + service.mergeOutputsIntoBlocks(blocks, outputs); + + assert.isUndefined(blocks[0].outputs); + }); + + test('should preserve blocks without matching outputs', () => { + const blocks: DeepnoteBlock[] = [ + { + id: 'block-1', + type: 'code', + sortingKey: 'a0', + content: 'print(1)', + outputs: [{ output_type: 'stream', text: 'old' }] + } + ]; + + const outputs = new Map(); + + outputs.set('block-2', [{ output_type: 'stream', text: 'new' }]); + + const result = service.mergeOutputsIntoBlocks(blocks, outputs); + + assert.deepEqual(result[0].outputs, [{ output_type: 'stream', text: 'old' }]); + }); + + test('should handle empty outputs map', () => { + const blocks: DeepnoteBlock[] = [{ id: 'block-1', type: 'code', sortingKey: 'a0', content: 'print(1)' }]; + + const outputs = new Map(); + + const result = service.mergeOutputsIntoBlocks(blocks, outputs); + + assert.lengthOf(result, 1); + assert.isUndefined(result[0].outputs); + }); + + test('should handle empty blocks array', () => { + const blocks: DeepnoteBlock[] = []; + const outputs = new Map(); + + outputs.set('block-1', [{ output_type: 'stream', text: 'test' }]); + + const result = service.mergeOutputsIntoBlocks(blocks, outputs); + + assert.lengthOf(result, 0); + }); + }); + + suite('stripOutputsFromBlocks', () => { + test('should remove outputs from all blocks', () => { + const blocks: DeepnoteBlock[] = [ + { + id: 'block-1', + type: 'code', + sortingKey: 'a0', + content: 'print(1)', + outputs: [{ output_type: 'stream', text: '1' }] + }, + { + id: 'block-2', + type: 'code', + sortingKey: 'a1', + content: 'print(2)', + outputs: [{ output_type: 'stream', text: '2' }] + } + ]; + + const result = service.stripOutputsFromBlocks(blocks); + + assert.lengthOf(result, 2); + assert.isUndefined(result[0].outputs); + assert.isUndefined(result[1].outputs); + }); + + test('should preserve other block properties', () => { + const blocks: DeepnoteBlock[] = [ + { + id: 'block-1', + type: 'code', + sortingKey: 'a0', + content: 'print(1)', + contentHash: 'sha256:abc123', + executionStartedAt: '2025-01-01T00:00:00Z', + outputs: [{ output_type: 'stream', text: '1' }] + } + ]; + + const result = service.stripOutputsFromBlocks(blocks); + + assert.strictEqual(result[0].id, 'block-1'); + assert.strictEqual(result[0].type, 'code'); + assert.strictEqual(result[0].content, 'print(1)'); + assert.strictEqual(result[0].contentHash, 'sha256:abc123'); + assert.strictEqual(result[0].executionStartedAt, '2025-01-01T00:00:00Z'); + assert.isUndefined(result[0].outputs); + }); + + test('should not modify original blocks', () => { + const blocks: DeepnoteBlock[] = [ + { + id: 'block-1', + type: 'code', + sortingKey: 'a0', + content: 'print(1)', + outputs: [{ output_type: 'stream', text: '1' }] + } + ]; + + service.stripOutputsFromBlocks(blocks); + + assert.isDefined(blocks[0].outputs); + }); + + test('should handle blocks without outputs', () => { + const blocks: DeepnoteBlock[] = [{ id: 'block-1', type: 'code', sortingKey: 'a0', content: 'print(1)' }]; + + const result = service.stripOutputsFromBlocks(blocks); + + assert.lengthOf(result, 1); + assert.isUndefined(result[0].outputs); + }); + + test('should handle empty array', () => { + const blocks: DeepnoteBlock[] = []; + + const result = service.stripOutputsFromBlocks(blocks); + + assert.lengthOf(result, 0); + }); + }); + + suite('extractOutputsFromBlocks', () => { + test('should extract outputs into a map', () => { + const blocks: DeepnoteBlock[] = [ + { + id: 'block-1', + type: 'code', + sortingKey: 'a0', + content: 'print(1)', + outputs: [{ output_type: 'stream', text: '1' }] + }, + { + id: 'block-2', + type: 'code', + sortingKey: 'a1', + content: 'print(2)', + outputs: [{ output_type: 'stream', text: '2' }] + } + ]; + + const result = service.extractOutputsFromBlocks(blocks); + + assert.strictEqual(result.size, 2); + assert.deepEqual(result.get('block-1'), [{ output_type: 'stream', text: '1' }]); + assert.deepEqual(result.get('block-2'), [{ output_type: 'stream', text: '2' }]); + }); + + test('should skip blocks without outputs', () => { + const blocks: DeepnoteBlock[] = [ + { + id: 'block-1', + type: 'code', + sortingKey: 'a0', + content: 'print(1)', + outputs: [{ output_type: 'stream', text: '1' }] + }, + { id: 'block-2', type: 'code', sortingKey: 'a1', content: 'print(2)' } + ]; + + const result = service.extractOutputsFromBlocks(blocks); + + assert.strictEqual(result.size, 1); + assert.isTrue(result.has('block-1')); + assert.isFalse(result.has('block-2')); + }); + + test('should skip blocks without ID', () => { + const blocks = [ + { type: 'code', sortingKey: 'a0', content: 'print(1)', outputs: [{ output_type: 'stream', text: '1' }] } + ] as unknown as DeepnoteBlock[]; + + const result = service.extractOutputsFromBlocks(blocks); + + assert.strictEqual(result.size, 0); + }); + + test('should handle empty array', () => { + const blocks: DeepnoteBlock[] = []; + + const result = service.extractOutputsFromBlocks(blocks); + + assert.strictEqual(result.size, 0); + }); + + test('should handle complex outputs', () => { + const complexOutput: DeepnoteOutput = { + output_type: 'execute_result', + execution_count: 1, + data: { + 'text/html': '...
', + 'text/plain': 'DataFrame...' + }, + metadata: { table_state_spec: '{}' } + }; + + const blocks: DeepnoteBlock[] = [ + { id: 'block-1', type: 'code', sortingKey: 'a0', content: 'df', outputs: [complexOutput] } + ]; + + const result = service.extractOutputsFromBlocks(blocks); + + assert.deepEqual(result.get('block-1'), [complexOutput]); + }); + }); + + suite('isSnapshotsEnabled', () => { + test('should return true when snapshots.enabled is true', () => { + const mockConfig = mock(); + when(mockConfig.get('snapshots.enabled', false)).thenReturn(true); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + + const result = service.isSnapshotsEnabled(); + + assert.isTrue(result); + }); + + test('should return false when snapshots.enabled is false', () => { + const mockConfig = mock(); + when(mockConfig.get('snapshots.enabled', false)).thenReturn(false); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + + const result = service.isSnapshotsEnabled(); + + assert.isFalse(result); + }); + }); + + suite('readSnapshot', () => { + const projectId = 'test-project-id-123'; + + test('should return undefined when no workspace folders exist', async () => { + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn(undefined); + + const result = await service.readSnapshot(projectId); + + assert.isUndefined(result); + }); + + test('should return undefined when workspace folders array is empty', async () => { + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([]); + + const result = await service.readSnapshot(projectId); + + assert.isUndefined(result); + }); + + test('should find and parse latest snapshot file', async () => { + const workspaceFolder: WorkspaceFolder = { + uri: Uri.file('/workspace'), + name: 'workspace', + index: 0 + }; + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder]); + + const snapshotUri = Uri.file('/workspace/snapshots/project_test-project-id-123_latest.snapshot.deepnote'); + when(mockedVSCodeNamespaces.workspace.findFiles(anything(), anything(), anything())).thenResolve([ + snapshotUri + ] as any); + + const snapshotYaml = ` +version: '1.0' +project: + id: test-project-id-123 + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + type: code + content: print(1) + outputs: + - output_type: stream + name: stdout + text: '1' + - id: block-2 + type: markdown + content: '# Hello' +`; + const mockFs = mock(); + when(mockFs.readFile(anything())).thenResolve(new TextEncoder().encode(snapshotYaml) as any); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + const result = await service.readSnapshot(projectId); + + assert.isDefined(result); + assert.strictEqual(result!.size, 1); + assert.deepEqual(result!.get('block-1'), [{ output_type: 'stream', name: 'stdout', text: '1' }]); + }); + + test('should return undefined when no snapshot files found', async () => { + const workspaceFolder: WorkspaceFolder = { + uri: Uri.file('/workspace'), + name: 'workspace', + index: 0 + }; + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder]); + when(mockedVSCodeNamespaces.workspace.findFiles(anything(), anything(), anything())).thenResolve([] as any); + + const result = await service.readSnapshot(projectId); + + assert.isUndefined(result); + }); + + test('should fall back to most recent timestamped snapshot when no latest exists', async () => { + const workspaceFolder: WorkspaceFolder = { + uri: Uri.file('/workspace'), + name: 'workspace', + index: 0 + }; + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder]); + + // First call for latest - returns empty + // Second call for timestamped - returns files + const timestampedUri1 = Uri.file( + '/workspace/snapshots/project_test-project-id-123_2025-01-01T10-00-00.snapshot.deepnote' + ); + const timestampedUri2 = Uri.file( + '/workspace/snapshots/project_test-project-id-123_2025-01-02T10-00-00.snapshot.deepnote' + ); + + let callCount = 0; + when(mockedVSCodeNamespaces.workspace.findFiles(anything(), anything(), anything())).thenCall(() => { + callCount++; + if (callCount === 1) { + return Promise.resolve([]); + } + + return Promise.resolve([timestampedUri1, timestampedUri2]); + }); + + const snapshotYaml = ` +version: '1.0' +project: + notebooks: + - blocks: + - id: block-1 + outputs: + - output_type: stream + text: 'from timestamped' +`; + const mockFs = mock(); + when(mockFs.readFile(anything())).thenResolve(new TextEncoder().encode(snapshotYaml) as any); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + const result = await service.readSnapshot(projectId); + + assert.isDefined(result); + assert.strictEqual(result!.size, 1); + }); + + test('should return empty map when snapshot file read fails', async () => { + const workspaceFolder: WorkspaceFolder = { + uri: Uri.file('/workspace'), + name: 'workspace', + index: 0 + }; + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder]); + + const snapshotUri = Uri.file('/workspace/snapshots/project_test-project-id-123_latest.snapshot.deepnote'); + when(mockedVSCodeNamespaces.workspace.findFiles(anything(), anything(), anything())).thenResolve([ + snapshotUri + ] as any); + + const mockFs = mock(); + when(mockFs.readFile(anything())).thenReject(new Error('File read error')); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + const result = await service.readSnapshot(projectId); + + // parseSnapshotFile catches read errors and returns empty map + assert.isDefined(result); + assert.strictEqual(result!.size, 0); + }); + + test('should return empty map when snapshot has invalid structure', async () => { + const workspaceFolder: WorkspaceFolder = { + uri: Uri.file('/workspace'), + name: 'workspace', + index: 0 + }; + when(mockedVSCodeNamespaces.workspace.workspaceFolders).thenReturn([workspaceFolder]); + + const snapshotUri = Uri.file('/workspace/snapshots/project_test-project-id-123_latest.snapshot.deepnote'); + when(mockedVSCodeNamespaces.workspace.findFiles(anything(), anything(), anything())).thenResolve([ + snapshotUri + ] as any); + + const invalidYaml = 'not_an_object'; + const mockFs = mock(); + when(mockFs.readFile(anything())).thenResolve(new TextEncoder().encode(invalidYaml) as any); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + const result = await service.readSnapshot(projectId); + + assert.isDefined(result); + assert.strictEqual(result!.size, 0); + }); + }); + + suite('createSnapshot', () => { + const projectUri = Uri.file('/workspace/my-project.deepnote'); + const projectId = 'test-project-id-123'; + const projectName = 'My Project'; + + test('should create snapshot files when there are changes', async () => { + const projectData = createProjectData(); + + const mockFs = mock(); + // Directory doesn't exist - stat throws + when(mockFs.stat(anything())).thenReject(new Error('ENOENT')); + when(mockFs.createDirectory(anything())).thenResolve(); + when(mockFs.readFile(anything())).thenReject(new Error('ENOENT')); + when(mockFs.writeFile(anything(), anything())).thenResolve(); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + const result = await service.createSnapshot(projectUri, projectId, projectName, projectData); + + assert.isDefined(result); + assert.include(result!.fsPath, 'snapshot.deepnote'); + }); + + test('should return undefined when project name is invalid', async () => { + const projectData = createProjectData(); + + const result = await service.createSnapshot(projectUri, projectId, '', projectData); + + assert.isUndefined(result); + }); + + test('should return undefined when directory creation fails', async () => { + const projectData = createProjectData(); + + const mockFs = mock(); + when(mockFs.stat(anything())).thenReject(new Error('ENOENT')); + when(mockFs.createDirectory(anything())).thenReject(new Error('Permission denied')); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + const result = await service.createSnapshot(projectUri, projectId, projectName, projectData); + + assert.isUndefined(result); + }); + + test('should skip snapshot creation when no changes detected', async () => { + const projectData = createProjectData(); + + const mockFs = mock(); + // Directory exists + when(mockFs.stat(anything())).thenResolve({ type: FileType.Directory } as any); + // Return same content as existing + const existingYaml = ` +metadata: + createdAt: '2025-01-01T00:00:00Z' +version: '1.0' +project: + id: test-project-id-123 + name: My Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + type: code + sortingKey: a0 + content: print(1) + outputs: + - output_type: stream + text: '1' +`; + when(mockFs.readFile(anything())).thenResolve(new TextEncoder().encode(existingYaml) as any); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + const result = await service.createSnapshot(projectUri, projectId, projectName, projectData); + + assert.isUndefined(result); + }); + + test('should return undefined when timestamped file write fails', async () => { + const projectData = createProjectData(); + + const mockFs = mock(); + when(mockFs.stat(anything())).thenResolve({ type: FileType.Directory } as any); + when(mockFs.readFile(anything())).thenReject(new Error('ENOENT')); + when(mockFs.writeFile(anything(), anything())).thenReject(new Error('Write failed')); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + const result = await service.createSnapshot(projectUri, projectId, projectName, projectData); + + assert.isUndefined(result); + }); + + test('should return timestamped path even if latest write fails', async () => { + const projectData = createProjectData(); + + const mockFs = mock(); + when(mockFs.stat(anything())).thenResolve({ type: FileType.Directory } as any); + when(mockFs.readFile(anything())).thenReject(new Error('ENOENT')); + + let writeCallCount = 0; + when(mockFs.writeFile(anything(), anything())).thenCall(() => { + writeCallCount++; + if (writeCallCount === 1) { + // First write (timestamped) succeeds + return Promise.resolve(); + } + // Second write (latest) fails + return Promise.reject(new Error('Write failed')); + }); + when(mockedVSCodeNamespaces.workspace.fs).thenReturn(instance(mockFs)); + + const result = await service.createSnapshot(projectUri, projectId, projectName, projectData); + + assert.isDefined(result); + assert.include(result!.fsPath, 'snapshot.deepnote'); + assert.notInclude(result!.fsPath, 'latest'); + }); + }); + + // Metadata tracking tests (now using serviceAny for private methods) + suite('execution metadata tracking', () => { + const notebookUri = 'file:///path/to/notebook.deepnote'; + const cellId = 'cell-123'; + + suite('recordCellExecutionStart (private)', () => { + test('should record cell execution start time', () => { + const startTime = Date.now(); + + serviceAny.recordCellExecutionStart(notebookUri, cellId, startTime); + + const metadata = service.getBlockExecutionMetadata(notebookUri, cellId); + assert.isDefined(metadata); + assert.isDefined(metadata!.executionStartedAt); + assert.isUndefined(metadata!.executionFinishedAt); + }); + + test('should initialize notebook execution state', () => { + const startTime = Date.now(); + + serviceAny.recordCellExecutionStart(notebookUri, cellId, startTime); + + const executionMetadata = service.getExecutionMetadata(notebookUri); + // Should not have execution metadata yet since no cells have completed + assert.isUndefined(executionMetadata); + }); + + test('should handle multiple cells in same notebook', () => { + const startTime = Date.now(); + + serviceAny.recordCellExecutionStart(notebookUri, 'cell-1', startTime); + serviceAny.recordCellExecutionStart(notebookUri, 'cell-2', startTime + 1000); + + const metadata1 = service.getBlockExecutionMetadata(notebookUri, 'cell-1'); + const metadata2 = service.getBlockExecutionMetadata(notebookUri, 'cell-2'); + + assert.isDefined(metadata1); + assert.isDefined(metadata2); + assert.notStrictEqual(metadata1!.executionStartedAt, metadata2!.executionStartedAt); + }); + }); + + suite('recordCellExecutionEnd (private)', () => { + test('should record successful cell execution end', () => { + const startTime = Date.now(); + const endTime = startTime + 1000; + + serviceAny.recordCellExecutionStart(notebookUri, cellId, startTime); + serviceAny.recordCellExecutionEnd(notebookUri, cellId, endTime, true); + + const metadata = service.getBlockExecutionMetadata(notebookUri, cellId); + assert.isDefined(metadata); + assert.isDefined(metadata!.executionStartedAt); + assert.isDefined(metadata!.executionFinishedAt); + }); + + test('should update execution summary on success', () => { + const startTime = Date.now(); + const endTime = startTime + 1000; + + serviceAny.recordCellExecutionStart(notebookUri, cellId, startTime); + serviceAny.recordCellExecutionEnd(notebookUri, cellId, endTime, true); + + const executionMetadata = service.getExecutionMetadata(notebookUri); + assert.isDefined(executionMetadata); + assert.isDefined(executionMetadata!.summary); + assert.strictEqual(executionMetadata!.summary!.blocksExecuted, 1); + assert.strictEqual(executionMetadata!.summary!.blocksSucceeded, 1); + assert.strictEqual(executionMetadata!.summary!.blocksFailed, 0); + }); + + test('should update execution summary on failure', () => { + const startTime = Date.now(); + const endTime = startTime + 1000; + + serviceAny.recordCellExecutionStart(notebookUri, cellId, startTime); + serviceAny.recordCellExecutionEnd(notebookUri, cellId, endTime, false); + + const executionMetadata = service.getExecutionMetadata(notebookUri); + assert.isDefined(executionMetadata); + assert.isDefined(executionMetadata!.summary); + assert.strictEqual(executionMetadata!.summary!.blocksExecuted, 1); + assert.strictEqual(executionMetadata!.summary!.blocksSucceeded, 0); + assert.strictEqual(executionMetadata!.summary!.blocksFailed, 1); + }); + + test('should record error details on failure', () => { + const startTime = Date.now(); + const endTime = startTime + 1000; + const error = { name: 'TypeError', message: 'undefined is not a function' }; + + serviceAny.recordCellExecutionStart(notebookUri, cellId, startTime); + serviceAny.recordCellExecutionEnd(notebookUri, cellId, endTime, false, error); + + const executionMetadata = service.getExecutionMetadata(notebookUri); + assert.isDefined(executionMetadata); + assert.isDefined(executionMetadata!.error); + assert.strictEqual(executionMetadata!.error!.name, 'TypeError'); + assert.strictEqual(executionMetadata!.error!.message, 'undefined is not a function'); + }); + + test('should accumulate multiple cell executions', () => { + const startTime = Date.now(); + + // Execute 3 cells: 2 successful, 1 failed + serviceAny.recordCellExecutionStart(notebookUri, 'cell-1', startTime); + serviceAny.recordCellExecutionEnd(notebookUri, 'cell-1', startTime + 100, true); + + serviceAny.recordCellExecutionStart(notebookUri, 'cell-2', startTime + 200); + serviceAny.recordCellExecutionEnd(notebookUri, 'cell-2', startTime + 300, true); + + serviceAny.recordCellExecutionStart(notebookUri, 'cell-3', startTime + 400); + serviceAny.recordCellExecutionEnd(notebookUri, 'cell-3', startTime + 500, false); + + const executionMetadata = service.getExecutionMetadata(notebookUri); + assert.isDefined(executionMetadata); + assert.isDefined(executionMetadata!.summary); + assert.strictEqual(executionMetadata!.summary!.blocksExecuted, 3); + assert.strictEqual(executionMetadata!.summary!.blocksSucceeded, 2); + assert.strictEqual(executionMetadata!.summary!.blocksFailed, 1); + }); + + test('should calculate total duration', () => { + const startTime = Date.now(); + const endTime = startTime + 5000; + + serviceAny.recordCellExecutionStart(notebookUri, cellId, startTime); + serviceAny.recordCellExecutionEnd(notebookUri, cellId, endTime, true); + + const executionMetadata = service.getExecutionMetadata(notebookUri); + assert.isDefined(executionMetadata); + assert.isDefined(executionMetadata!.summary); + assert.strictEqual(executionMetadata!.summary!.totalDurationMs, 5000); + }); + }); + + suite('getExecutionMetadata', () => { + test('should return undefined for unknown notebook', () => { + const metadata = service.getExecutionMetadata('unknown-notebook'); + assert.isUndefined(metadata); + }); + + test('should return undefined if no cells have been executed', () => { + const startTime = Date.now(); + serviceAny.recordCellExecutionStart(notebookUri, cellId, startTime); + + const metadata = service.getExecutionMetadata(notebookUri); + assert.isUndefined(metadata); + }); + + test('should include ISO timestamps', () => { + const startTime = Date.now(); + serviceAny.recordCellExecutionStart(notebookUri, cellId, startTime); + serviceAny.recordCellExecutionEnd(notebookUri, cellId, startTime + 1000, true); + + const metadata = service.getExecutionMetadata(notebookUri); + assert.isDefined(metadata); + assert.isDefined(metadata!.startedAt); + assert.isDefined(metadata!.finishedAt); + // Should be valid ISO date strings + assert.doesNotThrow(() => new Date(metadata!.startedAt!)); + assert.doesNotThrow(() => new Date(metadata!.finishedAt!)); + }); + }); + + suite('getBlockExecutionMetadata', () => { + test('should return undefined for unknown notebook', () => { + const metadata = service.getBlockExecutionMetadata('unknown-notebook', cellId); + assert.isUndefined(metadata); + }); + + test('should return undefined for unknown cell', () => { + const startTime = Date.now(); + serviceAny.recordCellExecutionStart(notebookUri, cellId, startTime); + + const metadata = service.getBlockExecutionMetadata(notebookUri, 'unknown-cell'); + assert.isUndefined(metadata); + }); + }); + + suite('clearExecutionState', () => { + test('should clear all state for a notebook', () => { + const startTime = Date.now(); + serviceAny.recordCellExecutionStart(notebookUri, cellId, startTime); + serviceAny.recordCellExecutionEnd(notebookUri, cellId, startTime + 1000, true); + + service.clearExecutionState(notebookUri); + + const executionMetadata = service.getExecutionMetadata(notebookUri); + const blockMetadata = service.getBlockExecutionMetadata(notebookUri, cellId); + + assert.isUndefined(executionMetadata); + assert.isUndefined(blockMetadata); + }); + + test('should only clear state for specified notebook', () => { + const startTime = Date.now(); + const otherNotebookUri = 'file:///other/notebook.deepnote'; + + serviceAny.recordCellExecutionStart(notebookUri, cellId, startTime); + serviceAny.recordCellExecutionEnd(notebookUri, cellId, startTime + 1000, true); + + serviceAny.recordCellExecutionStart(otherNotebookUri, 'other-cell', startTime); + serviceAny.recordCellExecutionEnd(otherNotebookUri, 'other-cell', startTime + 1000, true); + + service.clearExecutionState(notebookUri); + + // First notebook should be cleared + assert.isUndefined(service.getExecutionMetadata(notebookUri)); + + // Second notebook should still have state + assert.isDefined(service.getExecutionMetadata(otherNotebookUri)); + }); + }); + + suite('multiple notebooks', () => { + test('should track state independently for different notebooks', () => { + const notebook1 = 'file:///notebook1.deepnote'; + const notebook2 = 'file:///notebook2.deepnote'; + const startTime = Date.now(); + + // Execute cells in different notebooks + serviceAny.recordCellExecutionStart(notebook1, 'cell-1', startTime); + serviceAny.recordCellExecutionEnd(notebook1, 'cell-1', startTime + 100, true); + + serviceAny.recordCellExecutionStart(notebook2, 'cell-2', startTime); + serviceAny.recordCellExecutionEnd(notebook2, 'cell-2', startTime + 200, false); + + const metadata1 = service.getExecutionMetadata(notebook1); + const metadata2 = service.getExecutionMetadata(notebook2); + + assert.isDefined(metadata1); + assert.isDefined(metadata1!.summary); + assert.strictEqual(metadata1!.summary!.blocksSucceeded, 1); + assert.strictEqual(metadata1!.summary!.blocksFailed, 0); + + assert.isDefined(metadata2); + assert.isDefined(metadata2!.summary); + assert.strictEqual(metadata2!.summary!.blocksSucceeded, 0); + assert.strictEqual(metadata2!.summary!.blocksFailed, 1); + }); + }); + + suite('setRunAllMode / isRunAllMode', () => { + test('should return false by default', () => { + assert.isFalse(service.isRunAllMode(notebookUri)); + }); + + test('should return true after setting run all mode', () => { + service.setRunAllMode(notebookUri); + + assert.isTrue(service.isRunAllMode(notebookUri)); + }); + + test('should track run all mode per notebook', () => { + const notebook1 = 'file:///notebook1.deepnote'; + const notebook2 = 'file:///notebook2.deepnote'; + + service.setRunAllMode(notebook1); + + assert.isTrue(service.isRunAllMode(notebook1)); + assert.isFalse(service.isRunAllMode(notebook2)); + }); + }); + + suite('captureEnvironmentBeforeExecution', () => { + test('should not throw for valid notebook URI', async () => { + await service.captureEnvironmentBeforeExecution(notebookUri); + // Should complete without error + }); + }); + + suite('getEnvironmentMetadata', () => { + test('should return undefined when no environment captured', async () => { + const result = await service.getEnvironmentMetadata(notebookUri); + + assert.isUndefined(result); + }); + }); + }); +}); diff --git a/src/notebooks/notebookCommandListener.ts b/src/notebooks/notebookCommandListener.ts index 04907cb1ca..61588839a4 100644 --- a/src/notebooks/notebookCommandListener.ts +++ b/src/notebooks/notebookCommandListener.ts @@ -34,6 +34,7 @@ import { KernelConnector } from './controllers/kernelConnector'; import { IControllerRegistration } from './controllers/types'; import { IExtensionSyncActivationService } from '../platform/activation/types'; import { IKernelStatusProvider } from '../kernels/kernelStatusProvider'; +import { ISnapshotMetadataService } from './deepnote/snapshots/snapshotService'; export const INotebookCommandHandler = Symbol('INotebookCommandHandler'); export interface INotebookCommandHandler { @@ -54,7 +55,8 @@ export class NotebookCommandListener implements INotebookCommandHandler, IExtens @inject(IDataScienceErrorHandler) private errorHandler: IDataScienceErrorHandler, @inject(INotebookEditorProvider) private notebookEditorProvider: INotebookEditorProvider, @inject(IServiceContainer) private serviceContainer: IServiceContainer, - @inject(IKernelStatusProvider) private kernelStatusProvider: IKernelStatusProvider + @inject(IKernelStatusProvider) private kernelStatusProvider: IKernelStatusProvider, + @inject(ISnapshotMetadataService) private snapshotMetadataService: ISnapshotMetadataService ) {} activate(): void { @@ -114,6 +116,9 @@ export class NotebookCommandListener implements INotebookCommandHandler, IExtens private runAllCells() { if (window.activeNotebookEditor) { + const notebookUri = window.activeNotebookEditor.notebook.uri.toString(); + + this.snapshotMetadataService.setRunAllMode(notebookUri); commands.executeCommand('notebook.execute').then(noop, noop); } } diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index f7c8fd6bf6..53ac77c504 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -88,8 +88,8 @@ import { DeepnoteCellCopyHandler } from './deepnote/deepnoteCellCopyHandler'; import { DeepnoteEnvironmentTreeDataProvider } from '../kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.node'; import { OpenInDeepnoteHandler } from './deepnote/openInDeepnoteHandler.node'; import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler'; -import { ISnapshotMetadataService, SnapshotMetadataService } from './deepnote/snapshotMetadataService'; -import { EnvironmentCapture, IEnvironmentCapture } from './deepnote/environmentCapture.node'; +import { ISnapshotMetadataService, SnapshotService } from './deepnote/snapshots/snapshotService'; +import { EnvironmentCapture, IEnvironmentCapture } from './deepnote/snapshots/environmentCapture.node'; export function registerTypes(serviceManager: IServiceManager, isDevMode: boolean) { registerControllerTypes(serviceManager, isDevMode); @@ -242,10 +242,11 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea DeepnoteNotebookEnvironmentMapper ); - // Snapshot metadata services + // Snapshot service serviceManager.addSingleton(IEnvironmentCapture, EnvironmentCapture); - serviceManager.addSingleton(ISnapshotMetadataService, SnapshotMetadataService); - serviceManager.addBinding(ISnapshotMetadataService, IExtensionSyncActivationService); + serviceManager.addSingleton(SnapshotService, SnapshotService); + serviceManager.addBinding(SnapshotService, IExtensionSyncActivationService); + serviceManager.addBinding(SnapshotService, ISnapshotMetadataService); // File export/import serviceManager.addSingleton(IFileConverter, FileConverter); diff --git a/src/platform/common/constants.ts b/src/platform/common/constants.ts index f046de6187..c85719faec 100644 --- a/src/platform/common/constants.ts +++ b/src/platform/common/constants.ts @@ -223,6 +223,8 @@ export namespace Commands { export const OpenDeepnoteNotebook = 'deepnote.openNotebook'; export const OpenDeepnoteFile = 'deepnote.openFile'; export const RevealInDeepnoteExplorer = 'deepnote.revealInExplorer'; + export const EnableSnapshots = 'deepnote.enableSnapshots'; + export const DisableSnapshots = 'deepnote.disableSnapshots'; export const ManageIntegrations = 'deepnote.manageIntegrations'; export const AddSqlBlock = 'deepnote.addSqlBlock'; export const AddBigNumberChartBlock = 'deepnote.addBigNumberChartBlock'; diff --git a/src/platform/errors/invalidProjectNameError.ts b/src/platform/errors/invalidProjectNameError.ts new file mode 100644 index 0000000000..8126345738 --- /dev/null +++ b/src/platform/errors/invalidProjectNameError.ts @@ -0,0 +1,22 @@ +import { l10n } from 'vscode'; + +import { BaseError } from './types'; + +/** + * Error thrown when a project name is invalid for slug generation. + * This occurs when the project name is empty or contains only special characters + * that are removed during slugification. + * + * Cause: + * The project name provided cannot be converted to a valid slug for use in filenames. + * + * Handled by: + * The error should be caught and logged. The snapshot operation should be skipped + * if the project name cannot be slugified. + */ +export class InvalidProjectNameError extends BaseError { + constructor() { + super('unknown', l10n.t('Project name cannot be empty or contain only special characters')); + this.name = 'InvalidProjectNameError'; + } +} diff --git a/src/platform/notebooks/cellExecutionStateService.ts b/src/platform/notebooks/cellExecutionStateService.ts index f8289b2bc1..25b25dd376 100644 --- a/src/platform/notebooks/cellExecutionStateService.ts +++ b/src/platform/notebooks/cellExecutionStateService.ts @@ -38,6 +38,16 @@ export interface NotebookCellExecutionStateChangeEvent { readonly state: NotebookCellExecutionState; } +/** + * An event describing completion of a notebook's cell execution queue. + */ +export interface NotebookQueueCompletionEvent { + /** + * The URI of the notebook whose execution queue has completed. + */ + readonly notebookUri: string; +} + const STATE_NAMES: Record = { [NotebookCellExecutionState.Idle]: 'Idle', [NotebookCellExecutionState.Pending]: 'Pending', @@ -46,6 +56,7 @@ const STATE_NAMES: Record = { export namespace notebookCellExecutions { const eventEmitter = trackDisposable(new EventEmitter()); + const queueCompletionEmitter = trackDisposable(new EventEmitter()); /** * An {@link Event} which fires when the execution state of a cell has changed. @@ -54,6 +65,21 @@ export namespace notebookCellExecutions { // how a correct consumer works, e.g the consumer could have been late and missed an event? export const onDidChangeNotebookCellExecutionState = eventEmitter.event; + /** + * An {@link Event} which fires when a notebook's cell execution queue has completed. + * This is fired after all queued cells have finished executing (success, error, or cancel). + */ + export const onDidCompleteQueueExecution = queueCompletionEmitter.event; + + /** + * Notify listeners that a notebook's cell execution queue has completed. + * @param notebookUri The URI of the notebook whose queue completed + */ + export function notifyQueueComplete(notebookUri: string) { + logger.debug(`[CellExecState] Queue execution complete for ${notebookUri}`); + queueCompletionEmitter.fire({ notebookUri }); + } + export function changeCellState(cell: NotebookCell, state: NotebookCellExecutionState, executionOrder?: number) { const cellId = cell.metadata?.id as string | undefined; const stateName = STATE_NAMES[state] || String(state); diff --git a/src/platform/notebooks/deepnote/types.ts b/src/platform/notebooks/deepnote/types.ts index 3f2a20274c..6ae06bc731 100644 --- a/src/platform/notebooks/deepnote/types.ts +++ b/src/platform/notebooks/deepnote/types.ts @@ -117,22 +117,3 @@ export const IPlatformDeepnoteNotebookManager = Symbol('IPlatformDeepnoteNoteboo export interface IPlatformDeepnoteNotebookManager { getOriginalProject(projectId: string): DeepnoteProject | undefined; } - -/** - * Platform-layer interface for snapshot metadata service. - * Used by the kernel execution layer to capture environment before cell execution. - */ -export const ISnapshotMetadataService = Symbol('ISnapshotMetadataService'); -export interface ISnapshotMetadataService { - /** - * Capture environment before execution starts. - * Called at the start of a cell execution batch. - * This is blocking and should complete before cells execute. - */ - captureEnvironmentBeforeExecution(notebookUri: string): Promise; - - /** - * Clear execution state for a notebook (e.g., when kernel restarts). - */ - clearExecutionState(notebookUri: string): void; -}