feat: Track the environment information when executing blocks (#254)

Author: Artmann

build/esbuild/build.ts

@@ -59,10 +59,12 @@ const commonExternals = [
     '@opentelemetry/instrumentation',
     '@azure/functions-core',
     // Node.js builtins (with node: prefix)
+    'node:child_process',
     'node:crypto',
     'node:fs/promises',
     'node:os',
     'node:path',
+    'node:util',
     'ansi-regex' // Used by regexp utils
 ];
 // Create separate copies to avoid shared-state mutations

cspell.json

@@ -23,12 +23,18 @@
     "words": [
         "alloydb",
         "altdb",
+        "altgraph",
         "altuser",
+        "anyio",
+        "artmann",
         "blockgroup",
+        "boto",
         "channeldef",
+        "cpython",
         "clickhouse",
         "dataframe",
         "datascience",
+        "dateutil",
         "dbname",
         "deepnote",
         "deepnoteserver",
@@ -53,6 +59,8 @@
         "matplotlib",
         "millis",
         "mindsdb",
+        "msgspec",
+        "mypackage",
         "mydb",
         "nbformat",
         "nbinsx",
@@ -61,12 +69,16 @@
         "pids",
         "Pids",
         "plotly",
+        "pyenv",
         "pylsp",
         "PYTHONHOME",
+        "pyyaml",
         "Reselecting",
         "rootpass",
+        "scikit",
         "scipy",
         "sklearn",
+        "sqlalchemy",
         "taskkill",
         "testdb",
         "toolsai",

package-lock.json

@@ -2500,7 +2500,7 @@
     },
     "dependencies": {
         "@c4312/evt": "^0.1.1",
-        "@deepnote/blocks": "^1.3.5",
+        "@deepnote/blocks": "^1.4.0",
         "@deepnote/convert": "^1.3.0",
         "@deepnote/database-integrations": "^1.3.0",
         "@deepnote/sql-language-server": "^3.0.0",

package.json

@@ -1,7 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { CancellationToken, Disposable, EventEmitter, NotebookCell } from 'vscode';
+import { CancellationToken, Disposable, EventEmitter, NotebookCell, NotebookDocument } from 'vscode';
+
 import { logger } from '../../platform/logging';
 import { noop } from '../../platform/common/utils/misc';
 import { traceCellMessage } from './helpers';
@@ -13,6 +14,7 @@ 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';
 
 /**
  * A queue responsible for execution of cells.
@@ -48,7 +50,9 @@ export class CellExecutionQueue implements Disposable {
         private readonly session: Promise<IKernelSession>,
         private readonly executionFactory: CellExecutionFactory,
         readonly metadata: Readonly<KernelConnectionMetadata>,
-        readonly resourceUri: Resource
+        readonly resourceUri: Resource,
+        private readonly notebook?: NotebookDocument,
+        private readonly snapshotService?: ISnapshotMetadataService
     ) {}
 
     public dispose() {
@@ -189,6 +193,18 @@ export class CellExecutionQueue implements Disposable {
     private async executeQueuedCells() {
         let notebookClosed: boolean | undefined;
         const kernelConnection = await this.session;
+
+        // Capture environment before executing any cells
+        if (this.snapshotService && this.notebook) {
+            const notebookUri = this.notebook.uri.toString();
+
+            try {
+                await this.snapshotService.captureEnvironmentBeforeExecution(notebookUri);
+            } catch (error) {
+                logger.warn('[CellExecutionQueue] Failed to capture environment before execution', error);
+            }
+        }
+
         this.queueOfItemsToExecute.forEach(
             (exec) => exec.type === 'cell' && traceCellMessage(exec.cell, 'Ready to execute')
         );

src/kernels/execution/cellExecutionQueue.ts

@@ -45,6 +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';
 
 /**
  * Everything in this classes gets disposed via the `onWillCancel` hook.
@@ -65,7 +66,8 @@ export class NotebookKernelExecution implements INotebookKernelExecution {
         private readonly kernel: IKernel,
         context: IExtensionContext,
         formatters: ITracebackFormatter[],
-        private readonly notebook: NotebookDocument
+        private readonly notebook: NotebookDocument,
+        private readonly snapshotService?: ISnapshotMetadataService
     ) {
         const requestListener = new CellExecutionMessageHandlerService(
             kernel.controller,
@@ -361,7 +363,9 @@ export class NotebookKernelExecution implements INotebookKernelExecution {
             sessionPromise,
             this.executionFactory,
             this.kernel.kernelConnectionMetadata,
-            this.kernel.resourceUri
+            this.kernel.resourceUri,
+            document,
+            this.snapshotService
         );
         this.disposables.push(newCellExecutionQueue);
 

src/kernels/kernelExecution.ts

@@ -1,8 +1,9 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
-import { inject, injectable, multiInject, named } from 'inversify';
+import { inject, injectable, multiInject, named, optional } from 'inversify';
 import { Memento, NotebookDocument, Uri } from 'vscode';
+
 import {
     IAsyncDisposableRegistry,
     IConfigurationService,
@@ -31,6 +32,7 @@ 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';
 
 /**
  * Node version of a kernel provider. Needed in order to create the node version of a kernel.
@@ -50,7 +52,8 @@ export class KernelProvider extends BaseCoreKernelProvider {
         @inject(IMemento) @named(WORKSPACE_MEMENTO) private readonly workspaceStorage: Memento,
         @inject(IReplNotebookTrackerService) private readonly replTracker: IReplNotebookTrackerService,
         @inject(IKernelWorkingDirectory) private readonly kernelWorkingDirectory: IKernelWorkingDirectory,
-        @inject(IRawNotebookSupportedService) private readonly rawKernelSupported: IRawNotebookSupportedService
+        @inject(IRawNotebookSupportedService) private readonly rawKernelSupported: IRawNotebookSupportedService,
+        @inject(ISnapshotMetadataService) @optional() private readonly snapshotService?: ISnapshotMetadataService
     ) {
         super(asyncDisposables, disposables);
         disposables.push(jupyterServerUriStorage.onDidRemove(this.handleServerRemoval.bind(this)));
@@ -109,7 +112,10 @@ export class KernelProvider extends BaseCoreKernelProvider {
             this.disposables
         );
 
-        this.executions.set(kernel, new NotebookKernelExecution(kernel, this.context, this.formatters, notebook));
+        this.executions.set(
+            kernel,
+            new NotebookKernelExecution(kernel, this.context, this.formatters, notebook, this.snapshotService)
+        );
         this.asyncDisposables.push(kernel);
         this.storeKernel(notebook, options, kernel);
         this.deleteMappingIfKernelIsDisposed(kernel);

src/kernels/kernelProvider.node.ts

@@ -1,7 +1,6 @@
+import type { DeepnoteBlock } from '@deepnote/blocks';
 import type { NotebookCellData } from 'vscode';
 
-import type { DeepnoteBlock } from '../../../platform/deepnote/deepnoteTypes';
-
 export interface BlockConverter {
     applyChangesToBlock(block: DeepnoteBlock, cell: NotebookCellData): void;
     canConvert(blockType: string): boolean;

src/notebooks/deepnote/converters/blockConverter.ts

@@ -1,7 +1,7 @@
+import type { DeepnoteBlock } from '@deepnote/blocks';
 import { NotebookCellData, NotebookCellKind } from 'vscode';
 
 import type { BlockConverter } from './blockConverter';
-import type { DeepnoteBlock } from '../../../platform/deepnote/deepnoteTypes';
 import { DeepnoteBigNumberMetadataSchema } from '../deepnoteSchemas';
 import { DEEPNOTE_VSCODE_RAW_CONTENT_KEY } from './constants';
 

src/notebooks/deepnote/converters/chartBigNumberBlockConverter.ts

@@ -1,7 +1,6 @@
+import type { DeepnoteBlock } from '@deepnote/blocks';
 import { assert } from 'chai';
 import { NotebookCellData, NotebookCellKind } from 'vscode';
-
-import type { DeepnoteBlock } from '../../../platform/deepnote/deepnoteTypes';
 import { ChartBigNumberBlockConverter } from './chartBigNumberBlockConverter';
 import { DEEPNOTE_VSCODE_RAW_CONTENT_KEY } from './constants';