From b12b77580edd09b418e7ff5ad0b6d3a665738d82 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Filip=20P=C3=BDrek?= Date: Wed, 8 Oct 2025 12:58:13 +0200 Subject: [PATCH 1/4] feat: implement init notebook behavior --- DEEPNOTE_KERNEL_IMPLEMENTATION.md | 419 +++++++++++++++--- README.md | 2 + .../deepnote/deepnoteToolkitInstaller.node.ts | 2 +- src/kernels/deepnote/types.ts | 7 + .../deepnoteInitNotebookRunner.node.ts | 230 ++++++++++ .../deepnoteKernelAutoSelector.node.ts | 108 ++++- .../deepnote/deepnoteNotebookManager.ts | 18 + .../deepnoteRequirementsHelper.node.ts | 45 ++ src/notebooks/serviceRegistry.node.ts | 2 + src/notebooks/types.ts | 2 + 10 files changed, 761 insertions(+), 74 deletions(-) create mode 100644 src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts create mode 100644 src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts diff --git a/DEEPNOTE_KERNEL_IMPLEMENTATION.md b/DEEPNOTE_KERNEL_IMPLEMENTATION.md index ade1164c25..1fc516a52b 100644 --- a/DEEPNOTE_KERNEL_IMPLEMENTATION.md +++ b/DEEPNOTE_KERNEL_IMPLEMENTATION.md @@ -75,14 +75,40 @@ This implementation adds automatic kernel selection and startup for `.deepnote` - `provideJupyterServers(token)`: Lists all registered Deepnote servers - `resolveJupyterServer(server, token)`: Resolves server connection info by handle -#### 5. **Deepnote Kernel Auto-Selector** (`src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts`) +#### 5. **Deepnote Init Notebook Runner** (`src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts`) + +- Runs initialization notebooks automatically before user code executes +- Checks for `project.initNotebookId` in the Deepnote project YAML +- Executes init notebook code blocks sequentially in the kernel +- Shows progress notification to user +- Caches execution per project (runs only once per session) +- Logs errors but allows user to continue on failure +- Blocks user cell execution until init notebook completes + +**Key Methods:** + +- `runInitNotebookIfNeeded(notebook, projectId)`: Main entry point, checks cache and executes if needed +- `executeInitNotebook(notebook, initNotebook)`: Executes all code blocks from init notebook + +#### 6. **Deepnote Requirements Helper** (`src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts`) + +- Static helper class for creating `requirements.txt` from project settings +- Extracts `project.settings.requirements` array from Deepnote YAML +- Creates `requirements.txt` in workspace root before init notebook runs +- Ensures dependencies are available for pip installation in init notebook + +**Key Methods:** + +- `createRequirementsFile(project)`: Creates requirements.txt from project settings + +#### 7. **Deepnote Kernel Auto-Selector** (`src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts`) - Activation service that listens for notebook open events and controller selection changes - Automatically selects Deepnote kernel for `.deepnote` files -- Queries the Deepnote server for available kernel specs -- **Selects a server-native Python kernel spec** (e.g., `python3-venv` or any available Python kernel) +- Queries the Deepnote server for available kernel specs using **matching hash function** +- **Selects a server-native Python kernel spec** (e.g., `deepnote-venv-{hash}` or fallback to any available Python kernel) - The Deepnote server is started with the venv's Python interpreter, ensuring the kernel uses the venv environment -- Environment variables (`PATH`, `VIRTUAL_ENV`) are configured so the server and kernel use the venv's Python +- Environment variables (`PATH`, `VIRTUAL_ENV`, `JUPYTER_PATH`) are configured so the server and kernel use the venv's Python - Registers the server with the server provider - Creates kernel connection metadata - Registers the controller with VSCode @@ -90,74 +116,77 @@ This implementation adds automatic kernel selection and startup for `.deepnote` - **Reuses existing controllers and servers** for persistent kernel sessions - **Automatically reselects kernel** if it becomes unselected after errors - Tracks controllers per notebook file for efficient reuse +- **Coordinates init notebook execution** via event-driven approach **Key Methods:** -- `activate()`: Registers event listeners for notebook open/close and controller selection changes -- `ensureKernelSelected(notebook)`: Main logic for auto-selection, kernel spec selection, and kernel reuse +- `activate()`: Registers event listeners for notebook open/close, controller selection changes, and kernel starts +- `ensureKernelSelected(notebook)`: Main logic for auto-selection, kernel spec selection, kernel reuse, and init notebook preparation - `onDidOpenNotebook(notebook)`: Event handler for notebook opens - `onControllerSelectionChanged(event)`: Event handler for controller selection changes (auto-reselects if needed) - `onDidCloseNotebook(notebook)`: Event handler for notebook closes (preserves controllers for reuse) +- `onKernelStarted(kernel)`: Event handler for kernel starts (triggers init notebook execution) +- `getVenvHash(fileUri)`: Creates consistent hash for kernel spec naming (must match installer) -#### 6. **Service Registry Updates** (`src/notebooks/serviceRegistry.node.ts`) +#### 8. **Service Registry Updates** (`src/notebooks/serviceRegistry.node.ts`) - Registers all new Deepnote kernel services - Binds `DeepnoteServerProvider` as an activation service - Binds `IDeepnoteKernelAutoSelector` as an activation service -#### 7. **Kernel Types Updates** (`src/kernels/types.ts`) +#### 9. **Kernel Types Updates** (`src/kernels/types.ts`) - Adds `DeepnoteKernelConnectionMetadata` to `RemoteKernelConnectionMetadata` union type - Adds deserialization support for `'startUsingDeepnoteKernel'` kind ## Flow Diagram -```text -User opens .deepnote file - ↓ -DeepnoteKernelAutoSelector.onDidOpenNotebook() - ↓ -Check if kernel already selected → Yes → Exit - ↓ No -Get active Python interpreter - ↓ -DeepnoteToolkitInstaller.ensureInstalled() - ↓ -Extract base file URI (remove query params) - ↓ -Check if venv exists for this file → Yes → Skip to server - ↓ No -Create venv for this .deepnote file - ↓ -Upgrade pip to latest version in venv - ↓ -pip install deepnote-toolkit[server] and ipykernel in venv - ↓ -Verify installation - ↓ -DeepnoteServerStarter.getOrStartServer(venv, fileUri) - ↓ -Check if server running for this file → Yes → Return info - ↓ No -Find available port - ↓ -Start: python -m deepnote_toolkit server --jupyter-port - ↓ -Wait for server to be ready (poll /api endpoint) - ↓ -Register server with DeepnoteServerProvider - ↓ -Query server for available kernel specs - ↓ -Select any available Python kernel spec (e.g., python3-venv) - ↓ -Create DeepnoteKernelConnectionMetadata with server kernel spec - ↓ -Register controller with IControllerRegistration - ↓ -Set controller affinity to Preferred (auto-selects kernel) - ↓ -User runs cell → Executes on Deepnote kernel +```mermaid +flowchart TD + Start([User opens .deepnote file]) --> OpenEvent[DeepnoteKernelAutoSelector.onDidOpenNotebook] + OpenEvent --> CheckSelected{Kernel already selected?} + CheckSelected -->|Yes| Exit([Exit]) + CheckSelected -->|No| GetInterpreter[Get active Python interpreter] + GetInterpreter --> EnsureInstalled[DeepnoteToolkitInstaller.ensureInstalled] + EnsureInstalled --> ExtractURI[Extract base file URI] + ExtractURI --> CheckVenv{Venv exists?} + CheckVenv -->|Yes| GetServer[DeepnoteServerStarter.getOrStartServer] + CheckVenv -->|No| CreateVenv[Create venv for .deepnote file] + CreateVenv --> UpgradePip[Upgrade pip to latest version] + UpgradePip --> InstallToolkit[pip install deepnote-toolkit & ipykernel] + InstallToolkit --> InstallKernelSpec[Install kernel spec with --user] + InstallKernelSpec --> GetServer + GetServer --> CheckServer{Server running?} + CheckServer -->|Yes| RegisterServer[Register server with provider] + CheckServer -->|No| FindPort[Find available port] + FindPort --> StartServer[Start: python -m deepnote_toolkit server] + StartServer --> SetEnv[Set PATH, VIRTUAL_ENV, JUPYTER_PATH] + SetEnv --> WaitServer[Wait for server to be ready] + WaitServer --> RegisterServer + RegisterServer --> QuerySpecs[Query server for kernel specs] + QuerySpecs --> SelectSpec[Select deepnote-venv-hash spec] + SelectSpec --> CreateMetadata[Create DeepnoteKernelConnectionMetadata] + CreateMetadata --> RegisterController[Register controller] + RegisterController --> SetAffinity[Set controller affinity to Preferred] + SetAffinity --> CheckInit{Has initNotebookId?} + CheckInit -->|No| Ready([Controller ready!]) + CheckInit -->|Yes| CreateReq[Create requirements.txt] + CreateReq --> StorePending[Store in projectsPendingInitNotebook] + StorePending --> Ready + Ready --> UserRuns([User runs first cell]) + UserRuns --> CreateKernel[VS Code creates kernel lazily] + CreateKernel --> StartKernel[kernel.start] + StartKernel --> FireEvent[Fire onDidStartKernel event] + FireEvent --> EventHandler[onKernelStarted handler] + EventHandler --> CheckPending{Has pending init?} + CheckPending -->|No| ExecuteUser[Execute user's cell] + CheckPending -->|Yes| RunInit[Run init notebook] + RunInit --> GetKernel[Get kernel - guaranteed to exist!] + GetKernel --> ExecBlocks[Execute code blocks sequentially] + ExecBlocks --> PipInstall[pip install -r requirements.txt] + PipInstall --> MarkRun[Mark as run - cached] + MarkRun --> ExecuteUser + ExecuteUser --> ImportWorks([import pandas ✅]) ``` ## Configuration @@ -184,17 +213,19 @@ User runs cell → Executes on Deepnote kernel - Cells will wait for the kernel to be ready before executing - No kernel selection dialog will appear -3. **Once the progress notification shows "Kernel ready!"**: - -- The loading controller is automatically replaced with the real Deepnote kernel -- Your cells start executing - -4. The extension automatically: +3. **The extension automatically sets up the environment:** - Installs deepnote-toolkit in a dedicated virtual environment (first time only) - Starts a Deepnote server on an available port (if not already running) - Selects the appropriate Deepnote kernel -- Executes your cells +- Creates `requirements.txt` from project settings (if defined) +- Runs the init notebook (if `project.initNotebookId` is defined) + +4. **Once the progress notification shows "Kernel ready!"**: + +- The loading controller is automatically replaced with the real Deepnote kernel +- Init notebook code has been executed (if present) +- Your cells start executing **First-time setup** takes 15-30 seconds. **Subsequent opens** of the same file reuse the existing environment and server, taking less than 1 second. @@ -211,6 +242,254 @@ User runs cell → Executes on Deepnote kernel - **Persistent kernel sessions**: Controllers and servers remain available even after errors - **Automatic recovery**: If kernel becomes unselected, it's automatically reselected - **Seamless reusability**: Run cells as many times as you want without manual kernel selection +- **Init notebook support**: Automatically runs initialization code before user notebooks execute +- **Dependency management**: Creates `requirements.txt` from project settings for easy package installation + +## Init Notebook Feature + +### Overview + +Deepnote projects can define an **initialization notebook** that runs automatically before any user code executes. This feature ensures that the environment is properly configured with required packages and setup code before the main notebook runs. + +### How It Works + +When you open a Deepnote notebook, the extension: + +1. **Checks for `initNotebookId`** in the project YAML (`project.initNotebookId`) +2. **Creates `requirements.txt`** from `project.settings.requirements` array +3. **Finds the init notebook** in the project's notebooks array by ID +4. **Executes all code blocks** from the init notebook sequentially +5. **Shows progress** with a notification: "Running init notebook..." +6. **Caches the execution** so it only runs once per project per session +7. **Allows user code to run** after init notebook completes + +### Example Project Structure + +Here's an example of a Deepnote project YAML with an init notebook: + +```yaml +metadata: + createdAt: 2025-07-21T14:50:41.160Z + modifiedAt: 2025-10-07T11:28:09.117Z +project: + id: 4686ec79-9341-4ac4-8aba-ec0ea497f818 + name: My Data Science Project + initNotebookId: a5356b1e77b34793a815faa71e75aad5 # <-- Init notebook ID + notebooks: + - id: a5356b1e77b34793a815faa71e75aad5 # <-- This is the init notebook + name: Init + blocks: + - type: code + content: | + %%bash + # Install requirements from requirements.txt + if test -f requirements.txt + then + pip install -r ./requirements.txt + else echo "No requirements.txt found." + fi + - id: d8403aaa3cd9462a8051a75b8c1eec42 # <-- This is the main notebook + name: Main Analysis + blocks: + - type: code + content: | + import pandas as pd + import numpy as np + # User code starts here + settings: + requirements: + - pandas + - numpy + - matplotlib + - scikit-learn +version: 1.0.0 +``` + +### Behavior Details + +**Execution Flow:** +1. User opens any notebook from the project +2. Kernel controller is registered and selected +3. `requirements.txt` is created with: `pandas`, `numpy`, `matplotlib`, `scikit-learn` +4. **Init notebook info stored as "pending"** (will run when kernel starts) +5. Controller setup completes (no kernel exists yet) +6. **User runs first cell** → Triggers kernel creation +7. **Kernel starts** → `onDidStartKernel` event fires +8. **Init notebook runs automatically** and blocks cell execution until complete +9. Init notebook installs packages from `requirements.txt` +10. User's cell executes → Packages available! + +**Critical Implementation Details - Event-Driven Approach:** +- **Kernel is NOT explicitly started** (was causing `CancellationError`) +- **Listen to `kernelProvider.onDidStartKernel` event** instead +- Init notebook info is **stored as pending** during controller setup +- When user runs first cell, kernel is created lazily by VS Code +- `onDidStartKernel` fires, triggering init notebook execution +- Init notebook execution is **awaited** (blocks user's cell execution) +- Kernel is guaranteed to exist when init notebook runs (no retry logic needed) +- User's cell waits for init notebook to complete before executing + +**Caching:** +- Init notebook runs **once per project** per VS Code session +- If you open multiple notebooks from the same project, init runs only once +- Cache is cleared when VS Code restarts +- Only marks as "run" if execution actually succeeds + +**Error Handling:** +- If init notebook fails, the error is logged but user can still run cells +- Progress continues even if some blocks fail +- User receives notification about any errors +- Returns `false` if kernel unavailable (won't mark as run, can retry) +- Returns `true` if execution completes (marks as run) + +**Progress Indicators:** +- "Creating requirements.txt..." - File creation phase (during controller setup) +- "Kernel ready!" - Controller selection complete +- [User runs first cell] - Triggers kernel creation +- "Running init notebook..." - Execution phase (appears when kernel starts) +- Shows current block progress: "Executing block 1/3..." +- Displays per-block incremental progress +- User's cell queued until init notebook completes + +**Kernel Spec Discovery:** +- Uses correct hash function to match kernel spec name (djb2-style) +- Both installer and auto-selector use **identical** hash function +- Looks for `deepnote-venv-{hash}` (where hash matches installer exactly) +- Falls back with warnings if venv kernel not found +- Sets `JUPYTER_PATH` so server can find user-installed kernel specs + +**Event-Driven Execution:** +- Init notebook does NOT run during controller setup (kernel doesn't exist yet) +- Stores pending init info in `projectsPendingInitNotebook` map +- Listens to `kernelProvider.onDidStartKernel` event +- When user runs first cell and kernel is created, event fires +- Init notebook runs in the event handler when kernel is guaranteed to exist +- User's cell execution is blocked until init notebook completes + +### Use Cases + +**1. Package Installation** +```python +%%bash +pip install -r requirements.txt +``` + +**2. Environment Setup** +```python +import os +os.environ['DATA_PATH'] = '/workspace/data' +os.environ['MODEL_PATH'] = '/workspace/models' +``` + +**3. Database Connections** +```python +import sqlalchemy +engine = sqlalchemy.create_engine('postgresql://...') +``` + +**4. Data Preprocessing** +```python +import pandas as pd +# Preload common datasets +df = pd.read_csv('data/master_dataset.csv') +``` + +**5. Custom Imports** +```python +import sys +sys.path.append('./custom_modules') +from my_utils import helper_functions +``` + +### Benefits of Init Notebooks + +- ✅ **Consistent environment**: Every notebook in the project gets the same setup +- ✅ **Automated dependency installation**: No manual pip install commands needed +- ✅ **One-time setup**: Runs once per session, not every time you open a notebook +- ✅ **Transparent to users**: Happens automatically in the background +- ✅ **Error resilient**: Failures don't block user from running their code +- ✅ **Progress visibility**: Clear notifications show what's happening +- ✅ **Session persistence**: Init state is preserved across notebook switches + +### Configuration in Deepnote Projects + +To add an init notebook to your Deepnote project: + +1. Create a notebook named "Init" (or any name) +2. Add your setup code blocks to this notebook +3. Note the notebook's ID from the YAML +4. Add `initNotebookId: ` to your `project` section +5. Optionally add a `requirements` array under `project.settings` + +The extension will automatically detect and run this notebook when the project is opened. + +### Limitations + +- Init notebook must be a notebook within the same project +- Only code blocks are executed (markdown blocks are skipped) +- Runs in the same kernel as the main notebook (shared state) +- Cannot be cancelled once started +- Runs only once per project per VS Code session +- Requires user to run at least one cell to trigger kernel creation and init notebook execution + +### Technical Notes + +**Why Event-Driven Approach?** + +VS Code creates notebook kernels **lazily** - they don't exist when controllers are selected. The kernel is only created when: +1. User executes a cell, OR +2. Controller explicitly calls `startKernel()` (but this can throw `CancellationError`) + +The event-driven approach is more reliable: +- No `CancellationError` issues +- Kernel is guaranteed to exist when `onDidStartKernel` fires +- Natural integration with VS Code's lazy kernel creation +- Simpler code without complex retry logic + +**Event-Driven Flow Diagram:** + +```mermaid +sequenceDiagram + participant User + participant VSCode + participant AutoSelector as Kernel Auto-Selector + participant KernelProvider + participant InitRunner as Init Notebook Runner + participant Kernel + + User->>VSCode: Opens .deepnote file + VSCode->>AutoSelector: onDidOpenNotebook() + AutoSelector->>AutoSelector: ensureKernelSelected() + AutoSelector->>AutoSelector: Register controller + AutoSelector->>AutoSelector: Create requirements.txt + + alt Has initNotebookId + AutoSelector->>AutoSelector: Store in projectsPendingInitNotebook + Note over AutoSelector: "Init notebook will run when kernel starts" + end + + AutoSelector-->>User: Kernel ready! (but kernel not created yet) + + User->>VSCode: Runs first cell + VSCode->>KernelProvider: getOrCreate(notebook) + KernelProvider->>Kernel: Create kernel instance + Kernel->>Kernel: start() + Kernel->>KernelProvider: Fire onDidStartKernel event + KernelProvider->>AutoSelector: onKernelStarted(kernel) + + alt Has pending init notebook + AutoSelector->>InitRunner: runInitNotebookIfNeeded(notebook, projectId) + InitRunner->>KernelProvider: get(notebook) ✅ exists! + InitRunner->>Kernel: executeHidden(block.content) + Note over InitRunner,Kernel: pip install -r requirements.txt + InitRunner->>InitRunner: Mark as run (cached) + InitRunner-->>AutoSelector: Init complete ✅ + end + + AutoSelector-->>VSCode: Init notebook done + VSCode->>Kernel: Execute user's cell + Kernel-->>User: import pandas ✅ works! +``` ## UI Customization @@ -290,11 +569,21 @@ To test the implementation: - `src/kernels/deepnote/deepnoteServerStarter.node.ts` - Server lifecycle management - `src/kernels/deepnote/deepnoteServerProvider.node.ts` - Jupyter server provider implementation - `src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts` - Automatic kernel selection +- `src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts` - Init notebook execution service +- `src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts` - Requirements.txt creation helper ### Modified: - `src/kernels/types.ts` - Added DeepnoteKernelConnectionMetadata to union types -- `src/notebooks/serviceRegistry.node.ts` - Registered new services +- `src/notebooks/serviceRegistry.node.ts` - Registered new services including init notebook runner +- `src/notebooks/deepnote/deepnoteNotebookManager.ts` - Added init notebook execution tracking +- `src/notebooks/deepnote/deepnoteTypes.ts` - Added initNotebookId field to project type +- `src/notebooks/types.ts` - Updated IDeepnoteNotebookManager interface +- `src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts` - Removed duplicate hash function (now calls `toolkitInstaller.getVenvHash()`), added event-driven init notebook execution via `onDidStartKernel`, added pending init notebook tracking, changed to await init notebook +- `src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts` - Changed to return boolean for success/failure, simplified by removing retry logic (kernel guaranteed to exist), added dual progress bars +- `src/kernels/deepnote/deepnoteServerStarter.node.ts` - Added `JUPYTER_PATH` environment variable for kernel spec discovery +- `src/kernels/deepnote/deepnoteToolkitInstaller.node.ts` - Made `getVenvHash()` public for reuse by auto-selector +- `src/kernels/deepnote/types.ts` - Added `getVenvHash()` to `IDeepnoteToolkitInstaller` interface ## Dependencies @@ -336,12 +625,16 @@ The implementation uses a robust hashing approach for virtual environment direct 1. **Path Hashing**: Uses `getVenvHash()` to create short, unique identifiers from file paths 2. **Hash Algorithm**: Implements a djb2-style hash function for better distribution 3. **Format**: Generates identifiers like `venv_a1b2c3d4` (max 16 characters) -4. **Benefits**: +4. **Shared Implementation**: The hash function is defined in `DeepnoteToolkitInstaller` and exposed as a public method +5. **Consistent Usage**: Both the toolkit installer and kernel auto-selector call `toolkitInstaller.getVenvHash()` +6. **Benefits**: - Avoids Windows MAX_PATH (260 character) limitations - Prevents directory structure leakage into extension storage - Provides consistent naming for both venv directories and kernel specs - Reduces collision risk with better hash distribution +- **Single source of truth** - hash function defined in one place only +- **Guaranteed consistency** - both components use the exact same implementation ## Troubleshooting & Key Fixes diff --git a/README.md b/README.md index e862677183..fb5819eec1 100644 --- a/README.md +++ b/README.md @@ -9,6 +9,8 @@ This extension allows you to work with Deepnote notebooks in VS Code. Deepnote n - **More block types** - Choose from SQL blocks, chart blocks, and more specialized data science blocks - **Seamless language switching** - Switch between Python and SQL seamlessly within the same notebook - **Database integrations** - Connect directly to Postgres, Snowflake, BigQuery and more data sources +- **Init notebooks** - Automatically runs initialization code (like dependency installation) before your notebooks execute +- **Project requirements** - Automatically creates `requirements.txt` from your project settings for easy dependency management ## Useful commands diff --git a/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts b/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts index 7aa5c19177..38ca4b3f00 100644 --- a/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts +++ b/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts @@ -262,7 +262,7 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { } } - private getVenvHash(deepnoteFileUri: Uri): string { + public getVenvHash(deepnoteFileUri: Uri): string { // Create a short hash from the file path for kernel naming and venv directory // This provides better uniqueness and prevents directory structure leakage const path = deepnoteFileUri.fsPath; diff --git a/src/kernels/deepnote/types.ts b/src/kernels/deepnote/types.ts index 82df017dd4..622ae0abc9 100644 --- a/src/kernels/deepnote/types.ts +++ b/src/kernels/deepnote/types.ts @@ -85,6 +85,13 @@ export interface IDeepnoteToolkitInstaller { * @param deepnoteFileUri The URI of the .deepnote file */ getVenvInterpreter(deepnoteFileUri: vscode.Uri): Promise; + + /** + * Gets the hash for the venv directory/kernel spec name based on file path. + * @param deepnoteFileUri The URI of the .deepnote file + * @returns The hash string used for venv directory and kernel spec naming + */ + getVenvHash(deepnoteFileUri: vscode.Uri): string; } export const IDeepnoteServerStarter = Symbol('IDeepnoteServerStarter'); diff --git a/src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts b/src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts new file mode 100644 index 0000000000..20fe6931dd --- /dev/null +++ b/src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts @@ -0,0 +1,230 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { inject, injectable } from 'inversify'; +import { NotebookDocument, ProgressLocation, window, CancellationTokenSource } from 'vscode'; +import { logger } from '../../platform/logging'; +import { IDeepnoteNotebookManager } from '../types'; +import { DeepnoteProject, DeepnoteNotebook } from './deepnoteTypes'; +import { IKernelProvider } from '../../kernels/types'; +import { getDisplayPath } from '../../platform/common/platform/fs-paths'; + +/** + * Service responsible for running init notebooks before the main notebook starts. + * Init notebooks typically contain setup code like pip installs. + */ +@injectable() +export class DeepnoteInitNotebookRunner { + constructor( + @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, + @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider + ) {} + + /** + * Runs the init notebook if it exists and hasn't been run yet for this project. + * This should be called after the kernel is started but before user code executes. + * @param notebook The notebook document + * @param projectId The Deepnote project ID + */ + async runInitNotebookIfNeeded(notebook: NotebookDocument, projectId: string): Promise { + try { + // Check if init notebook has already run for this project + if (this.notebookManager.hasInitNotebookRun(projectId)) { + logger.info(`Init notebook already run for project ${projectId}, skipping`); + return; + } + + // Get the project data + const project = this.notebookManager.getOriginalProject(projectId) as DeepnoteProject | undefined; + if (!project) { + logger.warn(`Project ${projectId} not found, cannot run init notebook`); + return; + } + + // Check if project has an init notebook ID + const initNotebookId = (project.project as { initNotebookId?: string }).initNotebookId; + if (!initNotebookId) { + logger.info(`No init notebook configured for project ${projectId}`); + // Mark as run so we don't check again + this.notebookManager.markInitNotebookAsRun(projectId); + return; + } + + // Find the init notebook + const initNotebook = project.project.notebooks.find((nb) => nb.id === initNotebookId); + if (!initNotebook) { + logger.warn( + `Init notebook ${initNotebookId} not found in project ${projectId}, skipping initialization` + ); + this.notebookManager.markInitNotebookAsRun(projectId); + return; + } + + logger.info(`Running init notebook "${initNotebook.name}" (${initNotebookId}) for project ${projectId}`); + + // Execute the init notebook with progress + const success = await this.executeInitNotebook(notebook, initNotebook); + + if (success) { + // Mark as run so we don't run it again + this.notebookManager.markInitNotebookAsRun(projectId); + logger.info(`Init notebook completed successfully for project ${projectId}`); + } else { + logger.warn(`Init notebook did not execute for project ${projectId} - kernel not available`); + } + } catch (error) { + // Log error but don't throw - we want to let user continue anyway + logger.error(`Error running init notebook for project ${projectId}:`, error); + console.error(`Failed to run init notebook:`, error); + // Still mark as run to avoid retrying on every notebook open + this.notebookManager.markInitNotebookAsRun(projectId); + } + } + + /** + * Executes the init notebook's code blocks in the kernel. + * @param notebook The notebook document (for kernel context) + * @param initNotebook The init notebook to execute + * @returns True if execution completed, false if kernel was not available + */ + private async executeInitNotebook(notebook: NotebookDocument, initNotebook: DeepnoteNotebook): Promise { + // Show progress in both notification AND window for maximum visibility + const cancellationTokenSource = new CancellationTokenSource(); + + // Create a wrapper that reports to both progress locations + const executeWithDualProgress = async () => { + return window.withProgress( + { + location: ProgressLocation.Notification, + title: `🚀 Initializing project environment`, + cancellable: false + }, + async (notificationProgress) => { + return window.withProgress( + { + location: ProgressLocation.Window, + title: `Init: "${initNotebook.name}"`, + cancellable: false + }, + async (windowProgress) => { + // Helper to report to both progress bars + const reportProgress = (message: string, increment: number) => { + notificationProgress.report({ message, increment }); + windowProgress.report({ message, increment }); + }; + + return this.executeInitNotebookImpl( + notebook, + initNotebook, + reportProgress, + cancellationTokenSource.token + ); + } + ); + } + ); + }; + + try { + return await executeWithDualProgress(); + } finally { + cancellationTokenSource.dispose(); + } + } + + private async executeInitNotebookImpl( + notebook: NotebookDocument, + initNotebook: DeepnoteNotebook, + progress: (message: string, increment: number) => void, + _token: unknown + ): Promise { + try { + progress(`Running init notebook "${initNotebook.name}"...`, 0); + + // Get the kernel for this notebook + // Note: This should always exist because onKernelStarted already fired + const kernel = this.kernelProvider.get(notebook); + if (!kernel) { + logger.error( + `No kernel found for ${getDisplayPath( + notebook.uri + )} even after onDidStartKernel fired - this should not happen` + ); + return false; + } + + logger.info(`Kernel found for ${getDisplayPath(notebook.uri)}, starting init notebook execution`); + + // Filter out non-code blocks + const codeBlocks = initNotebook.blocks.filter((block) => block.type === 'code'); + + if (codeBlocks.length === 0) { + logger.info(`Init notebook has no code blocks, skipping execution`); + return true; // Not an error - just nothing to execute + } + + logger.info(`Executing ${codeBlocks.length} code blocks from init notebook`); + progress( + `Preparing to execute ${codeBlocks.length} initialization ${ + codeBlocks.length === 1 ? 'block' : 'blocks' + }...`, + 5 + ); + + // Get kernel execution + const kernelExecution = this.kernelProvider.getKernelExecution(kernel); + + // Execute each code block sequentially + for (let i = 0; i < codeBlocks.length; i++) { + const block = codeBlocks[i]; + const percentComplete = Math.floor((i / codeBlocks.length) * 100); + + // Show more detailed progress with percentage + progress( + `[${percentComplete}%] Executing block ${i + 1} of ${codeBlocks.length}...`, + 90 / codeBlocks.length // Reserve 5% for start, 5% for finish + ); + + logger.info(`Executing init notebook block ${i + 1}/${codeBlocks.length}`); + + try { + // Execute the code silently in the background + const outputs = await kernelExecution.executeHidden(block.content ?? ''); + + // Log outputs for debugging + if (outputs && outputs.length > 0) { + logger.info(`Init notebook block ${i + 1} produced ${outputs.length} outputs`); + + // Check for errors in outputs + const errors = outputs.filter( + (output: { output_type?: string }) => output.output_type === 'error' + ); + if (errors.length > 0) { + logger.warn(`Init notebook block ${i + 1} produced errors:`, errors); + } + } + } catch (blockError) { + // Log error but continue with next block + logger.error(`Error executing init notebook block ${i + 1}:`, blockError); + console.error(`Error in init notebook block ${i + 1}:`, blockError); + } + } + + logger.info(`Completed executing all init notebook blocks`); + progress(`✓ Initialization complete! Environment ready.`, 5); + + // Give user a moment to see the completion message + await new Promise((resolve) => setTimeout(resolve, 1000)); + + return true; + } catch (error) { + logger.error(`Error in executeInitNotebook:`, error); + throw error; + } + } +} + +export const IDeepnoteInitNotebookRunner = Symbol('IDeepnoteInitNotebookRunner'); +export interface IDeepnoteInitNotebookRunner { + runInitNotebookIfNeeded(notebook: NotebookDocument, projectId: string): Promise; +} diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index ac9521d7c9..bdac560834 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -34,6 +34,11 @@ import { createJupyterConnectionInfo } from '../../kernels/jupyter/jupyterUtils' import { IJupyterRequestCreator, IJupyterRequestAgentCreator } from '../../kernels/jupyter/types'; import { IConfigurationService } from '../../platform/common/types'; import { disposeAsync } from '../../platform/common/utils'; +import { IDeepnoteInitNotebookRunner } from './deepnoteInitNotebookRunner.node'; +import { IDeepnoteNotebookManager } from '../types'; +import { DeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; +import { DeepnoteProject } from './deepnoteTypes'; +import { IKernelProvider, IKernel } from '../../kernels/types'; /** * Automatically selects and starts Deepnote kernel for .deepnote notebooks @@ -48,6 +53,11 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, private readonly notebookConnectionMetadata = new Map(); // Track temporary loading controllers that get disposed when real controller is ready private readonly loadingControllers = new Map(); + // Track projects where we need to run init notebook (set during controller setup) + private readonly projectsPendingInitNotebook = new Map< + string, + { notebook: NotebookDocument; project: DeepnoteProject } + >(); constructor( @inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry, @@ -61,7 +71,10 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, @inject(IJupyterRequestAgentCreator) @optional() private readonly requestAgentCreator: IJupyterRequestAgentCreator | undefined, - @inject(IConfigurationService) private readonly configService: IConfigurationService + @inject(IConfigurationService) private readonly configService: IConfigurationService, + @inject(IDeepnoteInitNotebookRunner) private readonly initNotebookRunner: IDeepnoteInitNotebookRunner, + @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, + @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider ) {} public activate() { @@ -79,6 +92,10 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, this.disposables ); + // Listen to kernel starts to run init notebooks + // Kernels are created lazily when cells are executed, so this is the right time to run init notebook + this.kernelProvider.onDidStartKernel(this.onKernelStarted, this, this.disposables); + // Handle currently open notebooks - await all async operations Promise.all(workspace.notebookDocuments.map((d) => this.onDidOpenNotebook(d))).catch((error) => { logger.error(`Error handling open notebooks during activation: ${error}`); @@ -165,6 +182,45 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, this.serverProvider.unregisterServer(serverHandle); this.notebookServerHandles.delete(notebookKey); } + + // Clean up pending init notebook tracking + const projectId = notebook.metadata?.deepnoteProjectId; + if (projectId) { + this.projectsPendingInitNotebook.delete(projectId); + } + } + + private async onKernelStarted(kernel: IKernel) { + // Only handle deepnote notebooks + if (kernel.notebook?.notebookType !== DEEPNOTE_NOTEBOOK_TYPE) { + return; + } + + const notebook = kernel.notebook; + const projectId = notebook.metadata?.deepnoteProjectId; + + if (!projectId) { + return; + } + + // Check if we have a pending init notebook for this project + const pendingInit = this.projectsPendingInitNotebook.get(projectId); + if (!pendingInit) { + return; // No init notebook to run + } + + logger.info(`Kernel started for Deepnote notebook, running init notebook for project ${projectId}`); + + // Remove from pending list + this.projectsPendingInitNotebook.delete(projectId); + + // Run init notebook now that kernel is available + try { + await this.initNotebookRunner.runInitNotebookIfNeeded(notebook, projectId); + } catch (error) { + logger.error(`Error running init notebook: ${error}`); + // Continue anyway - don't block user if init fails + } } public async ensureKernelSelected(notebook: NotebookDocument, _token?: CancellationToken): Promise { @@ -316,25 +372,33 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, `Available kernel specs on Deepnote server: ${kernelSpecs.map((s) => s.name).join(', ')}` ); - // Create expected kernel name based on file path (same logic as in installer) - const safePath = baseFileUri.fsPath.replace(/[^a-zA-Z0-9]/g, '_'); - const venvHash = safePath.substring(0, 16); + // Create expected kernel name based on file path (uses installer's hash logic) + const venvHash = this.toolkitInstaller.getVenvHash(baseFileUri); const expectedKernelName = `deepnote-venv-${venvHash}`; logger.info(`Looking for venv kernel spec: ${expectedKernelName}`); // Prefer the venv kernel spec that uses the venv's Python interpreter // This ensures packages installed via pip are available to the kernel - kernelSpec = - kernelSpecs.find((s) => s.name === expectedKernelName) || - kernelSpecs.find((s) => s.language === 'python') || - kernelSpecs.find((s) => s.name === 'python3-venv') || - kernelSpecs[0]; + kernelSpec = kernelSpecs.find((s) => s.name === expectedKernelName); + + if (!kernelSpec) { + logger.warn( + `⚠️ Venv kernel spec '${expectedKernelName}' not found! Falling back to generic Python kernel.` + ); + logger.warn( + `This may cause import errors if packages are installed to the venv but kernel uses system Python.` + ); + kernelSpec = + kernelSpecs.find((s) => s.language === 'python') || + kernelSpecs.find((s) => s.name === 'python3-venv') || + kernelSpecs[0]; + } if (!kernelSpec) { throw new Error('No kernel specs available on Deepnote server'); } - logger.info(`Using kernel spec: ${kernelSpec.name} (${kernelSpec.display_name})`); + logger.info(`✓ Using kernel spec: ${kernelSpec.name} (${kernelSpec.display_name})`); } finally { await disposeAsync(sessionManager); } @@ -395,6 +459,30 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, this.loadingControllers.delete(notebookKey); logger.info(`Disposed loading controller for ${notebookKey}`); } + + // Prepare init notebook execution for when kernel starts + const projectId = notebook.metadata?.deepnoteProjectId; + if (projectId) { + // Get the project and create requirements.txt before kernel starts + const project = this.notebookManager.getOriginalProject(projectId) as + | DeepnoteProject + | undefined; + if (project) { + // Create requirements.txt first (needs to be ready for init notebook) + progress.report({ message: 'Creating requirements.txt...' }); + await DeepnoteRequirementsHelper.createRequirementsFile(project); + + // Check if project has an init notebook + if (project.project.initNotebookId && !this.notebookManager.hasInitNotebookRun(projectId)) { + // Store for execution when kernel actually starts + // Kernels are created lazily when cells execute, so we can't run init notebook now + this.projectsPendingInitNotebook.set(projectId, { notebook, project }); + logger.info( + `Init notebook will run automatically when kernel starts for project ${projectId}` + ); + } + } + } } catch (ex) { logger.error(`Failed to auto-select Deepnote kernel: ${ex}`); throw ex; diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.ts b/src/notebooks/deepnote/deepnoteNotebookManager.ts index e556ff4e46..be69780716 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.ts @@ -12,6 +12,7 @@ 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(); /** * Gets the currently selected notebook ID for a project. @@ -73,4 +74,21 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { updateCurrentNotebookId(projectId: string, notebookId: string): void { this.currentNotebookId.set(projectId, notebookId); } + + /** + * Checks if the init notebook has already been run for a project. + * @param projectId Project identifier + * @returns True if init notebook has been run, false otherwise + */ + hasInitNotebookRun(projectId: string): boolean { + return this.projectsWithInitNotebookRun.has(projectId); + } + + /** + * Marks the init notebook as having been run for a project. + * @param projectId Project identifier + */ + markInitNotebookAsRun(projectId: string): void { + this.projectsWithInitNotebookRun.add(projectId); + } } diff --git a/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts b/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts new file mode 100644 index 0000000000..b9afdb53dc --- /dev/null +++ b/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts @@ -0,0 +1,45 @@ +// Copyright (c) Microsoft Corporation. +// Licensed under the MIT License. + +import { workspace } from 'vscode'; +import * as fs from 'fs'; +import * as path from 'path'; +import type { DeepnoteProject } from './deepnoteTypes'; + +/** + * Helper class for creating requirements.txt files from Deepnote project settings. + */ +export class DeepnoteRequirementsHelper { + /** + * Extracts requirements from project settings and creates a local requirements.txt file. + * @param project The Deepnote project data containing requirements in settings + */ + static async createRequirementsFile(project: DeepnoteProject): Promise { + try { + const requirements = project.project.settings?.requirements; + if (!requirements || !Array.isArray(requirements) || requirements.length === 0) { + console.log(`No requirements found in project ${project.project.id}`); + return; + } + + // Get the workspace folder to determine where to create the requirements.txt file + const workspaceFolders = workspace.workspaceFolders; + if (!workspaceFolders || workspaceFolders.length === 0) { + console.log('No workspace folder found, cannot create requirements.txt'); + return; + } + + const workspaceRoot = workspaceFolders[0].uri.fsPath; + const requirementsPath = path.join(workspaceRoot, 'requirements.txt'); + + // Convert requirements array to text format + const requirementsText = requirements.join('\n') + '\n'; + + // Write the requirements.txt file + await fs.promises.writeFile(requirementsPath, requirementsText, 'utf8'); + console.log(`Created requirements.txt with ${requirements.length} dependencies at ${requirementsPath}`); + } catch (error) { + console.error(`Error creating requirements.txt:`, error); + } + } +} diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index 391aaf4679..81335e4d1e 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -53,6 +53,7 @@ import { DeepnoteToolkitInstaller } from '../kernels/deepnote/deepnoteToolkitIns import { DeepnoteServerStarter } from '../kernels/deepnote/deepnoteServerStarter.node'; import { DeepnoteKernelAutoSelector } from './deepnote/deepnoteKernelAutoSelector.node'; import { DeepnoteServerProvider } from '../kernels/deepnote/deepnoteServerProvider.node'; +import { DeepnoteInitNotebookRunner, IDeepnoteInitNotebookRunner } from './deepnote/deepnoteInitNotebookRunner.node'; export function registerTypes(serviceManager: IServiceManager, isDevMode: boolean) { registerControllerTypes(serviceManager, isDevMode); @@ -135,6 +136,7 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea serviceManager.addBinding(IDeepnoteServerProvider, IExtensionSyncActivationService); serviceManager.addSingleton(IDeepnoteKernelAutoSelector, DeepnoteKernelAutoSelector); serviceManager.addBinding(IDeepnoteKernelAutoSelector, IExtensionSyncActivationService); + serviceManager.addSingleton(IDeepnoteInitNotebookRunner, DeepnoteInitNotebookRunner); // File export/import serviceManager.addSingleton(IFileConverter, FileConverter); diff --git a/src/notebooks/types.ts b/src/notebooks/types.ts index 1a1cab0c3f..6e37b30d8e 100644 --- a/src/notebooks/types.ts +++ b/src/notebooks/types.ts @@ -32,4 +32,6 @@ export interface IDeepnoteNotebookManager { selectNotebookForProject(projectId: string, notebookId: string): void; storeOriginalProject(projectId: string, project: unknown, notebookId: string): void; updateCurrentNotebookId(projectId: string, notebookId: string): void; + hasInitNotebookRun(projectId: string): boolean; + markInitNotebookAsRun(projectId: string): void; } From 6b3b152bb2f2fef481cf134b07920ff8869bb025 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Filip=20P=C3=BDrek?= Date: Wed, 8 Oct 2025 12:58:52 +0200 Subject: [PATCH 2/4] fix: coderabbit feedback --- src/extension.common.ts | 2 + .../deepnoteInitNotebookRunner.node.ts | 2 - .../deepnoteKernelAutoSelector.node.ts | 7 +-- .../deepnoteRequirementsHelper.node.ts | 46 +++++++++++++++---- src/notebooks/serviceRegistry.node.ts | 2 + src/platform/logging/types.ts | 1 + 6 files changed, 47 insertions(+), 13 deletions(-) diff --git a/src/extension.common.ts b/src/extension.common.ts index 6b3052a0f4..5f98c14739 100644 --- a/src/extension.common.ts +++ b/src/extension.common.ts @@ -37,6 +37,7 @@ import { import { Common } from './platform/common/utils/localize'; import { IServiceContainer, IServiceManager } from './platform/ioc/types'; import { initializeLoggers as init, logger } from './platform/logging'; +import { ILogger } from './platform/logging/types'; import { getJupyterOutputChannel } from './standalone/devTools/jupyterOutputChannel'; import { isUsingPylance } from './standalone/intellisense/notebookPythonPathService'; import { noop } from './platform/common/utils/misc'; @@ -127,6 +128,7 @@ export function initializeGlobals( getJupyterOutputChannel(context.subscriptions), JUPYTER_OUTPUT_CHANNEL ); + serviceManager.addSingletonInstance(ILogger, logger); return [serviceManager, serviceContainer]; } diff --git a/src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts b/src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts index 20fe6931dd..efbcea9208 100644 --- a/src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts +++ b/src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts @@ -75,7 +75,6 @@ export class DeepnoteInitNotebookRunner { } catch (error) { // Log error but don't throw - we want to let user continue anyway logger.error(`Error running init notebook for project ${projectId}:`, error); - console.error(`Failed to run init notebook:`, error); // Still mark as run to avoid retrying on every notebook open this.notebookManager.markInitNotebookAsRun(projectId); } @@ -206,7 +205,6 @@ export class DeepnoteInitNotebookRunner { } catch (blockError) { // Log error but continue with next block logger.error(`Error executing init notebook block ${i + 1}:`, blockError); - console.error(`Error in init notebook block ${i + 1}:`, blockError); } } diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index bdac560834..1b5c9f4e88 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -36,7 +36,7 @@ import { IConfigurationService } from '../../platform/common/types'; import { disposeAsync } from '../../platform/common/utils'; import { IDeepnoteInitNotebookRunner } from './deepnoteInitNotebookRunner.node'; import { IDeepnoteNotebookManager } from '../types'; -import { DeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; +import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; import { DeepnoteProject } from './deepnoteTypes'; import { IKernelProvider, IKernel } from '../../kernels/types'; @@ -74,7 +74,8 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, @inject(IConfigurationService) private readonly configService: IConfigurationService, @inject(IDeepnoteInitNotebookRunner) private readonly initNotebookRunner: IDeepnoteInitNotebookRunner, @inject(IDeepnoteNotebookManager) private readonly notebookManager: IDeepnoteNotebookManager, - @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider + @inject(IKernelProvider) private readonly kernelProvider: IKernelProvider, + @inject(IDeepnoteRequirementsHelper) private readonly requirementsHelper: IDeepnoteRequirementsHelper ) {} public activate() { @@ -470,7 +471,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, if (project) { // Create requirements.txt first (needs to be ready for init notebook) progress.report({ message: 'Creating requirements.txt...' }); - await DeepnoteRequirementsHelper.createRequirementsFile(project); + await this.requirementsHelper.createRequirementsFile(project, progressToken); // Check if project has an init notebook if (project.project.initNotebookId && !this.notebookManager.hasInitNotebookRun(projectId)) { diff --git a/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts b/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts index b9afdb53dc..a79124aca5 100644 --- a/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts +++ b/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts @@ -1,31 +1,47 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -import { workspace } from 'vscode'; +import { inject, injectable } from 'inversify'; +import { workspace, CancellationToken } from 'vscode'; import * as fs from 'fs'; -import * as path from 'path'; +import * as path from '../../platform/vscode-path/path'; import type { DeepnoteProject } from './deepnoteTypes'; +import { ILogger } from '../../platform/logging/types'; /** * Helper class for creating requirements.txt files from Deepnote project settings. */ +@injectable() export class DeepnoteRequirementsHelper { + constructor(@inject(ILogger) private readonly logger: ILogger) {} + /** * Extracts requirements from project settings and creates a local requirements.txt file. * @param project The Deepnote project data containing requirements in settings + * @param token Cancellation token to abort the operation if needed */ - static async createRequirementsFile(project: DeepnoteProject): Promise { + async createRequirementsFile(project: DeepnoteProject, token: CancellationToken): Promise { try { - const requirements = project.project.settings?.requirements; + // Check if the operation has been cancelled + if (token.isCancellationRequested) { + return; + } + + const requirements = project.project.settings.requirements; if (!requirements || !Array.isArray(requirements) || requirements.length === 0) { - console.log(`No requirements found in project ${project.project.id}`); + this.logger.info(`No requirements found in project ${project.project.id}`); return; } // Get the workspace folder to determine where to create the requirements.txt file const workspaceFolders = workspace.workspaceFolders; if (!workspaceFolders || workspaceFolders.length === 0) { - console.log('No workspace folder found, cannot create requirements.txt'); + this.logger.info('No workspace folder found, cannot create requirements.txt'); + return; + } + + // Check cancellation before performing I/O + if (token.isCancellationRequested) { return; } @@ -37,9 +53,23 @@ export class DeepnoteRequirementsHelper { // Write the requirements.txt file await fs.promises.writeFile(requirementsPath, requirementsText, 'utf8'); - console.log(`Created requirements.txt with ${requirements.length} dependencies at ${requirementsPath}`); + + // Check cancellation after I/O operation + if (token.isCancellationRequested) { + this.logger.info('Requirements file creation was cancelled after write'); + return; + } + + this.logger.info( + `Created requirements.txt with ${requirements.length} dependencies at ${requirementsPath}` + ); } catch (error) { - console.error(`Error creating requirements.txt:`, error); + this.logger.error(`Error creating requirements.txt:`, error); } } } + +export const IDeepnoteRequirementsHelper = Symbol('IDeepnoteRequirementsHelper'); +export interface IDeepnoteRequirementsHelper { + createRequirementsFile(project: DeepnoteProject, token: CancellationToken): Promise; +} diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index 81335e4d1e..6cb19391a5 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -54,6 +54,7 @@ import { DeepnoteServerStarter } from '../kernels/deepnote/deepnoteServerStarter import { DeepnoteKernelAutoSelector } from './deepnote/deepnoteKernelAutoSelector.node'; import { DeepnoteServerProvider } from '../kernels/deepnote/deepnoteServerProvider.node'; import { DeepnoteInitNotebookRunner, IDeepnoteInitNotebookRunner } from './deepnote/deepnoteInitNotebookRunner.node'; +import { DeepnoteRequirementsHelper, IDeepnoteRequirementsHelper } from './deepnote/deepnoteRequirementsHelper.node'; export function registerTypes(serviceManager: IServiceManager, isDevMode: boolean) { registerControllerTypes(serviceManager, isDevMode); @@ -137,6 +138,7 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea serviceManager.addSingleton(IDeepnoteKernelAutoSelector, DeepnoteKernelAutoSelector); serviceManager.addBinding(IDeepnoteKernelAutoSelector, IExtensionSyncActivationService); serviceManager.addSingleton(IDeepnoteInitNotebookRunner, DeepnoteInitNotebookRunner); + serviceManager.addSingleton(IDeepnoteRequirementsHelper, DeepnoteRequirementsHelper); // File export/import serviceManager.addSingleton(IFileConverter, FileConverter); diff --git a/src/platform/logging/types.ts b/src/platform/logging/types.ts index bb8717b654..d03f3c1870 100644 --- a/src/platform/logging/types.ts +++ b/src/platform/logging/types.ts @@ -6,6 +6,7 @@ export type Arguments = unknown[]; +export const ILogger = Symbol('ILogger'); export interface ILogger { error(message: string, ...data: Arguments): void; warn(message: string, ...data: Arguments): void; From fe5819a4a351dde89e69defb58051782b66d86da Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Filip=20P=C3=BDrek?= Date: Wed, 8 Oct 2025 15:14:38 +0200 Subject: [PATCH 3/4] fix: address code review feedback --- DEEPNOTE_KERNEL_IMPLEMENTATION.md | 6 +- .../deepnoteInitNotebookRunner.node.ts | 80 ++++++++++++++--- .../deepnoteKernelAutoSelector.node.ts | 85 +++++++++++++------ .../deepnote/deepnoteNotebookManager.ts | 2 +- .../deepnoteRequirementsHelper.node.ts | 5 +- src/notebooks/types.ts | 2 +- 6 files changed, 133 insertions(+), 47 deletions(-) diff --git a/DEEPNOTE_KERNEL_IMPLEMENTATION.md b/DEEPNOTE_KERNEL_IMPLEMENTATION.md index 1fc516a52b..7a125e6d26 100644 --- a/DEEPNOTE_KERNEL_IMPLEMENTATION.md +++ b/DEEPNOTE_KERNEL_IMPLEMENTATION.md @@ -87,7 +87,7 @@ This implementation adds automatic kernel selection and startup for `.deepnote` **Key Methods:** -- `runInitNotebookIfNeeded(notebook, projectId)`: Main entry point, checks cache and executes if needed +- `runInitNotebookIfNeeded(projectId, notebook)`: Main entry point, checks cache and executes if needed - `executeInitNotebook(notebook, initNotebook)`: Executes all code blocks from init notebook #### 6. **Deepnote Requirements Helper** (`src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts`) @@ -213,7 +213,7 @@ flowchart TD - Cells will wait for the kernel to be ready before executing - No kernel selection dialog will appear -3. **The extension automatically sets up the environment:** +1. **The extension automatically sets up the environment:** - Installs deepnote-toolkit in a dedicated virtual environment (first time only) - Starts a Deepnote server on an available port (if not already running) @@ -221,7 +221,7 @@ flowchart TD - Creates `requirements.txt` from project settings (if defined) - Runs the init notebook (if `project.initNotebookId` is defined) -4. **Once the progress notification shows "Kernel ready!"**: +1. **Once the progress notification shows "Kernel ready!"**: - The loading controller is automatically replaced with the real Deepnote kernel - Init notebook code has been executed (if present) diff --git a/src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts b/src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts index efbcea9208..668d9a0482 100644 --- a/src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts +++ b/src/notebooks/deepnote/deepnoteInitNotebookRunner.node.ts @@ -1,8 +1,5 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - import { inject, injectable } from 'inversify'; -import { NotebookDocument, ProgressLocation, window, CancellationTokenSource } from 'vscode'; +import { NotebookDocument, ProgressLocation, window, CancellationTokenSource, CancellationToken } from 'vscode'; import { logger } from '../../platform/logging'; import { IDeepnoteNotebookManager } from '../types'; import { DeepnoteProject, DeepnoteNotebook } from './deepnoteTypes'; @@ -25,12 +22,28 @@ export class DeepnoteInitNotebookRunner { * This should be called after the kernel is started but before user code executes. * @param notebook The notebook document * @param projectId The Deepnote project ID + * @param token Optional cancellation token to stop execution if notebook is closed */ - async runInitNotebookIfNeeded(notebook: NotebookDocument, projectId: string): Promise { + async runInitNotebookIfNeeded( + projectId: string, + notebook: NotebookDocument, + token?: CancellationToken + ): Promise { try { + // Check for cancellation before starting + if (token?.isCancellationRequested) { + logger.info(`Init notebook cancelled before start for project ${projectId}`); + return; + } + // Check if init notebook has already run for this project - if (this.notebookManager.hasInitNotebookRun(projectId)) { - logger.info(`Init notebook already run for project ${projectId}, skipping`); + if (this.notebookManager.hasInitNotebookBeenRun(projectId)) { + logger.info(`Init notebook already ran for project ${projectId}, skipping`); + return; + } + + if (token?.isCancellationRequested) { + logger.info(`Init notebook cancelled for project ${projectId}`); return; } @@ -60,10 +73,15 @@ export class DeepnoteInitNotebookRunner { return; } + if (token?.isCancellationRequested) { + logger.info(`Init notebook cancelled before execution for project ${projectId}`); + return; + } + logger.info(`Running init notebook "${initNotebook.name}" (${initNotebookId}) for project ${projectId}`); // Execute the init notebook with progress - const success = await this.executeInitNotebook(notebook, initNotebook); + const success = await this.executeInitNotebook(notebook, initNotebook, token); if (success) { // Mark as run so we don't run it again @@ -73,6 +91,11 @@ export class DeepnoteInitNotebookRunner { logger.warn(`Init notebook did not execute for project ${projectId} - kernel not available`); } } catch (error) { + // Check if this is a cancellation error + if (error instanceof Error && error.message === 'Cancelled') { + logger.info(`Init notebook cancelled for project ${projectId}`); + return; + } // Log error but don't throw - we want to let user continue anyway logger.error(`Error running init notebook for project ${projectId}:`, error); // Still mark as run to avoid retrying on every notebook open @@ -84,12 +107,28 @@ export class DeepnoteInitNotebookRunner { * Executes the init notebook's code blocks in the kernel. * @param notebook The notebook document (for kernel context) * @param initNotebook The init notebook to execute + * @param token Optional cancellation token from parent operation * @returns True if execution completed, false if kernel was not available */ - private async executeInitNotebook(notebook: NotebookDocument, initNotebook: DeepnoteNotebook): Promise { + private async executeInitNotebook( + notebook: NotebookDocument, + initNotebook: DeepnoteNotebook, + token?: CancellationToken + ): Promise { + // Check for cancellation before starting + if (token?.isCancellationRequested) { + logger.info(`Init notebook execution cancelled before start`); + return false; + } + // Show progress in both notification AND window for maximum visibility const cancellationTokenSource = new CancellationTokenSource(); + // Link parent token to our local token if provided + const tokenDisposable = token?.onCancellationRequested(() => { + cancellationTokenSource.cancel(); + }); + // Create a wrapper that reports to both progress locations const executeWithDualProgress = async () => { return window.withProgress( @@ -127,6 +166,7 @@ export class DeepnoteInitNotebookRunner { try { return await executeWithDualProgress(); } finally { + tokenDisposable?.dispose(); cancellationTokenSource.dispose(); } } @@ -135,9 +175,15 @@ export class DeepnoteInitNotebookRunner { notebook: NotebookDocument, initNotebook: DeepnoteNotebook, progress: (message: string, increment: number) => void, - _token: unknown + token: CancellationToken ): Promise { try { + // Check for cancellation + if (token.isCancellationRequested) { + logger.info(`Init notebook execution cancelled`); + return false; + } + progress(`Running init notebook "${initNotebook.name}"...`, 0); // Get the kernel for this notebook @@ -170,11 +216,23 @@ export class DeepnoteInitNotebookRunner { 5 ); + // Check for cancellation + if (token.isCancellationRequested) { + logger.info(`Init notebook execution cancelled before starting blocks`); + return false; + } + // Get kernel execution const kernelExecution = this.kernelProvider.getKernelExecution(kernel); // Execute each code block sequentially for (let i = 0; i < codeBlocks.length; i++) { + // Check for cancellation between blocks + if (token.isCancellationRequested) { + logger.info(`Init notebook execution cancelled after block ${i}`); + return false; + } + const block = codeBlocks[i]; const percentComplete = Math.floor((i / codeBlocks.length) * 100); @@ -224,5 +282,5 @@ export class DeepnoteInitNotebookRunner { export const IDeepnoteInitNotebookRunner = Symbol('IDeepnoteInitNotebookRunner'); export interface IDeepnoteInitNotebookRunner { - runInitNotebookIfNeeded(notebook: NotebookDocument, projectId: string): Promise; + runInitNotebookIfNeeded(projectId: string, notebook: NotebookDocument, token?: CancellationToken): Promise; } diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 1b5c9f4e88..e6912cf3bc 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -10,7 +10,9 @@ import { window, ProgressLocation, notebooks, - NotebookController + NotebookController, + CancellationTokenSource, + Disposable } from 'vscode'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { IDisposableRegistry } from '../../platform/common/types'; @@ -215,12 +217,35 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Remove from pending list this.projectsPendingInitNotebook.delete(projectId); - // Run init notebook now that kernel is available + // Create a CancellationTokenSource tied to the notebook lifecycle + const cts = new CancellationTokenSource(); + const disposables: Disposable[] = []; + try { - await this.initNotebookRunner.runInitNotebookIfNeeded(notebook, projectId); + // Register handler to cancel the token if the notebook is closed + // Note: We check the URI to ensure we only cancel for the specific notebook that closed + const closeListener = workspace.onDidCloseNotebookDocument((closedNotebook) => { + if (closedNotebook.uri.toString() === notebook.uri.toString()) { + logger.info(`Notebook closed while init notebook was running, cancelling for project ${projectId}`); + cts.cancel(); + } + }); + disposables.push(closeListener); + + // Run init notebook with cancellation support + await this.initNotebookRunner.runInitNotebookIfNeeded(projectId, notebook, cts.token); } catch (error) { + // Check if this is a cancellation error - if so, just log and continue + if (error instanceof Error && error.message === 'Cancelled') { + logger.info(`Init notebook cancelled for project ${projectId}`); + return; + } logger.error(`Error running init notebook: ${error}`); // Continue anyway - don't block user if init fails + } finally { + // Always clean up the CTS and event listeners + cts.dispose(); + disposables.forEach((d) => d.dispose()); } } @@ -433,6 +458,35 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Store the controller for reuse this.notebookControllers.set(notebookKey, controller); + // Prepare init notebook execution for when kernel starts + // This MUST complete before marking controller as preferred to avoid race conditions + const projectId = notebook.metadata?.deepnoteProjectId; + if (projectId) { + // Get the project and create requirements.txt before kernel starts + const project = this.notebookManager.getOriginalProject(projectId) as + | DeepnoteProject + | undefined; + if (project) { + // Create requirements.txt first (needs to be ready for init notebook) + progress.report({ message: 'Creating requirements.txt...' }); + await this.requirementsHelper.createRequirementsFile(project, progressToken); + logger.info(`Created requirements.txt for project ${projectId}`); + + // Check if project has an init notebook + if ( + project.project.initNotebookId && + !this.notebookManager.hasInitNotebookBeenRun(projectId) + ) { + // Store for execution when kernel actually starts + // Kernels are created lazily when cells execute, so we can't run init notebook now + this.projectsPendingInitNotebook.set(projectId, { notebook, project }); + logger.info( + `Init notebook will run automatically when kernel starts for project ${projectId}` + ); + } + } + } + // Mark this controller as protected so it won't be automatically disposed // This is similar to how active interpreter controllers are protected this.controllerRegistration.trackActiveInterpreterControllers(controllers); @@ -448,6 +502,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Auto-select the controller for this notebook using affinity // Setting NotebookControllerAffinity.Preferred will make VSCode automatically select this controller + // This is done AFTER requirements.txt creation to avoid race conditions controller.controller.updateNotebookAffinity(notebook, NotebookControllerAffinity.Preferred); logger.info(`Successfully auto-selected Deepnote kernel for ${getDisplayPath(notebook.uri)}`); @@ -460,30 +515,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, this.loadingControllers.delete(notebookKey); logger.info(`Disposed loading controller for ${notebookKey}`); } - - // Prepare init notebook execution for when kernel starts - const projectId = notebook.metadata?.deepnoteProjectId; - if (projectId) { - // Get the project and create requirements.txt before kernel starts - const project = this.notebookManager.getOriginalProject(projectId) as - | DeepnoteProject - | undefined; - if (project) { - // Create requirements.txt first (needs to be ready for init notebook) - progress.report({ message: 'Creating requirements.txt...' }); - await this.requirementsHelper.createRequirementsFile(project, progressToken); - - // Check if project has an init notebook - if (project.project.initNotebookId && !this.notebookManager.hasInitNotebookRun(projectId)) { - // Store for execution when kernel actually starts - // Kernels are created lazily when cells execute, so we can't run init notebook now - this.projectsPendingInitNotebook.set(projectId, { notebook, project }); - logger.info( - `Init notebook will run automatically when kernel starts for project ${projectId}` - ); - } - } - } } catch (ex) { logger.error(`Failed to auto-select Deepnote kernel: ${ex}`); throw ex; diff --git a/src/notebooks/deepnote/deepnoteNotebookManager.ts b/src/notebooks/deepnote/deepnoteNotebookManager.ts index be69780716..3c324fb663 100644 --- a/src/notebooks/deepnote/deepnoteNotebookManager.ts +++ b/src/notebooks/deepnote/deepnoteNotebookManager.ts @@ -80,7 +80,7 @@ export class DeepnoteNotebookManager implements IDeepnoteNotebookManager { * @param projectId Project identifier * @returns True if init notebook has been run, false otherwise */ - hasInitNotebookRun(projectId: string): boolean { + hasInitNotebookBeenRun(projectId: string): boolean { return this.projectsWithInitNotebookRun.has(projectId); } diff --git a/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts b/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts index a79124aca5..b01c6fa53d 100644 --- a/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts +++ b/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts @@ -1,6 +1,3 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. - import { inject, injectable } from 'inversify'; import { workspace, CancellationToken } from 'vscode'; import * as fs from 'fs'; @@ -27,7 +24,7 @@ export class DeepnoteRequirementsHelper { return; } - const requirements = project.project.settings.requirements; + const requirements = project.project.settings?.requirements; if (!requirements || !Array.isArray(requirements) || requirements.length === 0) { this.logger.info(`No requirements found in project ${project.project.id}`); return; diff --git a/src/notebooks/types.ts b/src/notebooks/types.ts index 6e37b30d8e..377c27988a 100644 --- a/src/notebooks/types.ts +++ b/src/notebooks/types.ts @@ -32,6 +32,6 @@ export interface IDeepnoteNotebookManager { selectNotebookForProject(projectId: string, notebookId: string): void; storeOriginalProject(projectId: string, project: unknown, notebookId: string): void; updateCurrentNotebookId(projectId: string, notebookId: string): void; - hasInitNotebookRun(projectId: string): boolean; + hasInitNotebookBeenRun(projectId: string): boolean; markInitNotebookAsRun(projectId: string): void; } From 8a9c59b652d92062121c47298d319b7064a6bd1c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Filip=20P=C3=BDrek?= Date: Wed, 8 Oct 2025 17:45:28 +0200 Subject: [PATCH 4/4] feat: ask about overriding requirements.txt first --- .../deepnoteKernelAutoSelector.node.ts | 44 +++++------ .../deepnoteRequirementsHelper.node.ts | 79 ++++++++++++++++++- 2 files changed, 97 insertions(+), 26 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index e6912cf3bc..0d97a2162b 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -461,29 +461,27 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Prepare init notebook execution for when kernel starts // This MUST complete before marking controller as preferred to avoid race conditions const projectId = notebook.metadata?.deepnoteProjectId; - if (projectId) { - // Get the project and create requirements.txt before kernel starts - const project = this.notebookManager.getOriginalProject(projectId) as - | DeepnoteProject - | undefined; - if (project) { - // Create requirements.txt first (needs to be ready for init notebook) - progress.report({ message: 'Creating requirements.txt...' }); - await this.requirementsHelper.createRequirementsFile(project, progressToken); - logger.info(`Created requirements.txt for project ${projectId}`); - - // Check if project has an init notebook - if ( - project.project.initNotebookId && - !this.notebookManager.hasInitNotebookBeenRun(projectId) - ) { - // Store for execution when kernel actually starts - // Kernels are created lazily when cells execute, so we can't run init notebook now - this.projectsPendingInitNotebook.set(projectId, { notebook, project }); - logger.info( - `Init notebook will run automatically when kernel starts for project ${projectId}` - ); - } + const project = projectId + ? (this.notebookManager.getOriginalProject(projectId) as DeepnoteProject | undefined) + : undefined; + + if (project) { + // Create requirements.txt first (needs to be ready for init notebook) + progress.report({ message: 'Creating requirements.txt...' }); + await this.requirementsHelper.createRequirementsFile(project, progressToken); + logger.info(`Created requirements.txt for project ${projectId}`); + + // Check if project has an init notebook that hasn't been run yet + if ( + project.project.initNotebookId && + !this.notebookManager.hasInitNotebookBeenRun(projectId!) + ) { + // Store for execution when kernel actually starts + // Kernels are created lazily when cells execute, so we can't run init notebook now + this.projectsPendingInitNotebook.set(projectId!, { notebook, project }); + logger.info( + `Init notebook will run automatically when kernel starts for project ${projectId}` + ); } } diff --git a/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts b/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts index b01c6fa53d..4ea14115d9 100644 --- a/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts +++ b/src/notebooks/deepnote/deepnoteRequirementsHelper.node.ts @@ -1,16 +1,22 @@ import { inject, injectable } from 'inversify'; -import { workspace, CancellationToken } from 'vscode'; +import { workspace, CancellationToken, window } from 'vscode'; import * as fs from 'fs'; import * as path from '../../platform/vscode-path/path'; import type { DeepnoteProject } from './deepnoteTypes'; import { ILogger } from '../../platform/logging/types'; +import { IPersistentStateFactory } from '../../platform/common/types'; + +const DONT_ASK_OVERWRITE_REQUIREMENTS_KEY = 'DEEPNOTE_DONT_ASK_OVERWRITE_REQUIREMENTS'; /** * Helper class for creating requirements.txt files from Deepnote project settings. */ @injectable() export class DeepnoteRequirementsHelper { - constructor(@inject(ILogger) private readonly logger: ILogger) {} + constructor( + @inject(ILogger) private readonly logger: ILogger, + @inject(IPersistentStateFactory) private readonly persistentStateFactory: IPersistentStateFactory + ) {} /** * Extracts requirements from project settings and creates a local requirements.txt file. @@ -45,9 +51,76 @@ export class DeepnoteRequirementsHelper { const workspaceRoot = workspaceFolders[0].uri.fsPath; const requirementsPath = path.join(workspaceRoot, 'requirements.txt'); - // Convert requirements array to text format + // Convert requirements array to text format first const requirementsText = requirements.join('\n') + '\n'; + // Check if requirements.txt already exists + const fileExists = await fs.promises + .access(requirementsPath) + .then(() => true) + .catch(() => false); + + if (fileExists) { + // Read existing file contents and compare + const existingContent = await fs.promises.readFile(requirementsPath, 'utf8'); + + if (existingContent === requirementsText) { + this.logger.info('requirements.txt already has the correct content, skipping update'); + return; + } + + // File exists but content is different, check if we should prompt user + const dontAskState = this.persistentStateFactory.createGlobalPersistentState( + DONT_ASK_OVERWRITE_REQUIREMENTS_KEY, + false // default: ask user + ); + + if (!dontAskState.value) { + // User hasn't chosen "Don't Ask Again", so prompt them + const yes = 'Yes'; + const no = 'No'; + const dontAskAgain = "Don't Ask Again"; + + const response = await window.showWarningMessage( + `A requirements.txt file already exists in this workspace. Do you want to override it with requirements from your Deepnote project?`, + { modal: true }, + yes, + no, + dontAskAgain + ); + + // Check cancellation after showing the prompt + if (token.isCancellationRequested) { + return; + } + + switch (response) { + case yes: + // User wants to override, continue with writing + this.logger.info('User chose to override requirements.txt'); + break; + case no: + // User doesn't want to override + this.logger.info('User chose not to override requirements.txt'); + return; + case dontAskAgain: + // User chose "Don't Ask Again", save preference and override this time + await dontAskState.updateValue(true); + this.logger.info('User chose "Don\'t Ask Again" for requirements.txt override'); + break; + default: + // User dismissed the prompt (clicked X) + this.logger.info('User dismissed requirements.txt override prompt'); + return; + } + } else { + // User previously selected "Don't Ask Again", automatically override + this.logger.info( + 'Automatically overriding requirements.txt (user previously selected "Don\'t Ask Again")' + ); + } + } + // Write the requirements.txt file await fs.promises.writeFile(requirementsPath, requirementsText, 'utf8');