From ffe09588ac299d7ad88f852d707ee91eb5041181 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 9 Mar 2026 08:30:26 +0000 Subject: [PATCH 01/19] chore(dep): Update deepnote database integrations package --- package-lock.json | 14 +++++++------- package.json | 2 +- .../integrations/ConfigurationForm.tsx | 4 ++-- .../webview-side/integrations/TrinoForm.tsx | 17 +++++++++++++---- 4 files changed, 23 insertions(+), 14 deletions(-) diff --git a/package-lock.json b/package-lock.json index f05d4c434d..799d8f02a3 100644 --- a/package-lock.json +++ b/package-lock.json @@ -13,7 +13,7 @@ "@c4312/evt": "^0.1.1", "@deepnote/blocks": "^4.3.0", "@deepnote/convert": "^3.2.0", - "@deepnote/database-integrations": "^1.3.0", + "@deepnote/database-integrations": "^1.4.3", "@deepnote/sql-language-server": "^3.0.0", "@enonic/fnv-plus": "^1.3.0", "@jupyter-widgets/base": "^6.0.8", @@ -1954,9 +1954,9 @@ } }, "node_modules/@deepnote/database-integrations": { - "version": "1.3.0", - "resolved": "https://registry.npmjs.org/@deepnote/database-integrations/-/database-integrations-1.3.0.tgz", - "integrity": "sha512-Q08nyegvvrkZCbC/+hE7hxT+ISCx4ejHnx9D1/w9YW/cJ9iC7DubUR7vAR0DyiE6gBjfEco1xVBm/BMJRu/lqA==", + "version": "1.4.3", + "resolved": "https://registry.npmjs.org/@deepnote/database-integrations/-/database-integrations-1.4.3.tgz", + "integrity": "sha512-h12mkl4tX/0TSjF7wXq3e6YimxfcgvQzRjTr4eBg0kTbZizvhUjWEQw/9+JiQZCyNvanLaeg0LXVCRlp8LxqJQ==", "license": "Apache-2.0", "dependencies": { "zod": "3.25.76" @@ -32576,9 +32576,9 @@ } }, "@deepnote/database-integrations": { - "version": "1.3.0", - "resolved": "https://registry.npmjs.org/@deepnote/database-integrations/-/database-integrations-1.3.0.tgz", - "integrity": "sha512-Q08nyegvvrkZCbC/+hE7hxT+ISCx4ejHnx9D1/w9YW/cJ9iC7DubUR7vAR0DyiE6gBjfEco1xVBm/BMJRu/lqA==", + "version": "1.4.3", + "resolved": "https://registry.npmjs.org/@deepnote/database-integrations/-/database-integrations-1.4.3.tgz", + "integrity": "sha512-h12mkl4tX/0TSjF7wXq3e6YimxfcgvQzRjTr4eBg0kTbZizvhUjWEQw/9+JiQZCyNvanLaeg0LXVCRlp8LxqJQ==", "requires": { "zod": "3.25.76" }, diff --git a/package.json b/package.json index 5bfc1db66b..051493f10b 100644 --- a/package.json +++ b/package.json @@ -2675,7 +2675,7 @@ "@c4312/evt": "^0.1.1", "@deepnote/blocks": "^4.3.0", "@deepnote/convert": "^3.2.0", - "@deepnote/database-integrations": "^1.3.0", + "@deepnote/database-integrations": "^1.4.3", "@deepnote/sql-language-server": "^3.0.0", "@enonic/fnv-plus": "^1.3.0", "@jupyter-widgets/base": "^6.0.8", diff --git a/src/webviews/webview-side/integrations/ConfigurationForm.tsx b/src/webviews/webview-side/integrations/ConfigurationForm.tsx index dc79fd3d72..3c73ceb784 100644 --- a/src/webviews/webview-side/integrations/ConfigurationForm.tsx +++ b/src/webviews/webview-side/integrations/ConfigurationForm.tsx @@ -16,7 +16,7 @@ import { RedshiftForm } from './RedshiftForm'; import { SnowflakeForm } from './SnowflakeForm'; import { SpannerForm } from './SpannerForm'; import { SQLServerForm } from './SQLServerForm'; -import { TrinoForm } from './TrinoForm'; +import { isTrinoPasswordConfig, TrinoForm } from './TrinoForm'; import { ConfigurableDatabaseIntegrationConfig, ConfigurableDatabaseIntegrationType } from './types'; import { integrationTypeLabels } from './integrationUtils'; @@ -157,7 +157,7 @@ export const ConfigurationForm: React.FC = ({ return ( ; +export type TrinoPasswordConfig = TrinoConfig & { + metadata: Extract; +}; + +export function isTrinoPasswordConfig(config: TrinoConfig): config is TrinoPasswordConfig { + return config.metadata.authMethod !== 'trino-oauth'; +} + export interface ITrinoFormProps { integrationId: string; - existingConfig: Extract | null; + existingConfig: TrinoPasswordConfig | null; defaultName?: string; - onSave: (config: Extract) => void; + onSave: (config: TrinoPasswordConfig) => void; onCancel: () => void; } function createEmptyTrinoConfig(params: { id: string; name?: string; -}): Extract { +}): TrinoPasswordConfig { return { id: params.id, name: (params.name || getDefaultIntegrationName('trino')).trim(), @@ -38,7 +47,7 @@ export const TrinoForm: React.FC = ({ onSave, onCancel }) => { - const [pendingConfig, setPendingConfig] = React.useState>( + const [pendingConfig, setPendingConfig] = React.useState( existingConfig ? structuredClone(existingConfig) : createEmptyTrinoConfig({ id: integrationId, name: defaultName }) From 3bb72eb72378adb3fe1d95271ef3eba12d9187dc Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 9 Mar 2026 08:37:26 +0000 Subject: [PATCH 02/19] Reformat code --- .../webview-side/integrations/ConfigurationForm.tsx | 6 +++++- src/webviews/webview-side/integrations/TrinoForm.tsx | 5 +---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/webviews/webview-side/integrations/ConfigurationForm.tsx b/src/webviews/webview-side/integrations/ConfigurationForm.tsx index 3c73ceb784..a82699690d 100644 --- a/src/webviews/webview-side/integrations/ConfigurationForm.tsx +++ b/src/webviews/webview-side/integrations/ConfigurationForm.tsx @@ -157,7 +157,11 @@ export const ConfigurationForm: React.FC = ({ return ( void; } -function createEmptyTrinoConfig(params: { - id: string; - name?: string; -}): TrinoPasswordConfig { +function createEmptyTrinoConfig(params: { id: string; name?: string }): TrinoPasswordConfig { return { id: params.id, name: (params.name || getDefaultIntegrationName('trino')).trim(), From d5f67f62c90e42f6f57b028570cd77fbbaa09aea Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 10 Mar 2026 15:39:36 +0000 Subject: [PATCH 03/19] chore(runtime-core): Add deepnote runtime-core dependency, and use it instead of existing implementation --- package-lock.json | 55 ++ package.json | 1 + .../deepnote/deepnoteServerStarter.node.ts | 643 +++++------------- .../deepnoteServerStarter.unit.test.ts | 558 ++------------- .../deepnote/deepnoteToolkitInstaller.node.ts | 21 +- src/kernels/deepnote/types.ts | 3 + 6 files changed, 305 insertions(+), 976 deletions(-) diff --git a/package-lock.json b/package-lock.json index 8f63cf54b3..5b6043d0a5 100644 --- a/package-lock.json +++ b/package-lock.json @@ -14,6 +14,7 @@ "@deepnote/blocks": "^4.3.0", "@deepnote/convert": "^3.2.0", "@deepnote/database-integrations": "^1.4.3", + "@deepnote/runtime-core": "^0.2.0", "@deepnote/sql-language-server": "^3.0.0", "@enonic/fnv-plus": "^1.3.0", "@jupyter-widgets/base": "^6.0.8", @@ -1971,6 +1972,40 @@ "url": "https://github.com/sponsors/colinhacks" } }, + "node_modules/@deepnote/runtime-core": { + "version": "0.2.0", + "resolved": "https://registry.npmjs.org/@deepnote/runtime-core/-/runtime-core-0.2.0.tgz", + "integrity": "sha512-wIgUOSROSyFpfFd+Mx/9GA3mHdyJ7aIqs4bejS0SUr5ogC+wo1xj+ZfwfEzMQRse9M8f5SKn8qj6zjnykKRTJg==", + "license": "Apache-2.0", + "dependencies": { + "@deepnote/blocks": "4.3.0", + "@jupyterlab/nbformat": "^4.3.2", + "@jupyterlab/services": "^7.3.2", + "tcp-port-used": "^1.0.2", + "ws": "^8.18.0" + } + }, + "node_modules/@deepnote/runtime-core/node_modules/ws": { + "version": "8.19.0", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.19.0.tgz", + "integrity": "sha512-blAT2mjOEIi0ZzruJfIhb3nps74PRWTCz1IjglWEEpQl5XS/UNama6u2/rjFkDDouqr4L67ry+1aGIALViWjDg==", + "license": "MIT", + "engines": { + "node": ">=10.0.0" + }, + "peerDependencies": { + "bufferutil": "^4.0.1", + "utf-8-validate": ">=5.0.2" + }, + "peerDependenciesMeta": { + "bufferutil": { + "optional": true + }, + "utf-8-validate": { + "optional": true + } + } + }, "node_modules/@deepnote/sql-language-server": { "version": "3.0.0", "resolved": "https://registry.npmjs.org/@deepnote/sql-language-server/-/sql-language-server-3.0.0.tgz", @@ -32590,6 +32625,26 @@ } } }, + "@deepnote/runtime-core": { + "version": "0.2.0", + "resolved": "https://registry.npmjs.org/@deepnote/runtime-core/-/runtime-core-0.2.0.tgz", + "integrity": "sha512-wIgUOSROSyFpfFd+Mx/9GA3mHdyJ7aIqs4bejS0SUr5ogC+wo1xj+ZfwfEzMQRse9M8f5SKn8qj6zjnykKRTJg==", + "requires": { + "@deepnote/blocks": "4.3.0", + "@jupyterlab/nbformat": "^4.3.2", + "@jupyterlab/services": "^7.3.2", + "tcp-port-used": "^1.0.2", + "ws": "^8.18.0" + }, + "dependencies": { + "ws": { + "version": "8.19.0", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.19.0.tgz", + "integrity": "sha512-blAT2mjOEIi0ZzruJfIhb3nps74PRWTCz1IjglWEEpQl5XS/UNama6u2/rjFkDDouqr4L67ry+1aGIALViWjDg==", + "requires": {} + } + } + }, "@deepnote/sql-language-server": { "version": "3.0.0", "resolved": "https://registry.npmjs.org/@deepnote/sql-language-server/-/sql-language-server-3.0.0.tgz", diff --git a/package.json b/package.json index 1c78265433..46f57c49d4 100644 --- a/package.json +++ b/package.json @@ -2676,6 +2676,7 @@ "@deepnote/blocks": "^4.3.0", "@deepnote/convert": "^3.2.0", "@deepnote/database-integrations": "^1.4.3", + "@deepnote/runtime-core": "^0.2.0", "@deepnote/sql-language-server": "^3.0.0", "@enonic/fnv-plus": "^1.3.0", "@jupyter-widgets/base": "^6.0.8", diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index fe41270e0f..22bd170310 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -1,29 +1,36 @@ -// Copyright (c) Microsoft Corporation. -// Licensed under the MIT License. +/** + * @deepnote/runtime-core functions not currently exported that would be useful: + * - findConsecutiveAvailablePorts(startPort) — duplicated logic for multi-server port reservation + * - waitForServer(info, timeoutMs) — health-check polling on /api + * - createJsonWebSocketFactory() — forces JSON-only Jupyter WS protocol, potential stability improvement + * - ExecutionEngine.toPythonLiteral(value) — JS-to-Python literal conversion + */ import * as fs from 'fs-extra'; import { inject, injectable, named, optional } from 'inversify'; import * as os from 'os'; import { CancellationToken, l10n, Uri } from 'vscode'; + +import { startServer, stopServer, type ServerInfo as RuntimeCoreServerInfo } from '@deepnote/runtime-core'; + import { IExtensionSyncActivationService } from '../../platform/activation/types'; -import { Cancellation, raceCancellationError } from '../../platform/common/cancellation'; +import { Cancellation } from '../../platform/common/cancellation'; import { STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants'; -import { IProcessServiceFactory, ObservableExecutionResult } from '../../platform/common/process/types.node'; -import { IAsyncDisposableRegistry, IDisposable, IHttpClient, IOutputChannel } from '../../platform/common/types'; +import { IProcessServiceFactory } from '../../platform/common/process/types.node'; +import { IAsyncDisposableRegistry, IDisposable, IOutputChannel } from '../../platform/common/types'; import { sleep } from '../../platform/common/utils/async'; import { generateUuid } from '../../platform/common/uuid'; -import { DeepnoteServerStartupError, DeepnoteServerTimeoutError } from '../../platform/errors/deepnoteKernelErrors'; +import { DeepnoteServerStartupError } from '../../platform/errors/deepnoteKernelErrors'; import { logger } from '../../platform/logging'; import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import * as path from '../../platform/vscode-path/path'; -import { DEEPNOTE_DEFAULT_PORT, DeepnoteServerInfo, IDeepnoteServerStarter, IDeepnoteToolkitInstaller } from './types'; +import { DeepnoteServerInfo, IDeepnoteServerStarter, IDeepnoteToolkitInstaller } from './types'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; -import tcpPortUsed from 'tcp-port-used'; -/** - * Lock file data structure for tracking server ownership - */ +const SERVER_STARTUP_TIMEOUT_MS = 120_000; +const GRACEFUL_SHUTDOWN_TIMEOUT_MS = 3000; + interface ServerLockFile { sessionId: string; pid: number; @@ -42,65 +49,53 @@ type PendingOperation = interface ProjectContext { environmentId: string; - serverProcess: ObservableExecutionResult | null; + runtimeCoreServerInfo: RuntimeCoreServerInfo | null; serverInfo: DeepnoteServerInfo | null; } /** * Starts and manages the deepnote-toolkit Jupyter server. + * + * Uses @deepnote/runtime-core's `startServer`/`stopServer` for the core server + * lifecycle (process spawn, port discovery, health checks, shutdown), and layers + * extension-specific concerns on top: lock files, orphan cleanup, SQL integration + * env vars, output channel logging, and multi-server concurrency control. */ @injectable() export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtensionSyncActivationService { - private readonly serverProcesses: Map> = new Map(); - private readonly serverInfos: Map = new Map(); private readonly disposablesByFile: Map = new Map(); private readonly projectContexts: Map = new Map(); - // Track in-flight operations per file to prevent concurrent start/stop private readonly pendingOperations: Map = new Map(); - // Global lock for port allocation to prevent race conditions when multiple environments start concurrently private portAllocationLock: Promise = Promise.resolve(); - // Unique session ID for this VS Code window instance private readonly sessionId: string = generateUuid(); - // Directory for lock files private readonly lockFileDir: string = path.join(os.tmpdir(), 'vscode-deepnote-locks'); - // Track server output for error reporting - private readonly serverOutputByFile: Map = new Map(); constructor( @inject(IProcessServiceFactory) private readonly processServiceFactory: IProcessServiceFactory, @inject(IDeepnoteToolkitInstaller) private readonly toolkitInstaller: IDeepnoteToolkitInstaller, @inject(DeepnoteAgentSkillsManager) private readonly agentSkillsManager: DeepnoteAgentSkillsManager, @inject(IOutputChannel) @named(STANDARD_OUTPUT_CHANNEL) private readonly outputChannel: IOutputChannel, - @inject(IHttpClient) private readonly httpClient: IHttpClient, @inject(IAsyncDisposableRegistry) asyncRegistry: IAsyncDisposableRegistry, @inject(ISqlIntegrationEnvVarsProvider) @optional() private readonly sqlIntegrationEnvVars?: ISqlIntegrationEnvVarsProvider ) { - // Register for disposal when the extension deactivates asyncRegistry.push(this); } public activate(): void { - // Ensure lock file directory exists this.initializeLockFileDirectory().catch((ex) => { logger.warn('Failed to initialize lock file directory', ex); }); - // Clean up any orphaned deepnote-toolkit processes from previous sessions this.cleanupOrphanedProcesses().catch((ex) => { logger.warn('Failed to cleanup orphaned processes', ex); }); } /** - * Environment-based method: Start a server for a kernel environment. - * @param interpreter The Python interpreter to use - * @param venvPath The path to the venv - * @param managedVenv Whether the venv is managed by this extension (created by us) - * @param environmentId The environment ID (used as key for server management) - * @param token Cancellation token - * @returns Server connection information + * Start a server for a kernel environment. + * Serializes concurrent operations on the same environment to prevent race conditions. */ public async startServer( interpreter: PythonEnvironment, @@ -114,7 +109,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const fileKey = deepnoteFileUri.fsPath; const serverKey = `${fileKey}-${environmentId}`; - // Wait for any pending operations on this environment to complete let pendingOp = this.pendingOperations.get(serverKey); if (pendingOp) { logger.info(`Waiting for pending operation on ${serverKey} to complete...`); @@ -135,34 +129,28 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension return existingServerInfo; } - // Start the operation if not already pending pendingOp = this.pendingOperations.get(serverKey); if (pendingOp && pendingOp.type === 'start') { - // TODO - check pending operation environment id ? return await pendingOp.promise; } } else { - // Stop the existing server logger.info( `Stopping existing server for ${serverKey} with environmentId ${existingEnvironmentId} to start new one with environmentId ${environmentId}...` ); await this.stopServerForEnvironment(existingContext, deepnoteFileUri, token); - // TODO - Clear controllers for the notebook ? } } else { - const newContext = { + const newContext: ProjectContext = { environmentId, - serverProcess: null, + runtimeCoreServerInfo: null, serverInfo: null }; this.projectContexts.set(serverKey, newContext); - existingContext = newContext; } - // Start the operation and track it const operation = { type: 'start' as const, promise: this.startServerForEnvironment( @@ -181,11 +169,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension try { const result = await operation.promise; - // Update context with running server info existingContext.serverInfo = result; return result; } finally { - // Remove from pending operations when done if (this.pendingOperations.get(serverKey) === operation) { this.pendingOperations.delete(serverKey); } @@ -193,17 +179,14 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } /** - * Environment-based method: Stop the server for a kernel environment. - * @param environmentId The environment ID + * Stop the deepnote-toolkit server for a kernel environment. */ - // public async stopServer(environmentId: string, token?: CancellationToken): Promise { public async stopServer(deepnoteFileUri: Uri, token?: CancellationToken): Promise { Cancellation.throwIfCanceled(token); const fileKey = deepnoteFileUri.fsPath; const projectContext = this.projectContexts.get(fileKey) ?? null; - // Wait for any pending operations on this environment to complete const pendingOp = this.pendingOperations.get(fileKey); if (pendingOp) { logger.info(`Waiting for pending operation on ${fileKey} before stopping...`); @@ -216,7 +199,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); - // Start the stop operation and track it const operation = { type: 'stop' as const, promise: this.stopServerForEnvironment(projectContext, deepnoteFileUri, token) @@ -226,7 +208,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension try { await operation.promise; } finally { - // Remove from pending operations when done if (this.pendingOperations.get(fileKey) === operation) { this.pendingOperations.delete(fileKey); } @@ -234,7 +215,14 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } /** - * Environment-based server start implementation. + * Core server start using @deepnote/runtime-core's `startServer`. + * + * Extension-specific layers: + * - Toolkit/venv installation (before start) + * - SQL integration env var injection (via ServerOptions.env) + * - Lock file creation (after start, using returned PID) + * - Output channel logging (via process stdout/stderr streams) + * - Port allocation serialization across concurrent starts */ private async startServerForEnvironment( projectContext: ProjectContext, @@ -251,7 +239,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); - // Ensure toolkit is installed in venv and get venv's Python interpreter logger.info(`Ensuring deepnote-toolkit is installed in venv for environment ${environmentId}...`); const { pythonInterpreter: venvInterpreter } = await this.toolkitInstaller.ensureVenvAndToolkit( interpreter, @@ -268,171 +255,66 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); - // Allocate both ports with global lock to prevent race conditions - // Note: allocatePorts reserves both ports immediately in serverInfos - // const { jupyterPort, lspPort } = await this.allocatePorts(environmentId); - const { jupyterPort, lspPort } = await this.allocatePorts(serverKey); + // Serialize port allocation across concurrent server starts + const port = await this.reserveStartPort(serverKey); logger.info( - `Starting deepnote-toolkit server on jupyter port ${jupyterPort} and lsp port ${lspPort} for ${serverKey} with environmentId ${environmentId}` + `Starting deepnote-toolkit server on port ${port} for ${serverKey} with environmentId ${environmentId}` ); - this.outputChannel.appendLine( - l10n.t('Starting Deepnote server on jupyter port {0} and lsp port {1}...', jupyterPort, lspPort) - ); - - // Start the server with venv's Python in PATH - const processService = await this.processServiceFactory.create(undefined); - - // Set up environment to ensure the venv's Python is used for shell commands - const venvBinDir = path.dirname(venvInterpreter.uri.fsPath); - const env = { ...process.env }; - - // Prepend venv bin directory to PATH so shell commands use venv's Python - env.PATH = `${venvBinDir}${process.platform === 'win32' ? ';' : ':'}${env.PATH || ''}`; + this.outputChannel.appendLine(l10n.t('Starting Deepnote server on port {0}...', port)); - // Also set VIRTUAL_ENV to indicate we're in a venv - env.VIRTUAL_ENV = venvPath.fsPath; + // Gather SQL integration env vars to pass to the server + const extraEnv = await this.gatherSqlIntegrationEnvVars(deepnoteFileUri, environmentId, token); - // Enforce published pip constraints to prevent breaking Deepnote Toolkit's dependencies - env.DEEPNOTE_ENFORCE_PIP_CONSTRAINTS = 'true'; - - // Detached mode - env.DEEPNOTE_RUNTIME__RUNNING_IN_DETACHED_MODE = 'true'; - - // Detached mode ensures no requests are made to the backend (directly, or via proxy) - // as there is no backend running in the extension, therefore: - // 1. integration environment variables are injected here instead - // 2. post start hooks won't work / are not executed - env.DEEPNOTE_RUNTIME__RUNNING_IN_DETACHED_MODE = 'true'; - - // Inject SQL integration environment variables - if (this.sqlIntegrationEnvVars) { - logger.debug( - `DeepnoteServerStarter: Injecting SQL integration env vars for ${fileKey} with environmentId ${environmentId}` + let runtimeCoreInfo: RuntimeCoreServerInfo; + try { + runtimeCoreInfo = await startServer({ + pythonEnv: venvPath.fsPath, + workingDirectory: path.dirname(deepnoteFileUri.fsPath), + port, + startupTimeoutMs: SERVER_STARTUP_TIMEOUT_MS, + env: extraEnv + }); + } catch (error) { + throw new DeepnoteServerStartupError( + interpreter.uri.fsPath, + port, + 'unknown', + '', + error instanceof Error ? error.message : String(error), + error instanceof Error ? error : undefined ); - try { - const sqlEnvVars = await this.sqlIntegrationEnvVars.getEnvironmentVariables(deepnoteFileUri, token); - // const sqlEnvVars = {}; // TODO: update how environment variables are retrieved - if (sqlEnvVars && Object.keys(sqlEnvVars).length > 0) { - logger.debug(`DeepnoteServerStarter: Injecting ${Object.keys(sqlEnvVars).length} SQL env vars`); - Object.assign(env, sqlEnvVars); - } else { - logger.debug('DeepnoteServerStarter: No SQL integration env vars to inject'); - } - } catch (error) { - logger.error('DeepnoteServerStarter: Failed to get SQL integration env vars', error.message); - } - } else { - logger.debug('DeepnoteServerStarter: SqlIntegrationEnvironmentVariablesProvider not available'); } - // Remove PYTHONHOME if it exists (can interfere with venv) - delete env.PYTHONHOME; - - const serverProcess = processService.execObservable( - venvInterpreter.uri.fsPath, - [ - '-m', - 'deepnote_toolkit', - 'server', - '--jupyter-port', - jupyterPort.toString(), - '--ls-port', - lspPort.toString() - ], - { env, cwd: path.dirname(deepnoteFileUri.fsPath) } - ); - - projectContext.serverProcess = serverProcess; - - this.serverProcesses.set(serverKey, serverProcess); - - // Track disposables for this environment - const disposables: IDisposable[] = []; - this.disposablesByFile.set(serverKey, disposables); + projectContext.runtimeCoreServerInfo = runtimeCoreInfo; - // Initialize output tracking for error reporting - this.serverOutputByFile.set(serverKey, { stdout: '', stderr: '' }); - - // Monitor server output - serverProcess.out.onDidChange( - (output) => { - const outputTracking = this.serverOutputByFile.get(serverKey); - if (output.source === 'stdout') { - logger.trace(`Deepnote server (${serverKey}): ${output.out}`); - this.outputChannel.appendLine(output.out); - if (outputTracking) { - // Keep last 5000 characters of output for error reporting - outputTracking.stdout = (outputTracking.stdout + output.out).slice(-5000); - } - } else if (output.source === 'stderr') { - logger.warn(`Deepnote server stderr (${serverKey}): ${output.out}`); - this.outputChannel.appendLine(output.out); - if (outputTracking) { - // Keep last 5000 characters of error output for error reporting - outputTracking.stderr = (outputTracking.stderr + output.out).slice(-5000); - } - } - }, - this, - disposables - ); + const serverInfo: DeepnoteServerInfo = { + url: runtimeCoreInfo.url, + jupyterPort: runtimeCoreInfo.jupyterPort, + lspPort: runtimeCoreInfo.lspPort, + process: runtimeCoreInfo.process + }; - // Wait for server to be ready - const url = `http://localhost:${jupyterPort}`; - const serverInfo = { url, jupyterPort, lspPort }; - this.serverInfos.set(serverKey, serverInfo); + // Set up output channel logging from the server process + this.monitorServerOutput(serverKey, runtimeCoreInfo); - // Write lock file for the server process - const serverPid = serverProcess.proc?.pid; + // Write lock file for orphan-cleanup tracking + const serverPid = runtimeCoreInfo.process.pid; if (serverPid) { await this.writeLockFile(serverPid); } else { logger.warn(`Could not get PID for server process for ${serverKey}`); } - try { - const serverReady = await this.waitForServer(serverInfo, 120000, token); - if (!serverReady) { - const output = this.serverOutputByFile.get(serverKey); - - throw new DeepnoteServerTimeoutError(serverInfo.url, 120000, output?.stderr || undefined); - } - } catch (error) { - if (error instanceof DeepnoteServerTimeoutError || error instanceof DeepnoteServerStartupError) { - // await this.stopServerImpl(deepnoteFileUri); - await this.stopServerForEnvironment(projectContext, deepnoteFileUri); - throw error; - } - - // Capture output BEFORE cleaning up (stopServerImpl deletes it) - const output = this.serverOutputByFile.get(serverKey); - const capturedStdout = output?.stdout || ''; - const capturedStderr = output?.stderr || ''; - - // Clean up leaked server before rethrowing - await this.stopServerForEnvironment(projectContext, deepnoteFileUri); - - throw new DeepnoteServerStartupError( - interpreter.uri.fsPath, - serverInfo.jupyterPort, - 'unknown', - capturedStdout, - capturedStderr, - error instanceof Error ? error : undefined - ); - } - - logger.info(`Deepnote server started successfully at ${url} for ${serverKey}`); - this.outputChannel.appendLine(l10n.t('✓ Deepnote server running at {0}', url)); + logger.info(`Deepnote server started successfully at ${runtimeCoreInfo.url} for ${serverKey}`); + this.outputChannel.appendLine(l10n.t('✓ Deepnote server running at {0}', runtimeCoreInfo.url)); return serverInfo; } /** - * Environment-based server stop implementation. + * Stop the server using @deepnote/runtime-core's `stopServer` (SIGTERM -> wait -> SIGKILL). */ - // private async stopServerForEnvironment(environmentId: string, token?: CancellationToken): Promise { private async stopServerForEnvironment( projectContext: ProjectContext | null, deepnoteFileUri: Uri, @@ -442,23 +324,23 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); - // const serverProcess = this.serverProcesses.get(fileKey); - const serverProcess = projectContext?.serverProcess; + const runtimeCoreInfo = projectContext?.runtimeCoreServerInfo; - if (serverProcess) { - const serverPid = serverProcess.proc?.pid; + if (runtimeCoreInfo) { + const serverPid = runtimeCoreInfo.process.pid; try { logger.info(`Stopping Deepnote server for ${fileKey}...`); - serverProcess.proc?.kill(); - this.serverProcesses.delete(fileKey); - this.serverInfos.delete(fileKey); - this.serverOutputByFile.delete(fileKey); + await stopServer(runtimeCoreInfo); this.outputChannel.appendLine(l10n.t('Deepnote server stopped for {0}', fileKey)); } catch (ex) { logger.error('Error stopping Deepnote server', ex); } finally { - // Clean up lock file after stopping the server + if (projectContext) { + projectContext.runtimeCoreServerInfo = null; + projectContext.serverInfo = null; + } + if (serverPid) { await this.deleteLockFile(serverPid); } @@ -474,41 +356,26 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } - private async waitForServer( - serverInfo: DeepnoteServerInfo, - timeout: number, - token?: CancellationToken - ): Promise { - const startTime = Date.now(); - while (Date.now() - startTime < timeout) { - Cancellation.throwIfCanceled(token); - if (await this.isServerRunning(serverInfo)) { - return true; - } - await raceCancellationError(token, sleep(500)); - } - return false; - } - + /** + * Check if a server is still running by probing its /api endpoint. + */ private async isServerRunning(serverInfo: DeepnoteServerInfo): Promise { try { - // Try to connect to the Jupyter API endpoint - const exists = await this.httpClient.exists(`${serverInfo.url}/api`).catch(() => false); - return exists; + const response = await fetch(`${serverInfo.url}/api`); + return response.ok; } catch { return false; } } /** - * Allocate both Jupyter and LSP ports atomically with global serialization. - * When multiple environments start simultaneously, this ensures each gets unique ports. + * Serialize port reservation across concurrent server starts. * - * @param key The environment ID to reserve ports for - * @returns Object with jupyterPort and lspPort + * runtime-core's `startServer` finds its own consecutive ports, but when multiple + * servers start concurrently in the extension, they can race. This lock serializes + * the starts so each `startServer` call sees the ports bound by previous calls. */ - private async allocatePorts(key: string): Promise<{ jupyterPort: number; lspPort: number }> { - // Chain onto the existing lock promise to serialize allocations even when multiple calls start concurrently + private async reserveStartPort(serverKey: string): Promise { const previousLock = this.portAllocationLock; let releaseLock: () => void; const currentLock = new Promise((resolve) => { @@ -516,239 +383,135 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension }); this.portAllocationLock = previousLock.then(() => currentLock); - // Wait until all prior allocations have completed before proceeding await previousLock; try { - // Get all ports currently in use by our managed servers - const portsInUse = new Set(); - for (const serverInfo of this.serverInfos.values()) { - if (serverInfo.jupyterPort) { - portsInUse.add(serverInfo.jupyterPort); - } - if (serverInfo.lspPort) { - portsInUse.add(serverInfo.lspPort); + // Collect ports already in use by running servers to pick a non-conflicting start port + let maxPort = 8888; + for (const ctx of this.projectContexts.values()) { + if (ctx.serverInfo) { + maxPort = Math.max(maxPort, ctx.serverInfo.jupyterPort + 2, ctx.serverInfo.lspPort + 1); } } - // Find a pair of consecutive available ports - const { jupyterPort, lspPort } = await this.findConsecutiveAvailablePorts( - DEEPNOTE_DEFAULT_PORT, - portsInUse - ); - - // Reserve both ports by adding to serverInfos - // This prevents other concurrent allocations from getting the same ports - const serverInfo = { - url: `http://localhost:${jupyterPort}`, - jupyterPort, - lspPort - }; - this.serverInfos.set(key, serverInfo); - - logger.info( - `Allocated consecutive ports for ${key}: jupyter=${jupyterPort}, lsp=${lspPort} (excluded: ${ - portsInUse.size > 2 - ? Array.from(portsInUse) - .filter((p) => p !== jupyterPort && p !== lspPort) - .join(', ') - : 'none' - })` - ); - - return { jupyterPort, lspPort }; + logger.info(`Reserved start port ${maxPort} for ${serverKey}`); + return maxPort; } finally { - // Release the lock to allow next allocation in the chain to proceed releaseLock!(); } } /** - * Find a pair of consecutive available ports (port and port+1). - * This is critical for the deepnote-toolkit server which expects consecutive ports. - * - * @param startPort The port number to start searching from - * @param portsInUse Set of ports already allocated to other servers - * @returns A pair of consecutive ports { jupyterPort, lspPort } where lspPort = jupyterPort + 1 - * @throws DeepnoteServerStartupError if no consecutive ports can be found after maxAttempts + * Gather SQL integration environment variables for the deepnote-toolkit server. */ - private async findConsecutiveAvailablePorts( - startPort: number, - portsInUse: Set - ): Promise<{ jupyterPort: number; lspPort: number }> { - const maxAttempts = 100; - - for (let attempt = 0; attempt < maxAttempts; attempt++) { - // Try to find an available Jupyter port - const candidatePort = await this.findAvailablePort( - attempt === 0 ? startPort : startPort + attempt, - portsInUse - ); + private async gatherSqlIntegrationEnvVars( + deepnoteFileUri: Uri, + environmentId: string, + token?: CancellationToken + ): Promise> { + const extraEnv: Record = {}; - const nextPort = candidatePort + 1; + if (!this.sqlIntegrationEnvVars) { + logger.debug('DeepnoteServerStarter: SqlIntegrationEnvironmentVariablesProvider not available'); + return extraEnv; + } - // Check if the consecutive port (candidatePort + 1) is also available - const isNextPortInUse = portsInUse.has(nextPort); - const isNextPortAvailable = !isNextPortInUse && (await this.isPortAvailable(nextPort)); - logger.info( - `Consecutive port check for base ${candidatePort}: next=${nextPort}, inUseSet=${isNextPortInUse}, available=${isNextPortAvailable}` - ); + const fileKey = deepnoteFileUri.fsPath; - if (isNextPortAvailable) { - // Found a consecutive pair! - return { jupyterPort: candidatePort, lspPort: nextPort }; + logger.debug( + `DeepnoteServerStarter: Injecting SQL integration env vars for ${fileKey} with environmentId ${environmentId}` + ); + try { + const sqlEnvVars = await this.sqlIntegrationEnvVars.getEnvironmentVariables(deepnoteFileUri, token); + if (sqlEnvVars && Object.keys(sqlEnvVars).length > 0) { + logger.debug(`DeepnoteServerStarter: Injecting ${Object.keys(sqlEnvVars).length} SQL env vars`); + Object.assign(extraEnv, sqlEnvVars); + } else { + logger.debug('DeepnoteServerStarter: No SQL integration env vars to inject'); } - - // Consecutive port not available - mark both as unavailable and try next - portsInUse.add(candidatePort); - portsInUse.add(nextPort); + } catch (error) { + logger.error('DeepnoteServerStarter: Failed to get SQL integration env vars', error); } - // Failed to find consecutive ports after max attempts - throw new DeepnoteServerStartupError( - 'python', - startPort, - 'process_failed', - '', - l10n.t( - 'Failed to find consecutive available ports after {0} attempts starting from port {1}. Please close some applications using network ports and try again.', - maxAttempts, - startPort - ) - ); + return extraEnv; } /** - * Check if a specific port is available on the system by actually trying to bind to it. - * This is more reliable than get-port which doesn't test the exact port. + * Stream stdout/stderr from the server process to the VSCode output channel. */ - private async isPortAvailable(port: number): Promise { - try { - const inUse = await tcpPortUsed.check(port, '127.0.0.1'); - if (inUse) { - return false; - } + private monitorServerOutput(serverKey: string, runtimeCoreInfo: RuntimeCoreServerInfo): void { + const proc = runtimeCoreInfo.process; + const disposables: IDisposable[] = []; + this.disposablesByFile.set(serverKey, disposables); - // Also check IPv6 loopback to be safe - try { - const inUseIpv6 = await tcpPortUsed.check(port, '::1'); - return !inUseIpv6; - } catch (error: unknown) { - if (error instanceof Error && 'code' in error && error.code === 'EAFNOSUPPORT') { - logger.debug('IPv6 is not supported on this system'); - return true; + if (proc.stdout) { + const stdout = proc.stdout; + const onData = (data: Buffer) => { + const text = data.toString(); + logger.trace(`Deepnote server (${serverKey}): ${text}`); + this.outputChannel.appendLine(text); + }; + stdout.on('data', onData); + disposables.push({ + dispose: () => { + stdout.off('data', onData); } - logger.warn(`Failed to check IPv6 port availability for ${port}:`, error); - return false; - } - } catch (error) { - logger.warn(`Failed to check port availability for ${port}:`, error); - return false; + }); } - } - /** - * Find an available port starting from the given port number. - * Checks both our internal portsInUse set and system availability by actually binding to test. - */ - private async findAvailablePort(startPort: number, portsInUse: Set): Promise { - let port = startPort; - let attempts = 0; - const maxAttempts = 100; - - while (attempts < maxAttempts) { - // Skip ports already in use by our servers - if (!portsInUse.has(port)) { - // Check if this port is actually available on the system by binding to it - const available = await this.isPortAvailable(port); - - if (available) { - return port; + if (proc.stderr) { + const stderr = proc.stderr; + const onData = (data: Buffer) => { + const text = data.toString(); + logger.warn(`Deepnote server stderr (${serverKey}): ${text}`); + this.outputChannel.appendLine(text); + }; + stderr.on('data', onData); + disposables.push({ + dispose: () => { + stderr.off('data', onData); } - } - - // Try next port - port++; - attempts++; + }); } - - throw new DeepnoteServerStartupError( - 'python', // unknown here - startPort, - 'process_failed', - '', - l10n.t( - 'Failed to find available port after {0} attempts (started at {1}). Ports in use: {2}', - maxAttempts, - startPort, - Array.from(portsInUse).join(', ') - ) - ); } public async dispose(): Promise { logger.info('Disposing DeepnoteServerStarter - stopping all servers...'); - // Wait for any pending operations to complete (with timeout) const pendingOps = Array.from(this.pendingOperations.values()); if (pendingOps.length > 0) { logger.info(`Waiting for ${pendingOps.length} pending operations to complete...`); - await Promise.allSettled(pendingOps.map((op) => Promise.race([op, sleep(2000)]))); + await Promise.allSettled(pendingOps.map((op) => Promise.race([op, sleep(GRACEFUL_SHUTDOWN_TIMEOUT_MS)]))); } - // Stop all server processes and wait for them to exit - const killPromises: Promise[] = []; + const stopPromises: Promise[] = []; const pidsToCleanup: number[] = []; - for (const [fileKey, serverProcess] of this.serverProcesses.entries()) { - try { - logger.info(`Stopping Deepnote server for ${fileKey}...`); - const proc = serverProcess.proc; - if (proc && !proc.killed) { - const serverPid = proc.pid; - if (serverPid) { - pidsToCleanup.push(serverPid); - } - - // Create a promise that resolves when the process exits - const exitPromise = new Promise((resolve) => { - const timeout = setTimeout(() => { - logger.warn(`Process for ${fileKey} did not exit gracefully, force killing...`); - try { - proc.kill('SIGKILL'); - } catch { - // Ignore errors on force kill - } - resolve(); - }, 3000); // Wait up to 3 seconds for graceful exit - - proc.once('exit', () => { - clearTimeout(timeout); - resolve(); - }); - }); - - // Send SIGTERM for graceful shutdown - proc.kill('SIGTERM'); - killPromises.push(exitPromise); + for (const [key, ctx] of this.projectContexts.entries()) { + if (ctx.runtimeCoreServerInfo) { + const pid = ctx.runtimeCoreServerInfo.process.pid; + if (pid) { + pidsToCleanup.push(pid); } - } catch (ex) { - logger.error(`Error stopping Deepnote server for ${fileKey}`, ex); + + logger.info(`Stopping Deepnote server for ${key}...`); + stopPromises.push( + stopServer(ctx.runtimeCoreServerInfo).catch((ex) => { + logger.error(`Error stopping Deepnote server for ${key}`, ex); + }) + ); } } - // Wait for all processes to exit - if (killPromises.length > 0) { - logger.info(`Waiting for ${killPromises.length} server processes to exit...`); - await Promise.allSettled(killPromises); + if (stopPromises.length > 0) { + logger.info(`Waiting for ${stopPromises.length} server processes to exit...`); + await Promise.allSettled(stopPromises); } - // Clean up lock files for all stopped processes for (const pid of pidsToCleanup) { await this.deleteLockFile(pid); } - // Dispose all tracked disposables for (const [fileKey, disposables] of this.disposablesByFile.entries()) { try { disposables.forEach((d) => d.dispose()); @@ -757,19 +520,15 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } - // Clear all maps - this.serverProcesses.clear(); - this.serverInfos.clear(); + this.projectContexts.clear(); this.disposablesByFile.clear(); this.pendingOperations.clear(); - this.serverOutputByFile.clear(); logger.info('DeepnoteServerStarter disposed successfully'); } - /** - * Initialize the lock file directory - */ + // ── Lock file management (extension-specific) ── + private async initializeLockFileDirectory(): Promise { try { await fs.ensureDir(this.lockFileDir); @@ -779,16 +538,10 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } - /** - * Get the lock file path for a given PID - */ private getLockFilePath(pid: number): string { return path.join(this.lockFileDir, `server-${pid}.json`); } - /** - * Write a lock file for a server process - */ private async writeLockFile(pid: number): Promise { try { const lockData: ServerLockFile = { @@ -804,9 +557,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } - /** - * Read a lock file for a given PID - */ private async readLockFile(pid: number): Promise { try { const lockFilePath = this.getLockFilePath(pid); @@ -819,9 +569,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension return null; } - /** - * Delete a lock file for a given PID - */ private async deleteLockFile(pid: number): Promise { try { const lockFilePath = this.getLockFilePath(pid); @@ -834,15 +581,13 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } - /** - * Check if a process is orphaned by verifying its parent process - */ + // ── Orphaned process cleanup (extension-specific) ── + private async isProcessOrphaned(pid: number): Promise { try { const processService = await this.processServiceFactory.create(undefined); if (process.platform === 'win32') { - // Windows: use WMIC to get parent process ID const result = await processService.exec( 'wmic', ['process', 'where', `ProcessId=${pid}`, 'get', 'ParentProcessId'], @@ -856,36 +601,27 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension if (lines.length > 0) { const ppid = parseInt(lines[0].trim(), 10); if (!isNaN(ppid)) { - // PPID of 0 means orphaned if (ppid === 0) { return true; } - // Check if parent process exists const parentCheck = await processService.exec( 'tasklist', ['/FI', `PID eq ${ppid}`, '/FO', 'CSV', '/NH'], { throwOnStdErr: false } ); - // Normalize and check stdout const stdout = (parentCheck.stdout || '').trim(); - // Parent is missing if: - // 1. stdout is empty - // 2. stdout starts with "INFO:" (case-insensitive) - // 3. stdout contains "no tasks are running" (case-insensitive) if (stdout.length === 0 || /^INFO:/i.test(stdout) || /no tasks are running/i.test(stdout)) { - return true; // Parent missing, process is orphaned + return true; } - // Parent exists return false; } } } } else { - // Unix: use ps to get parent process ID const result = await processService.exec('ps', ['-o', 'ppid=', '-p', pid.toString()], { throwOnStdErr: false }); @@ -893,11 +629,10 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension if (result.stdout) { const ppid = parseInt(result.stdout.trim(), 10); if (!isNaN(ppid)) { - // PPID of 1 typically means orphaned (adopted by init/systemd) if (ppid === 1) { return true; } - // Check if parent process exists + const parentCheck = await processService.exec('ps', ['-p', ppid.toString(), '-o', 'pid='], { throwOnStdErr: false }); @@ -909,29 +644,21 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension logger.warn(`Failed to check if process ${pid} is orphaned`, ex); } - // If we can't determine, assume it's not orphaned (safer) return false; } - /** - * Cleans up any orphaned deepnote-toolkit processes from previous VS Code sessions. - * This prevents port conflicts when starting new servers. - */ private async cleanupOrphanedProcesses(): Promise { try { logger.info('Checking for orphaned deepnote-toolkit processes...'); const processService = await this.processServiceFactory.create(undefined); - // Find all deepnote-toolkit server processes let command: string; let args: string[]; if (process.platform === 'win32') { - // Windows: use tasklist and findstr command = 'tasklist'; args = ['/FI', 'IMAGENAME eq python.exe', '/FO', 'CSV', '/NH']; } else { - // Unix-like: use ps and grep command = 'ps'; args = ['aux']; } @@ -943,19 +670,15 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const candidatePids: number[] = []; for (const line of lines) { - // Look for processes running deepnote_toolkit server if (line.includes('deepnote_toolkit') && line.includes('server')) { - // Extract PID based on platform let pid: number | undefined; if (process.platform === 'win32') { - // Windows CSV format: "python.exe","12345",... const match = line.match(/"python\.exe","(\d+)"/); if (match) { pid = parseInt(match[1], 10); } } else { - // Unix format: user PID ... const parts = line.trim().split(/\s+/); if (parts.length > 1) { pid = parseInt(parts[1], 10); @@ -976,15 +699,11 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const pidsToKill: number[] = []; const pidsToSkip: Array<{ pid: number; reason: string }> = []; - // Check each process to determine if it should be killed for (const pid of candidatePids) { - // Check if there's a lock file for this PID const lockData = await this.readLockFile(pid); if (lockData) { - // Lock file exists - check if it belongs to a different session if (lockData.sessionId !== this.sessionId) { - // Different session - check if the process is actually orphaned const isOrphaned = await this.isProcessOrphaned(pid); if (isOrphaned) { logger.info( @@ -998,23 +717,19 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension }); } } else { - // Same session - this shouldn't happen during startup, but skip it pidsToSkip.push({ pid, reason: 'belongs to current session' }); } } else { - // No lock file - assume it's an external/non-managed process and skip it pidsToSkip.push({ pid, reason: 'no lock file (assuming external process)' }); } } - // Log skipped processes if (pidsToSkip.length > 0) { for (const { pid, reason } of pidsToSkip) { logger.info(`Skipping PID ${pid}: ${reason}`); } } - // Kill orphaned processes if (pidsToKill.length > 0) { logger.info(`Killing ${pidsToKill.length} orphaned process(es): ${pidsToKill.join(', ')}`); this.outputChannel.appendLine( @@ -1032,7 +747,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } logger.info(`Killed orphaned process ${pid}`); - // Clean up the lock file after killing await this.deleteLockFile(pid); } catch (ex) { logger.warn(`Failed to kill process ${pid}`, ex); @@ -1048,7 +762,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } } catch (ex) { - // Don't fail startup if cleanup fails logger.warn('Error during orphaned process cleanup', ex); } } diff --git a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts index b6f46df475..c63174c83b 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts @@ -1,46 +1,40 @@ import { assert } from 'chai'; -import * as sinon from 'sinon'; -import tcpPortUsed from 'tcp-port-used'; import { anything, instance, mock, when } from 'ts-mockito'; + import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; import { DeepnoteServerStarter } from './deepnoteServerStarter.node'; import { IProcessServiceFactory } from '../../platform/common/process/types.node'; -import { IAsyncDisposableRegistry, IHttpClient, IOutputChannel } from '../../platform/common/types'; +import { IAsyncDisposableRegistry, IOutputChannel } from '../../platform/common/types'; import { IDeepnoteToolkitInstaller } from './types'; import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnote/types'; -import { logger } from '../../platform/logging'; -import * as net from 'net'; /** - * Integration tests for DeepnoteServerStarter port allocation logic. - * These tests use real port checking to ensure consecutive ports are allocated. + * Unit tests for DeepnoteServerStarter. * - * Note: These are integration tests that actually check port availability on the system. - * They test the critical fix where consecutive ports must be available. + * Port allocation, server spawning, and health checks are now delegated to + * @deepnote/runtime-core's startServer/stopServer. These tests focus on the + * extension-specific layers: port reservation serialization, SQL env var + * gathering, and lifecycle orchestration. */ -suite('DeepnoteServerStarter - Port Allocation Integration Tests', () => { +suite('DeepnoteServerStarter', () => { let serverStarter: DeepnoteServerStarter; let mockProcessServiceFactory: IProcessServiceFactory; let mockToolkitInstaller: IDeepnoteToolkitInstaller; let mockAgentSkillsManager: DeepnoteAgentSkillsManager; let mockOutputChannel: IOutputChannel; - let mockHttpClient: IHttpClient; let mockAsyncRegistry: IAsyncDisposableRegistry; let mockSqlIntegrationEnvVars: ISqlIntegrationEnvVarsProvider; - // Helper to access private methods for testing // eslint-disable-next-line @typescript-eslint/no-explicit-any const getPrivateMethod = (obj: any, methodName: string) => { return obj[methodName].bind(obj); }; setup(() => { - // Create mocks mockProcessServiceFactory = mock(); mockToolkitInstaller = mock(); mockAgentSkillsManager = mock(); mockOutputChannel = mock(); - mockHttpClient = mock(); mockAsyncRegistry = mock(); mockSqlIntegrationEnvVars = mock(); @@ -52,527 +46,89 @@ suite('DeepnoteServerStarter - Port Allocation Integration Tests', () => { instance(mockToolkitInstaller), instance(mockAgentSkillsManager), instance(mockOutputChannel), - instance(mockHttpClient), instance(mockAsyncRegistry), instance(mockSqlIntegrationEnvVars) ); }); teardown(async () => { - // Dispose the serverStarter to clean up any allocated ports and state if (serverStarter) { await serverStarter.dispose(); } }); - suite('isPortAvailable', () => { - let checkStub: sinon.SinonStub; - - setup(() => { - checkStub = sinon.stub(tcpPortUsed, 'check'); - }); - - teardown(() => { - checkStub.restore(); - }); - - test('should return true when both IPv4 and IPv6 loopbacks are free', async () => { - const port = 54321; - checkStub.onFirstCall().resolves(false); // IPv4 - checkStub.onSecondCall().resolves(false); // IPv6 - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const isPortAvailable = getPrivateMethod(serverStarter as any, 'isPortAvailable'); - const result = await isPortAvailable(port); - - assert.isTrue(result, 'Expected port to be reported as available'); - assert.strictEqual(checkStub.callCount, 2, 'Should check both IPv4 and IPv6 loopbacks'); - assert.deepEqual(checkStub.getCall(0).args, [port, '127.0.0.1']); - assert.deepEqual(checkStub.getCall(1).args, [port, '::1']); - }); - - test('should return false when IPv4 loopback is already in use', async () => { - const port = 54322; - checkStub.onFirstCall().resolves(true); - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const isPortAvailable = getPrivateMethod(serverStarter as any, 'isPortAvailable'); - const result = await isPortAvailable(port); - - assert.isFalse(result, 'Expected port to be reported as in use'); - assert.strictEqual(checkStub.callCount, 1, 'IPv6 check should be skipped when IPv4 is busy'); - }); - - test('should return false and log when port checks throw', async () => { - const port = 54323; - const error = new Error('check failed'); - checkStub.rejects(error); - - const warnStub = sinon.stub(logger, 'warn'); - - try { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const isPortAvailable = getPrivateMethod(serverStarter as any, 'isPortAvailable'); - const result = await isPortAvailable(port); - - assert.isFalse(result, 'Expected port check to fail closed when an error occurs'); - assert.isTrue(warnStub.called, 'Expected warning to be logged when check fails'); - } finally { - warnStub.restore(); - } - }); - - test('should return true when IPv6 is disabled (EAFNOSUPPORT error)', async () => { - const port = 54324; - const ipv6Error = new Error('connect EAFNOSUPPORT ::1:54324'); - (ipv6Error as any).code = 'EAFNOSUPPORT'; - - // IPv4 check succeeds (port is available) - checkStub.onFirstCall().resolves(false); - - // IPv6 check throws EAFNOSUPPORT (IPv6 not supported) - checkStub.onSecondCall().rejects(ipv6Error); - - const debugStub = sinon.stub(logger, 'debug'); - - try { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const isPortAvailable = getPrivateMethod(serverStarter as any, 'isPortAvailable'); - const result = await isPortAvailable(port); - - assert.isTrue(result, 'Expected port to be available when IPv4 is free and IPv6 is not supported'); - assert.strictEqual(checkStub.callCount, 2, 'Should check both IPv4 and IPv6'); - assert.deepEqual(checkStub.getCall(0).args, [port, '127.0.0.1']); - assert.deepEqual(checkStub.getCall(1).args, [port, '::1']); - assert.isTrue( - debugStub.calledWith('IPv6 is not supported on this system'), - 'Should log debug message about IPv6 not being supported' - ); - } finally { - debugStub.restore(); - } - }); - - test('should return false when IPv6 check throws non-EAFNOSUPPORT error', async () => { - const port = 54325; - const ipv6Error = new Error('Some other IPv6 error'); - - // IPv4 check succeeds (port is available) - checkStub.onFirstCall().resolves(false); - - // IPv6 check throws a different error - checkStub.onSecondCall().rejects(ipv6Error); - - const warnStub = sinon.stub(logger, 'warn'); - - try { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const isPortAvailable = getPrivateMethod(serverStarter as any, 'isPortAvailable'); - const result = await isPortAvailable(port); - - assert.isFalse( - result, - 'Expected port check to fail closed when IPv6 check fails with non-EAFNOSUPPORT error' - ); - assert.strictEqual(checkStub.callCount, 2, 'Should check both IPv4 and IPv6'); - assert.isTrue(warnStub.called, 'Should log warning when IPv6 check fails'); - const warnCall = warnStub.getCall(0); - assert.include(warnCall.args[0], 'Failed to check IPv6 port availability'); - } finally { - warnStub.restore(); - } - }); - }); - - suite('findAvailablePort', () => { - test('should find an available port starting from given port', async () => { - const portsInUse = new Set(); - const startPort = 54400; + suite('reserveStartPort - Port Serialization', () => { + test('should return default port when no servers are running', async () => { + const reserveStartPort = getPrivateMethod(serverStarter, 'reserveStartPort'); + const port = await reserveStartPort('test-key'); - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const findAvailablePort = getPrivateMethod(serverStarter as any, 'findAvailablePort'); - const result = await findAvailablePort(startPort, portsInUse); - - // Should find a port at or after the start port - assert.isAtLeast(result, startPort); + assert.strictEqual(port, 8888); }); - test('should skip ports in portsInUse set', async () => { - const portsInUse = new Set([54500, 54501, 54502]); - const startPort = 54500; + test('should return ports beyond existing servers', async () => { + const reserveStartPort = getPrivateMethod(serverStarter, 'reserveStartPort'); + // Simulate a running server context by directly setting projectContexts // eslint-disable-next-line @typescript-eslint/no-explicit-any - const findAvailablePort = getPrivateMethod(serverStarter as any, 'findAvailablePort'); - const result = await findAvailablePort(startPort, portsInUse); - - // Should skip the ports in use - assert.isFalse(portsInUse.has(result), 'Should not return a port from portsInUse'); - assert.isAtLeast(result, 54503); - }); - - test('should find available port within reasonable attempts', async () => { - const portsInUse = new Set(); - const startPort = 54600; - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const findAvailablePort = getPrivateMethod(serverStarter as any, 'findAvailablePort'); - const result = await findAvailablePort(startPort, portsInUse); - - // Should find a port without error - assert.isNumber(result); - assert.isAtLeast(result, startPort); - }); - }); - - suite('allocatePorts - Consecutive Port Allocation (Critical Bug Fix)', () => { - test('should allocate consecutive ports (LSP = Jupyter + 1)', async () => { - const key = 'test-consecutive-1'; - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const allocatePorts = getPrivateMethod(serverStarter as any, 'allocatePorts'); - const result = await allocatePorts(key); - - // THIS IS THE CRITICAL ASSERTION: LSP port must be exactly Jupyter + 1 - assert.strictEqual( - result.lspPort, - result.jupyterPort + 1, - 'LSP port must be consecutive (Jupyter port + 1)' - ); - }); - - test('should allocate different consecutive port pairs for multiple servers', async () => { - const key1 = 'test-server-1'; - const key2 = 'test-server-2'; - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const allocatePorts = getPrivateMethod(serverStarter as any, 'allocatePorts'); - - const result1 = await allocatePorts(key1); - const result2 = await allocatePorts(key2); - - // Both should have consecutive ports - assert.strictEqual(result1.lspPort, result1.jupyterPort + 1); - assert.strictEqual(result2.lspPort, result2.jupyterPort + 1); - - // Ports should not overlap - assert.notEqual(result1.jupyterPort, result2.jupyterPort); - assert.notEqual(result1.lspPort, result2.lspPort); - assert.notEqual(result1.jupyterPort, result2.lspPort); - assert.notEqual(result1.lspPort, result2.jupyterPort); - }); - - test('CRITICAL REGRESSION TEST: should skip non-consecutive ports when LSP port is taken', async () => { - // This test simulates the EXACT bug scenario that was reported: - // - Port 8888 is available - // - Port 8889 (8888+1) is TAKEN by another process - // - System should NOT allocate 8888+8890 (non-consecutive) - // - System SHOULD find a different consecutive pair like 8890+8891 - - const blockingServer = net.createServer(); - const blockedPort = 54701; // We'll block this port to simulate 8889 being taken - - // Bind to port 54701 to block it - await new Promise((resolve) => { - blockingServer.listen(blockedPort, 'localhost', () => { - resolve(); - }); + const projectContexts = (serverStarter as any).projectContexts as Map; + projectContexts.set('existing-key', { + environmentId: 'env1', + runtimeCoreServerInfo: null, + serverInfo: { url: 'http://localhost:8888', jupyterPort: 8888, lspPort: 8889 } }); - try { - const key = 'test-blocked-lsp-port'; + const port = await reserveStartPort('test-key-2'); - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const allocatePorts = getPrivateMethod(serverStarter as any, 'allocatePorts'); - - // Try to allocate ports - it should skip 54700 because 54701 is taken - const result = await allocatePorts(key); - - // CRITICAL: Ports must be consecutive - assert.strictEqual( - result.lspPort, - result.jupyterPort + 1, - 'Even when some ports are blocked, allocated ports MUST be consecutive' - ); - - // Should not have allocated the blocked port or its predecessor - assert.notEqual(result.jupyterPort, blockedPort); - assert.notEqual(result.lspPort, blockedPort); - assert.isFalse( - result.jupyterPort === blockedPort - 1 && result.lspPort === blockedPort, - 'Should not allocate pair where second port is blocked' - ); - } finally { - // Clean up: close the blocking server - await new Promise((resolve) => { - blockingServer.close(() => resolve()); - }); - } + assert.isAtLeast(port, 8890, 'Should skip ports used by existing servers'); }); - test('should handle rapid sequential allocations', async () => { - const keys = ['seq-1', 'seq-2', 'seq-3', 'seq-4', 'seq-5']; + test('should serialize concurrent calls', async () => { + const reserveStartPort = getPrivateMethod(serverStarter, 'reserveStartPort'); - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const allocatePorts = getPrivateMethod(serverStarter as any, 'allocatePorts'); + // Launch concurrent port reservations + const [port1, port2, port3] = await Promise.all([ + reserveStartPort('key-1'), + reserveStartPort('key-2'), + reserveStartPort('key-3') + ]); - const results = []; - for (const key of keys) { - const result = await allocatePorts(key); - results.push(result); - } - - // All should have unique, consecutive port pairs - const allPorts = results.flatMap((r) => [r.jupyterPort, r.lspPort]); - const uniquePorts = new Set(allPorts); - assert.strictEqual(uniquePorts.size, results.length * 2, 'All ports should be unique'); - - // Each result should have consecutive ports - for (const result of results) { - assert.strictEqual(result.lspPort, result.jupyterPort + 1); - } - }); - - test('should update serverInfos map with allocated ports', async () => { - const key = 'test-server-info'; - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const allocatePorts = getPrivateMethod(serverStarter as any, 'allocatePorts'); - const result = await allocatePorts(key); - - // Check that serverInfos was updated - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const serverInfos = (serverStarter as any).serverInfos as Map; - assert.isTrue(serverInfos.has(key)); - - const info = serverInfos.get(key); - assert.strictEqual(info.jupyterPort, result.jupyterPort); - assert.strictEqual(info.lspPort, result.lspPort); - assert.strictEqual(info.url, `http://localhost:${result.jupyterPort}`); - }); - - test('should respect already allocated ports', async () => { - // First allocation - const key1 = 'first-server'; - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const allocatePorts = getPrivateMethod(serverStarter as any, 'allocatePorts'); - const result1 = await allocatePorts(key1); - - // Second allocation should get different ports - const key2 = 'second-server'; - const result2 = await allocatePorts(key2); - - // Verify no overlap - const ports1 = new Set([result1.jupyterPort, result1.lspPort]); - assert.isFalse(ports1.has(result2.jupyterPort), 'Second Jupyter port should not overlap'); - assert.isFalse(ports1.has(result2.lspPort), 'Second LSP port should not overlap'); + // All should return valid numbers (even if same, since no server info is stored between calls) + assert.isNumber(port1); + assert.isNumber(port2); + assert.isNumber(port3); }); }); - suite('Port Allocation Edge Cases', () => { - test('should allocate ports successfully even after multiple allocations', async () => { - // Allocate many port pairs to test robustness - const count = 10; - const keys = Array.from({ length: count }, (_, i) => `stress-test-${i}`); - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const allocatePorts = getPrivateMethod(serverStarter as any, 'allocatePorts'); - - const results = []; - for (const key of keys) { - const result = await allocatePorts(key); - results.push(result); - } - - // All should be successful and consecutive - assert.strictEqual(results.length, count); - for (const result of results) { - assert.strictEqual(result.lspPort, result.jupyterPort + 1); - } - - // All ports should be unique - const allPorts = results.flatMap((r) => [r.jupyterPort, r.lspPort]); - const uniquePorts = new Set(allPorts); - assert.strictEqual(uniquePorts.size, count * 2); - }); + suite('gatherSqlIntegrationEnvVars', () => { + test('should return empty object when no provider is available', async () => { + // Create a starter without SQL provider + const starterWithoutSql = new DeepnoteServerStarter( + instance(mockProcessServiceFactory), + instance(mockToolkitInstaller), + instance(mockAgentSkillsManager), + instance(mockOutputChannel), + instance(mockAsyncRegistry) + ); - test('should return valid port numbers', async () => { - const key = 'valid-ports'; + const gatherEnvVars = getPrivateMethod(starterWithoutSql, 'gatherSqlIntegrationEnvVars'); + const { Uri } = await import('vscode'); + const result = await gatherEnvVars(Uri.file('/test/file.deepnote'), 'env1'); - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const allocatePorts = getPrivateMethod(serverStarter as any, 'allocatePorts'); - const result = await allocatePorts(key); + assert.deepStrictEqual(result, {}); - // Ports should be in valid range - assert.isAtLeast(result.jupyterPort, 1024, 'Port should be above well-known ports'); - assert.isBelow(result.jupyterPort, 65536, 'Port should be below max port number'); - assert.isAtLeast(result.lspPort, 1024); - assert.isBelow(result.lspPort, 65536); + await starterWithoutSql.dispose(); }); }); - suite('Critical Bug Fix Verification', () => { - test('REGRESSION TEST: should never allocate non-consecutive ports', async () => { - // This is the critical regression test for the bug where - // if Jupyter port was available but LSP port (Jupyter+1) was not, - // the system would allocate non-consecutive ports causing server hangs - - // Use unique keys with timestamp to avoid conflicts with other tests - const timestamp = Date.now(); - const keys = [ - `concurrent-test-${timestamp}-1`, - `concurrent-test-${timestamp}-2`, - `concurrent-test-${timestamp}-3` - ]; + suite('dispose', () => { + test('should clear all internal state', async () => { + await serverStarter.dispose(); // eslint-disable-next-line @typescript-eslint/no-explicit-any - const allocatePorts = getPrivateMethod(serverStarter as any, 'allocatePorts'); - - const results = await Promise.all(keys.map((key) => allocatePorts(key))); - - // Verify each result has consecutive ports - for (let i = 0; i < results.length; i++) { - const result = results[i]; - assert.strictEqual( - result.lspPort, - result.jupyterPort + 1, - `Server ${i + 1} (${keys[i]}): LSP port MUST be Jupyter port + 1. ` + - `This prevents server startup hangs when toolkit expects consecutive ports.` - ); - } - - // Verify uniqueness: no two concurrent calls received the same port pair - const portPairs = new Set(results.map((r) => `${r.jupyterPort}:${r.lspPort}`)); - assert.strictEqual( - portPairs.size, - results.length, - 'All concurrent allocations must receive unique port pairs' - ); - - // Verify uniqueness of individual ports - const allPorts = results.flatMap((r) => [r.jupyterPort, r.lspPort]); - const uniquePorts = new Set(allPorts); - assert.strictEqual( - uniquePorts.size, - allPorts.length, - 'All allocated ports (both Jupyter and LSP) must be unique across concurrent calls' - ); - }); - }); - - suite('findConsecutiveAvailablePorts - Edge Cases', () => { - test('should mark both ports unavailable and continue when consecutive port is taken', async () => { - // This test covers the scenario where a candidate port is available - // but the next port (candidate + 1) is not available. - // The system should mark BOTH ports as unavailable in portsInUse and continue searching. - - const server1 = net.createServer(); - const server2 = net.createServer(); - const blockedPort1 = 54801; - const blockedPort2 = 54803; - - // Block ports 54801 and 54803 (leaving 54800 and 54802 available but not consecutive) - await new Promise((resolve) => { - server1.listen(blockedPort1, 'localhost', () => { - server2.listen(blockedPort2, 'localhost', () => { - resolve(); - }); - }); - }); - - try { - const portsInUse = new Set(); - const startPort = 54800; - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const findConsecutiveAvailablePorts = getPrivateMethod( - serverStarter as any, - 'findConsecutiveAvailablePorts' - ); - - // Should skip 54800 (since 54801 is blocked) and 54802 (since 54803 is blocked) - // and find the next consecutive pair like 54804+54805 - const result = await findConsecutiveAvailablePorts(startPort, portsInUse); - - // Verify ports are consecutive - assert.strictEqual(result.lspPort, result.jupyterPort + 1); - - // Should have found ports after the blocked ones - assert.isTrue( - result.jupyterPort > blockedPort2 || result.jupyterPort < blockedPort1 - 1, - 'Should skip blocked port ranges' - ); - } finally { - // Clean up - await new Promise((resolve) => { - server1.close(() => { - server2.close(() => resolve()); - }); - }); - } - }); - - test('should throw DeepnoteServerStartupError when max attempts exhausted', async () => { - // This test covers the scenario where we cannot find consecutive ports - // after maxAttempts (100 attempts). This should throw a DeepnoteServerStartupError. - // Strategy: Block every other port so individual ports are available, - // but no consecutive pairs exist (blocking the +1 port for each available port) - - const servers: any[] = []; - - try { - // Block every other port starting from 55001 (leaving 55000, 55002, 55004, etc. available) - // This ensures findAvailablePort succeeds, but the consecutive port is always blocked - const startPort = 55000; - const portsToBlock = 200; // Block 200 odd-numbered ports - - // Create servers blocking every other port (the +1 ports) - for (let i = 0; i < portsToBlock; i++) { - const portToBlock = startPort + i * 2 + 1; // Block 55001, 55003, 55005, etc. - const server = net.createServer(); - servers.push(server); - await new Promise((resolve) => { - server.listen(portToBlock, 'localhost', () => resolve()); - }); - } - - const portsInUse = new Set(); - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const findConsecutiveAvailablePorts = getPrivateMethod( - serverStarter as any, - 'findConsecutiveAvailablePorts' - ); - - // Should throw DeepnoteServerStartupError after maxAttempts - // Note: The error could come from either findConsecutiveAvailablePorts or findAvailablePort - // depending on port availability timing - let errorThrown = false; - try { - await findConsecutiveAvailablePorts(startPort, portsInUse); - } catch (error: any) { - errorThrown = true; - assert.strictEqual(error.constructor.name, 'DeepnoteServerStartupError'); - // Accept either error message since both indicate port exhaustion - const isConsecutiveError = error.stderr.includes('Failed to find consecutive available ports'); - const isSinglePortError = error.stderr.includes('Failed to find available port'); - assert.isTrue( - isConsecutiveError || isSinglePortError, - `Expected port exhaustion error, got: ${error.stderr}` - ); - } - - assert.isTrue(errorThrown, 'Expected DeepnoteServerStartupError to be thrown'); - } finally { - // Clean up all servers - await Promise.all( - servers.map( - (server) => - new Promise((resolve) => { - server.close(() => resolve()); - }) - ) - ); - } + const starter = serverStarter as any; + assert.strictEqual(starter.projectContexts.size, 0); + assert.strictEqual(starter.disposablesByFile.size, 0); + assert.strictEqual(starter.pendingOperations.size, 0); }); }); }); diff --git a/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts b/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts index c18c9d4a47..2c7ebdadbc 100644 --- a/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts +++ b/src/kernels/deepnote/deepnoteToolkitInstaller.node.ts @@ -4,6 +4,8 @@ import { inject, injectable, named } from 'inversify'; import { CancellationToken, l10n, Uri, workspace } from 'vscode'; +import { resolvePythonExecutable } from '@deepnote/runtime-core'; + import { Cancellation } from '../../platform/common/cancellation'; import { STANDARD_OUTPUT_CHANNEL } from '../../platform/common/constants'; import { IFileSystem } from '../../platform/common/platform/types'; @@ -44,6 +46,8 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { /** * Get the venv Python interpreter by direct venv path. + * Uses @deepnote/runtime-core's `resolvePythonExecutable` which handles + * venv root, bin dir, and bare command detection across all platforms. */ private async getVenvInterpreterByPath(venvPath: Uri): Promise { const cacheKey = venvPath.fsPath; @@ -52,18 +56,15 @@ export class DeepnoteToolkitInstaller implements IDeepnoteToolkitInstaller { return { uri: this.venvPythonPaths.get(cacheKey)!, id: this.venvPythonPaths.get(cacheKey)!.fsPath }; } - // Check if venv exists - const pythonInVenv = - process.platform === 'win32' - ? Uri.joinPath(venvPath, 'Scripts', 'python.exe') - : Uri.joinPath(venvPath, 'bin', 'python'); + try { + const resolvedPath = await resolvePythonExecutable(venvPath.fsPath); + const pythonUri = Uri.file(resolvedPath); - if (await this.fs.exists(pythonInVenv)) { - this.venvPythonPaths.set(cacheKey, pythonInVenv); - return { uri: pythonInVenv, id: pythonInVenv.fsPath }; + this.venvPythonPaths.set(cacheKey, pythonUri); + return { uri: pythonUri, id: pythonUri.fsPath }; + } catch { + return undefined; } - - return undefined; } public async getVenvInterpreter(deepnoteFileUri: Uri): Promise { diff --git a/src/kernels/deepnote/types.ts b/src/kernels/deepnote/types.ts index ef64ae04e0..5c63eacd1d 100644 --- a/src/kernels/deepnote/types.ts +++ b/src/kernels/deepnote/types.ts @@ -1,6 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. +import type { ChildProcess } from 'node:child_process'; import * as vscode from 'vscode'; import { serializePythonEnvironment } from '../../platform/api/pythonApi'; @@ -190,6 +191,8 @@ export interface DeepnoteServerInfo { jupyterPort: number; lspPort: number; token?: string; + /** The underlying server process from @deepnote/runtime-core, used for lifecycle management */ + process?: ChildProcess; } export const IDeepnoteServerProvider = Symbol('IDeepnoteServerProvider'); From ff6d44f4ee6715cb67cf37350d233374563469e9 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 11 Mar 2026 15:28:12 +0000 Subject: [PATCH 04/19] feat: Add Agent block visualization support --- .../deepnote/agentCellStatusBarProvider.ts | 249 +++++++++++++++++ .../agentCellStatusBarProvider.unit.test.ts | 251 ++++++++++++++++++ .../converters/agentBlockConverter.ts | 34 +++ .../agentBlockConverter.unit.test.ts | 198 ++++++++++++++ .../deepnote/deepnoteDataConverter.ts | 2 + .../ephemeralCellDecorationProvider.ts | 123 +++++++++ .../ephemeralCellStatusBarProvider.ts | 83 ++++++ ...phemeralCellStatusBarProvider.unit.test.ts | 169 ++++++++++++ src/notebooks/serviceRegistry.node.ts | 15 ++ src/notebooks/serviceRegistry.web.ts | 15 ++ src/renderers/client/markdown.ts | 48 +++- 11 files changed, 1186 insertions(+), 1 deletion(-) create mode 100644 src/notebooks/deepnote/agentCellStatusBarProvider.ts create mode 100644 src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts create mode 100644 src/notebooks/deepnote/converters/agentBlockConverter.ts create mode 100644 src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts create mode 100644 src/notebooks/deepnote/ephemeralCellDecorationProvider.ts create mode 100644 src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts create mode 100644 src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts new file mode 100644 index 0000000000..8d4ba8eba4 --- /dev/null +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -0,0 +1,249 @@ +import { + CancellationToken, + Disposable, + EventEmitter, + NotebookCell, + NotebookCellStatusBarItem, + NotebookCellStatusBarItemProvider, + NotebookEdit, + WorkspaceEdit, + commands, + l10n, + notebooks, + window, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import type { Pocket } from '../../platform/deepnote/pocket'; + +const DEFAULT_MAX_ITERATIONS = 20; +const MIN_ITERATIONS = 1; +const MAX_ITERATIONS = 100; +const AGENT_MODEL_OPTIONS = ['auto', 'gpt-4o', 'sonnet']; + +/** + * Provides status bar items for agent cells showing the block type indicator, + * AI model picker, and max iterations setting. + */ +@injectable() +export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService { + private readonly disposables: Disposable[] = []; + private readonly _onDidChangeCellStatusBarItems = new EventEmitter(); + + public readonly onDidChangeCellStatusBarItems = this._onDidChangeCellStatusBarItems.event; + + public activate(): void { + this.disposables.push(notebooks.registerNotebookCellStatusBarItemProvider('deepnote', this)); + + this.disposables.push( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook.notebookType === 'deepnote') { + this._onDidChangeCellStatusBarItems.fire(); + } + }) + ); + + this.disposables.push( + commands.registerCommand('deepnote.switchAgentModel', async (cell?: NotebookCell) => { + const activeCell = cell || this.getActiveCell(); + if (activeCell) { + await this.switchModel(activeCell); + } + }) + ); + + this.disposables.push( + commands.registerCommand('deepnote.setAgentMaxIterations', async (cell?: NotebookCell) => { + const activeCell = cell || this.getActiveCell(); + if (activeCell) { + await this.setMaxIterations(activeCell); + } + }) + ); + + this.disposables.push(this._onDidChangeCellStatusBarItems); + } + + public dispose(): void { + this.disposables.forEach((d) => d.dispose()); + } + + public provideCellStatusBarItems( + cell: NotebookCell, + token: CancellationToken + ): NotebookCellStatusBarItem[] | undefined { + if (token.isCancellationRequested) { + return undefined; + } + + if (!this.isAgentCell(cell)) { + return undefined; + } + + const metadata = cell.metadata as Record | undefined; + const model = this.getModel(metadata); + const maxIterations = this.getMaxIterations(metadata); + + return [ + this.createAgentIndicatorItem(), + this.createModelPickerItem(cell, model), + this.createMaxIterationsItem(cell, maxIterations) + ]; + } + + private createAgentIndicatorItem(): NotebookCellStatusBarItem { + return { + text: `$(hubot) ${l10n.t('Agent Block')}`, + alignment: 1, + priority: 100, + tooltip: l10n.t('Deepnote Agent Block\nAI-powered block that autonomously generates code and analysis') + }; + } + + private createMaxIterationsItem(cell: NotebookCell, maxIterations: number): NotebookCellStatusBarItem { + return { + text: l10n.t('$(iterations) Max iterations: {0}', maxIterations), + alignment: 1, + priority: 80, + tooltip: l10n.t('Maximum iterations for agent\nClick to change'), + command: { + title: l10n.t('Set Max Iterations'), + command: 'deepnote.setAgentMaxIterations', + arguments: [cell] + } + }; + } + + private createModelPickerItem(cell: NotebookCell, model: string): NotebookCellStatusBarItem { + return { + text: `$(symbol-enum) ${l10n.t('Model: {0}', model)}`, + alignment: 1, + priority: 90, + tooltip: l10n.t('AI Model: {0}\nClick to change', model), + command: { + title: l10n.t('Switch Model'), + command: 'deepnote.switchAgentModel', + arguments: [cell] + } + }; + } + + private getActiveCell(): NotebookCell | undefined { + const activeEditor = window.activeNotebookEditor; + if (activeEditor && activeEditor.selection) { + return activeEditor.notebook.cellAt(activeEditor.selection.start); + } + + return undefined; + } + + private getMaxIterations(metadata: Record | undefined): number { + const value = metadata?.deepnote_max_iterations; + if (typeof value === 'number' && Number.isInteger(value) && value >= MIN_ITERATIONS) { + return value; + } + + return DEFAULT_MAX_ITERATIONS; + } + + private getModel(metadata: Record | undefined): string { + const value = metadata?.deepnote_model; + if (typeof value === 'string' && value) { + return value; + } + + return 'auto'; + } + + private isAgentCell(cell: NotebookCell): boolean { + const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; + + return pocket?.type === 'agent'; + } + + private async setMaxIterations(cell: NotebookCell): Promise { + if (!this.isAgentCell(cell)) { + return; + } + + const metadata = cell.metadata as Record | undefined; + const currentValue = this.getMaxIterations(metadata); + + const input = await window.showInputBox({ + prompt: l10n.t('Enter maximum number of iterations ({0}-{1})', MIN_ITERATIONS, MAX_ITERATIONS), + value: String(currentValue), + validateInput: (value) => { + const num = parseInt(value, 10); + if (isNaN(num) || !Number.isInteger(num)) { + return l10n.t('Please enter a whole number'); + } + if (num < MIN_ITERATIONS || num > MAX_ITERATIONS) { + return l10n.t('Value must be between {0} and {1}', MIN_ITERATIONS, MAX_ITERATIONS); + } + + return undefined; + } + }); + + if (input === undefined) { + return; + } + + const newValue = parseInt(input, 10); + if (newValue === currentValue) { + return; + } + + await this.updateCellMetadata(cell, { deepnote_max_iterations: newValue }); + } + + private async switchModel(cell: NotebookCell): Promise { + if (!this.isAgentCell(cell)) { + return; + } + + const metadata = cell.metadata as Record | undefined; + const currentModel = this.getModel(metadata); + + const items = AGENT_MODEL_OPTIONS.map((option) => ({ + label: option, + description: option === currentModel ? l10n.t('Currently selected') : undefined + })); + + const selected = await window.showQuickPick(items, { + placeHolder: l10n.t('Select AI model for agent') + }); + + if (!selected || selected.label === currentModel) { + return; + } + + const newModel = selected.label === 'auto' ? undefined : selected.label; + + await this.updateCellMetadata(cell, { deepnote_model: newModel }); + } + + private async updateCellMetadata(cell: NotebookCell, updates: Record): Promise { + const updatedMetadata = { ...cell.metadata, ...updates }; + + // Remove keys set to undefined so they don't persist + for (const [key, value] of Object.entries(updates)) { + if (value === undefined) { + delete updatedMetadata[key]; + } + } + + const edit = new WorkspaceEdit(); + edit.set(cell.notebook.uri, [NotebookEdit.updateCellMetadata(cell.index, updatedMetadata)]); + + const success = await workspace.applyEdit(edit); + if (!success) { + void window.showErrorMessage(l10n.t('Failed to update agent cell metadata')); + return; + } + + this._onDidChangeCellStatusBarItems.fire(); + } +} diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts new file mode 100644 index 0000000000..e5397a1948 --- /dev/null +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -0,0 +1,251 @@ +import { expect } from 'chai'; +import { CancellationToken } from 'vscode'; + +import { AgentCellStatusBarProvider } from './agentCellStatusBarProvider'; +import { createMockCell } from './deepnoteTestHelpers'; + +suite('AgentCellStatusBarProvider', () => { + let provider: AgentCellStatusBarProvider; + let mockToken: CancellationToken; + + setup(() => { + mockToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + provider = new AgentCellStatusBarProvider(); + }); + + teardown(() => { + provider.dispose(); + }); + + suite('Agent Cell Detection', () => { + test('Should return status bar items for agent cell', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.not.be.undefined; + expect(items).to.have.lengthOf(3); + }); + + test('Should return undefined for code cell', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined for sql cell', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'sql' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined for markdown cell', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'markdown' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined for cell without pocket', () => { + const cell = createMockCell({ metadata: {} }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined for cell without metadata', () => { + const cell = createMockCell({ metadata: undefined }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined when cancellation is requested', () => { + const cancelledToken: CancellationToken = { + isCancellationRequested: true, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, cancelledToken); + + expect(items).to.be.undefined; + }); + }); + + suite('Agent Block Indicator', () => { + test('Should display agent block label with icon', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[0].text).to.include('$(hubot)'); + expect(items[0].text).to.include('Agent Block'); + expect(items[0].alignment).to.equal(1); + expect(items[0].priority).to.equal(100); + }); + + test('Should not have a command on the indicator', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[0].command).to.be.undefined; + }); + }); + + suite('Model Picker', () => { + test('Should display "auto" when no model is set', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: auto'); + expect(items[1].text).to.include('$(symbol-enum)'); + }); + + test('Should display configured model from metadata', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_model: 'gpt-4o' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: gpt-4o'); + }); + + test('Should display sonnet model', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_model: 'sonnet' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: sonnet'); + }); + + test('Should display "auto" when model is empty string', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_model: '' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: auto'); + }); + + test('Should have switch model command', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].command).to.not.be.undefined; + const cmd = items[1].command as any; + expect(cmd.command).to.equal('deepnote.switchAgentModel'); + }); + + test('Should have priority 90', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].priority).to.equal(90); + }); + }); + + suite('Max Iterations', () => { + test('Should display default max iterations (20) when not set', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + expect(items[2].text).to.include('$(iterations)'); + }); + + test('Should display configured max iterations from metadata', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: 10 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 10'); + }); + + test('Should display default when max iterations is not a number', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: 'invalid' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + }); + + test('Should display default when max iterations is zero', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: 0 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + }); + + test('Should display default when max iterations is a float', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: 5.5 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + }); + + test('Should have set max iterations command', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].command).to.not.be.undefined; + const cmd = items[2].command as any; + expect(cmd.command).to.equal('deepnote.setAgentMaxIterations'); + }); + + test('Should have priority 80', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].priority).to.equal(80); + }); + }); + + suite('Combined metadata', () => { + test('Should display both model and max iterations from metadata', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_model: 'gpt-4o', + deepnote_max_iterations: 50 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items).to.have.lengthOf(3); + expect(items[0].text).to.include('Agent Block'); + expect(items[1].text).to.include('Model: gpt-4o'); + expect(items[2].text).to.include('Max iterations: 50'); + }); + }); +}); diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.ts b/src/notebooks/deepnote/converters/agentBlockConverter.ts new file mode 100644 index 0000000000..357ab5d3c9 --- /dev/null +++ b/src/notebooks/deepnote/converters/agentBlockConverter.ts @@ -0,0 +1,34 @@ +import type { DeepnoteBlock } from '@deepnote/blocks'; +import { NotebookCellData, NotebookCellKind } from 'vscode'; + +import type { BlockConverter } from './blockConverter'; + +/** + * Converter for agent blocks. + * + * Agent blocks are rendered as code cells with markdown language so the natural-language + * prompt gets reasonable syntax highlighting while remaining visually distinct from + * Python code blocks. The prompt text is stored in `block.content`. + * + * Agent-specific metadata (model, MCP servers, max iterations, etc.) is preserved + * through the generic metadata pass-through in DeepnoteDataConverter. + */ +export class AgentBlockConverter implements BlockConverter { + applyChangesToBlock(block: DeepnoteBlock, cell: NotebookCellData): void { + block.content = cell.value || ''; + } + + canConvert(blockType: string): boolean { + return blockType.toLowerCase() === 'agent'; + } + + convertToCell(block: DeepnoteBlock): NotebookCellData { + const cell = new NotebookCellData(NotebookCellKind.Code, block.content || '', 'markdown'); + + return cell; + } + + getSupportedTypes(): string[] { + return ['agent']; + } +} diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts new file mode 100644 index 0000000000..43fac425db --- /dev/null +++ b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts @@ -0,0 +1,198 @@ +import type { DeepnoteBlock } from '@deepnote/blocks'; +import { assert } from 'chai'; +import { NotebookCellData, NotebookCellKind } from 'vscode'; +import { AgentBlockConverter } from './agentBlockConverter'; +import dedent from 'dedent'; + +suite('AgentBlockConverter', () => { + let converter: AgentBlockConverter; + + setup(() => { + converter = new AgentBlockConverter(); + }); + + suite('canConvert', () => { + test('returns true for "agent" type', () => { + assert.strictEqual(converter.canConvert('agent'), true); + }); + + test('returns true for "Agent" type (case insensitive)', () => { + assert.strictEqual(converter.canConvert('Agent'), true); + }); + + test('returns false for other types', () => { + assert.strictEqual(converter.canConvert('code'), false); + assert.strictEqual(converter.canConvert('markdown'), false); + assert.strictEqual(converter.canConvert('sql'), false); + }); + }); + + suite('getSupportedTypes', () => { + test('returns array with "agent"', () => { + const types = converter.getSupportedTypes(); + + assert.deepStrictEqual(types, ['agent']); + }); + }); + + suite('convertToCell', () => { + test('converts agent block to code cell with markdown language', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Analyze the dataset and create a summary report', + id: 'agent-block-123', + sortingKey: 'a0', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, 'Analyze the dataset and create a summary report'); + assert.strictEqual(cell.languageId, 'markdown'); + }); + + test('handles empty content', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: '', + id: 'agent-block-456', + sortingKey: 'a1', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, ''); + assert.strictEqual(cell.languageId, 'markdown'); + }); + + test('handles undefined content', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + id: 'agent-block-789', + sortingKey: 'a2', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, ''); + assert.strictEqual(cell.languageId, 'markdown'); + }); + + test('preserves multiline prompt', () => { + const prompt = dedent` + You are a senior data analyst. + + Perform a thorough exploratory analysis: + 1. Create a grouped bar chart of revenue by quarter + 2. Create a line chart showing churn rate trends + 3. Compute a pivot table of average revenue + `; + + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: prompt, + id: 'agent-block-multiline', + sortingKey: 'a3', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, prompt); + assert.strictEqual(cell.languageId, 'markdown'); + }); + + test('preserves agent block with metadata', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Analyze the data', + id: 'agent-block-with-metadata', + metadata: { + deepnote_agent_model: 'gpt-4o' + }, + sortingKey: 'a4', + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, 'Analyze the data'); + assert.strictEqual(cell.languageId, 'markdown'); + }); + }); + + suite('applyChangesToBlock', () => { + test('updates block content from cell value', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Old prompt', + id: 'agent-block-123', + sortingKey: 'a0', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + const cell = new NotebookCellData( + NotebookCellKind.Code, + 'New prompt with updated instructions', + 'markdown' + ); + + converter.applyChangesToBlock(block, cell); + + assert.strictEqual(block.content, 'New prompt with updated instructions'); + }); + + test('handles empty cell value', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Some prompt', + id: 'agent-block-456', + sortingKey: 'a1', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + const cell = new NotebookCellData(NotebookCellKind.Code, '', 'markdown'); + + converter.applyChangesToBlock(block, cell); + + assert.strictEqual(block.content, ''); + }); + + test('does not modify other block properties', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Old prompt', + id: 'agent-block-789', + metadata: { + deepnote_agent_model: 'gpt-4o', + custom: 'value' + }, + sortingKey: 'a2', + type: 'agent' + }; + const cell = new NotebookCellData(NotebookCellKind.Code, 'New prompt', 'markdown'); + + converter.applyChangesToBlock(block, cell); + + assert.strictEqual(block.content, 'New prompt'); + assert.strictEqual(block.id, 'agent-block-789'); + assert.strictEqual(block.type, 'agent'); + assert.strictEqual(block.sortingKey, 'a2'); + assert.deepStrictEqual(block.metadata, { + deepnote_agent_model: 'gpt-4o', + custom: 'value' + }); + }); + }); +}); diff --git a/src/notebooks/deepnote/deepnoteDataConverter.ts b/src/notebooks/deepnote/deepnoteDataConverter.ts index 9f71700a71..51007cae9d 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.ts @@ -11,6 +11,7 @@ import { MarkdownBlockConverter } from './converters/markdownBlockConverter'; import { VisualizationBlockConverter } from './converters/visualizationBlockConverter'; import { compile as convertVegaLiteSpecToVega, ensureVegaLiteLoaded } from './vegaLiteWrapper'; import { produce } from 'immer'; +import { AgentBlockConverter } from './converters/agentBlockConverter'; import { SqlBlockConverter } from './converters/sqlBlockConverter'; import { TextBlockConverter } from './converters/textBlockConverter'; // @ts-ignore - types_unstable subpath requires moduleResolution: "node16" which mandates module: "node16" and .js extensions on all imports @@ -38,6 +39,7 @@ export class DeepnoteDataConverter { private readonly registry = new ConverterRegistry(); constructor() { + this.registry.register(new AgentBlockConverter()); this.registry.register(new CodeBlockConverter()); this.registry.register(new MarkdownBlockConverter()); this.registry.register(new ChartBigNumberBlockConverter()); diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts new file mode 100644 index 0000000000..5c913700ff --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts @@ -0,0 +1,123 @@ +import { + Disposable, + NotebookCell, + NotebookDocument, + OverviewRulerLane, + Range, + TextEditor, + TextEditorDecorationType, + ThemeColor, + window, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; + +const NOTEBOOK_CELL_SCHEME = 'vscode-notebook-cell'; + +/** + * Applies visual decorations (left border, background tint, reduced opacity) to + * code cell editors that belong to ephemeral blocks (`is_ephemeral: true`). + * + * The left border is rendered via a `before` pseudo-element on each line, + * which avoids overlapping or shifting the code text. + * + * Markup cells are handled separately by the markdown-it renderer plugin in + * `src/renderers/client/markdown.ts`. + */ +@injectable() +export class EphemeralCellDecorationProvider implements IExtensionSyncActivationService { + private readonly disposables: Disposable[] = []; + + private ephemeralDecorationType!: TextEditorDecorationType; + + public activate(): void { + this.ephemeralDecorationType = window.createTextEditorDecorationType({ + opacity: '0.8', + isWholeLine: true, + overviewRulerColor: new ThemeColor('charts.yellow'), + overviewRulerLane: OverviewRulerLane.Left, + before: { + contentText: '\u200B', + width: '3px', + backgroundColor: new ThemeColor('charts.yellow'), + margin: '0 8px 0 0' + } + }); + + this.disposables.push(this.ephemeralDecorationType); + + this.disposables.push( + window.onDidChangeVisibleTextEditors(() => { + this.updateDecorations(); + }) + ); + + this.disposables.push( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook.notebookType === 'deepnote') { + this.updateDecorations(); + } + }) + ); + + this.updateDecorations(); + } + + public dispose(): void { + this.disposables.forEach((d) => d.dispose()); + } + + private findCellForEditor(editor: TextEditor): NotebookCell | undefined { + const uri = editor.document.uri; + if (uri.scheme !== NOTEBOOK_CELL_SCHEME) { + return undefined; + } + + for (const notebook of workspace.notebookDocuments) { + if (notebook.notebookType !== 'deepnote') { + continue; + } + + const cell = this.findMatchingCell(notebook, editor); + if (cell) { + return cell; + } + } + + return undefined; + } + + private findMatchingCell(notebook: NotebookDocument, editor: TextEditor): NotebookCell | undefined { + for (const cell of notebook.getCells()) { + if (cell.document.uri.toString() === editor.document.uri.toString()) { + return cell; + } + } + + return undefined; + } + + private updateDecorations(): void { + for (const editor of window.visibleTextEditors) { + if (editor.document.uri.scheme !== NOTEBOOK_CELL_SCHEME) { + continue; + } + + const cell = this.findCellForEditor(editor); + if (!cell || cell.metadata?.is_ephemeral !== true) { + editor.setDecorations(this.ephemeralDecorationType, []); + continue; + } + + const lineRanges: Range[] = []; + for (let i = 0; i < editor.document.lineCount; i++) { + const line = editor.document.lineAt(i); + lineRanges.push(line.range); + } + + editor.setDecorations(this.ephemeralDecorationType, lineRanges); + } + } +} diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts new file mode 100644 index 0000000000..7089391e7c --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts @@ -0,0 +1,83 @@ +import { + CancellationToken, + Disposable, + EventEmitter, + NotebookCell, + NotebookCellStatusBarItem, + NotebookCellStatusBarItemProvider, + l10n, + notebooks, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; + +const EPHEMERAL_INDICATOR_PRIORITY = 1000; + +/** + * Provides a status bar indicator for ephemeral cells — blocks that were + * auto-generated by an agent and marked with `is_ephemeral: true` in metadata. + */ +@injectable() +export class EphemeralCellStatusBarProvider + implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService +{ + private readonly disposables: Disposable[] = []; + private readonly _onDidChangeCellStatusBarItems = new EventEmitter(); + + public readonly onDidChangeCellStatusBarItems = this._onDidChangeCellStatusBarItems.event; + + public activate(): void { + this.disposables.push(notebooks.registerNotebookCellStatusBarItemProvider('deepnote', this)); + + this.disposables.push( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook.notebookType === 'deepnote') { + this._onDidChangeCellStatusBarItems.fire(); + } + }) + ); + + this.disposables.push(this._onDidChangeCellStatusBarItems); + } + + public dispose(): void { + this.disposables.forEach((d) => d.dispose()); + } + + public provideCellStatusBarItems( + cell: NotebookCell, + token: CancellationToken + ): NotebookCellStatusBarItem | undefined { + if (token.isCancellationRequested) { + return undefined; + } + + if (!this.isEphemeralCell(cell)) { + return undefined; + } + + const agentSourceBlockId = cell.metadata?.agent_source_block_id as string | undefined; + + return this.createEphemeralIndicatorItem(agentSourceBlockId); + } + + private createEphemeralIndicatorItem(agentSourceBlockId?: string): NotebookCellStatusBarItem { + const tooltipLines = [l10n.t('Auto-generated ephemeral block')]; + if (agentSourceBlockId) { + tooltipLines.push(l10n.t('Source agent block: {0}', agentSourceBlockId)); + } + + return { + text: `$(sparkle) ${l10n.t('Ephemeral')}`, + alignment: 1, + priority: EPHEMERAL_INDICATOR_PRIORITY, + tooltip: tooltipLines.join('\n') + }; + } + + private isEphemeralCell(cell: NotebookCell): boolean { + return cell.metadata?.is_ephemeral === true; + } +} diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts new file mode 100644 index 0000000000..27e53dcdf2 --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts @@ -0,0 +1,169 @@ +import { expect } from 'chai'; +import { CancellationToken } from 'vscode'; + +import { EphemeralCellStatusBarProvider } from './ephemeralCellStatusBarProvider'; +import { createMockCell } from './deepnoteTestHelpers'; + +suite('EphemeralCellStatusBarProvider', () => { + let provider: EphemeralCellStatusBarProvider; + let mockToken: CancellationToken; + + setup(() => { + mockToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + provider = new EphemeralCellStatusBarProvider(); + }); + + teardown(() => { + provider.dispose(); + }); + + suite('Ephemeral Cell Detection', () => { + test('Should return a status bar item for ephemeral cell', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + + test('Should return undefined when is_ephemeral is false', () => { + const cell = createMockCell({ metadata: { is_ephemeral: false } }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined when is_ephemeral is not set', () => { + const cell = createMockCell({ metadata: {} }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined for cell without metadata', () => { + const cell = createMockCell({ metadata: undefined }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined when is_ephemeral is a non-boolean truthy value', () => { + const cell = createMockCell({ metadata: { is_ephemeral: 'true' } }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined when cancellation is requested', () => { + const cancelledToken: CancellationToken = { + isCancellationRequested: true, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, cancelledToken); + + expect(item).to.be.undefined; + }); + }); + + suite('Status Bar Item Properties', () => { + test('Should display sparkle icon with Ephemeral label', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.text).to.include('$(sparkle)'); + expect(item.text).to.include('Ephemeral'); + }); + + test('Should have left alignment', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.alignment).to.equal(1); + }); + + test('Should have priority 1000 to appear before all other items', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.priority).to.equal(1000); + }); + + test('Should not have a command', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.command).to.be.undefined; + }); + }); + + suite('Tooltip', () => { + test('Should include auto-generated description in tooltip', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.tooltip).to.include('Auto-generated ephemeral block'); + }); + + test('Should include agent source block ID in tooltip when present', () => { + const cell = createMockCell({ + metadata: { + is_ephemeral: true, + agent_source_block_id: 'a0000000000000000000000000000004' + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.tooltip).to.include('a0000000000000000000000000000004'); + expect(item.tooltip).to.include('Source agent block'); + }); + + test('Should not include source block line in tooltip when agent_source_block_id is absent', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.tooltip).to.not.include('Source agent block'); + }); + }); + + suite('Coexistence with other cell types', () => { + test('Should return item for ephemeral agent cell', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + is_ephemeral: true, + agent_source_block_id: 'source-id' + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + + test('Should return item for ephemeral code cell', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'code' }, + is_ephemeral: true + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + + test('Should return item for ephemeral markdown cell', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'markdown' }, + is_ephemeral: true + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + }); +}); diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index cbd8b860fe..5de4c2c4d9 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -85,7 +85,10 @@ import { DeepnoteExtensionSidecarWriter } from '../kernels/deepnote/environments import { DeepnoteNotebookEnvironmentMapper } from '../kernels/deepnote/environments/deepnoteNotebookEnvironmentMapper.node'; import { DeepnoteNotebookCommandListener } from './deepnote/deepnoteNotebookCommandListener'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; +import { AgentCellStatusBarProvider } from './deepnote/agentCellStatusBarProvider'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; +import { EphemeralCellDecorationProvider } from './deepnote/ephemeralCellDecorationProvider'; +import { EphemeralCellStatusBarProvider } from './deepnote/ephemeralCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIntegrationStartupCodeProvider'; import { DeepnoteCellCopyHandler } from './deepnote/deepnoteCellCopyHandler'; @@ -230,6 +233,18 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellDecorationProvider + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNewCellLanguageService diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index 2488ff73d7..4be669266c 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/src/notebooks/serviceRegistry.web.ts @@ -50,7 +50,10 @@ import { IIntegrationWebviewProvider } from './deepnote/integrations/types'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; +import { AgentCellStatusBarProvider } from './deepnote/agentCellStatusBarProvider'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; +import { EphemeralCellDecorationProvider } from './deepnote/ephemeralCellDecorationProvider'; +import { EphemeralCellStatusBarProvider } from './deepnote/ephemeralCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler'; @@ -125,6 +128,18 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellDecorationProvider + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNewCellLanguageService diff --git a/src/renderers/client/markdown.ts b/src/renderers/client/markdown.ts index b5399de2df..8c69bfd618 100644 --- a/src/renderers/client/markdown.ts +++ b/src/renderers/client/markdown.ts @@ -1,3 +1,5 @@ +import type { ActivationFunction } from 'vscode-notebook-renderer'; + const styleContent = ` .alert { width: auto; @@ -31,13 +33,57 @@ const styleContent = ` background-color: rgb(255,205,210); color: rgb(183,28,28); } + +.ephemeral-cell { + border-left: 3px solid var(--vscode-charts-yellow, #cca700); + padding-left: 8px; + opacity: 0.8; +} +.ephemeral-badge { + display: inline-block; + font-size: 0.75em; + padding: 1px 6px; + border-radius: 3px; + background: var(--vscode-charts-yellow, #cca700); + color: var(--vscode-editor-background, #1e1e1e); + margin-bottom: 4px; + font-weight: 600; + letter-spacing: 0.03em; +} `; -export async function activate() { +export const activate: ActivationFunction = async (ctx) => { const style = document.createElement('style'); style.textContent = styleContent; const template = document.createElement('template'); template.classList.add('markdown-style'); template.content.appendChild(style); document.head.appendChild(template); + + const markdownRenderer = await ctx.getRenderer('vscode.markdown-it-renderer'); + if (markdownRenderer) { + (markdownRenderer as any).extendMarkdownIt((md: any) => { + addEphemeralCellWrapper(md); + }); + } + + return undefined; +}; + +function addEphemeralCellWrapper(md: any): void { + md.core.ruler.push('ephemeral_wrapper', (state: any) => { + const metadata = state.env?.outputItem?.metadata; + if (!metadata || metadata.is_ephemeral !== true) { + return; + } + + const openToken = new state.Token('html_block', '', 0); + openToken.content = '
\u2728 Ephemeral\n'; + + const closeToken = new state.Token('html_block', '', 0); + closeToken.content = '
\n'; + + state.tokens.unshift(openToken); + state.tokens.push(closeToken); + }); } From 957fdcdc43b55f014fbcc4b21763b716457e57b2 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 12 Mar 2026 11:32:19 +0000 Subject: [PATCH 05/19] Add a dummy agent block execution handler --- .../controllers/vscodeNotebookController.ts | 26 +- .../deepnote/agentCellExecutionHandler.ts | 51 ++++ .../agentCellExecutionHandler.unit.test.ts | 257 ++++++++++++++++++ .../converters/agentBlockConverter.ts | 8 +- .../agentBlockConverter.unit.test.ts | 18 +- .../deepnoteKernelAutoSelector.node.ts | 32 ++- 6 files changed, 368 insertions(+), 24 deletions(-) create mode 100644 src/notebooks/deepnote/agentCellExecutionHandler.ts create mode 100644 src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index eeeb2614e8..8b4f86db78 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -90,6 +90,7 @@ import { RemoteKernelReconnectBusyIndicator } from './remoteKernelReconnectBusyI import { IConnectionDisplayData, IConnectionDisplayDataProvider, IVSCodeNotebookController } from './types'; import { notebookPathToDeepnoteProjectFilePath } from '../../platform/deepnote/deepnoteProjectUtils'; import { DEEPNOTE_NOTEBOOK_TYPE, IDeepnoteKernelAutoSelector } from '../../kernels/deepnote/types'; +import { executeAgentCell, isAgentCell } from '../deepnote/agentCellExecutionHandler'; /** * Our implementation of the VSCode Notebook Controller. Called by VS code to execute cells in a notebook. Also displayed @@ -626,16 +627,29 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont // Start execution now (from the user's point of view) // Creating these execution objects marks the cell as queued for execution (vscode will update cell UI). type CellExec = { cell: NotebookCell; exec: NotebookCellExecution }; - const cellExecs: CellExec[] = (this.cellQueue.get(doc) || []).map((cell) => { - const exec = this.createCellExecutionIfNecessary(cell, new KernelController(this.controller)); - return { cell, exec }; - }); + const allCells = this.cellQueue.get(doc) || []; this.cellQueue.delete(doc); - const firstCell = cellExecs.length ? cellExecs[0].cell : undefined; - if (!firstCell) { + + const agentCells = allCells.filter((cell) => isAgentCell(cell)); + const kernelCells = allCells.filter((cell) => !isAgentCell(cell)); + + // Execute agent cells directly without kernel involvement + if (agentCells.length > 0) { + logger.trace(`Executing ${agentCells.length} agent cell(s) for ${getDisplayPath(doc.uri)} without kernel`); + await Promise.all(agentCells.map((cell) => executeAgentCell(cell, this.controller))).catch(noop); + } + + if (kernelCells.length === 0) { return; } + const cellExecs: CellExec[] = kernelCells.map((cell) => { + const exec = this.createCellExecutionIfNecessary(cell, new KernelController(this.controller)); + return { cell, exec }; + }); + + const firstCell = cellExecs[0].cell; + logger.trace(`Execute Notebook ${getDisplayPath(doc.uri)}. Step 1`); // Connect to a matching kernel if possible (but user may pick a different one) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts new file mode 100644 index 0000000000..51ab4b0ce5 --- /dev/null +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -0,0 +1,51 @@ +import { NotebookCell, NotebookCellOutput, NotebookCellOutputItem, NotebookController } from 'vscode'; + +import type { Pocket } from '../../platform/deepnote/pocket'; +import { logger } from '../../platform/logging'; + +export function isAgentCell(cell: NotebookCell): boolean { + const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; + + return pocket?.type === 'agent'; +} + +export async function executeAgentCell(cell: NotebookCell, controller: NotebookController): Promise { + const execution = controller.createNotebookCellExecution(cell); + execution.start(Date.now()); + + try { + await execution.clearOutput(); + const prompt = cell.document.getText(); + + const output = new NotebookCellOutput([ + NotebookCellOutputItem.text(`[Agent] Received prompt (${prompt.length} chars)...\n`) + ]); + await execution.replaceOutput([output]); + + const chunks = [ + { delay: 500, text: '[Agent] Analyzing prompt...\n' }, + { delay: 1000, text: '[Agent] Generating plan...\n' }, + { delay: 2000, text: '[Agent] Executing steps...\n' }, + { delay: 3000, text: `[Agent] Done.\n\nPrompt: ${prompt}\n` } + ]; + + let accumulated = `[Agent] Received prompt (${prompt.length} chars)...\n`; + for (const chunk of chunks) { + await delay(chunk.delay); + accumulated += chunk.text; + await execution.replaceOutputItems(NotebookCellOutputItem.text(accumulated), output); + } + + execution.end(true, Date.now()); + } catch (error) { + logger.error('Agent cell execution failed', error); + const message = error instanceof Error ? error.message : String(error); + const stderrOutput = new NotebookCellOutput([NotebookCellOutputItem.stderr(message)]); + await execution.replaceOutput([stderrOutput]).then(undefined, () => undefined); + execution.end(false, Date.now()); + } +} + +function delay(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts new file mode 100644 index 0000000000..77544aeafd --- /dev/null +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -0,0 +1,257 @@ +import { expect } from 'chai'; +import * as sinon from 'sinon'; +import { NotebookCellOutput, NotebookCellOutputItem, NotebookController } from 'vscode'; + +import { executeAgentCell, isAgentCell } from './agentCellExecutionHandler'; +import { createMockCell } from './deepnoteTestHelpers'; + +suite('AgentCellExecutionHandler', () => { + suite('isAgentCell', () => { + test('returns true for cell with agent pocket type', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + + expect(isAgentCell(cell)).to.be.true; + }); + + test('returns false for cell with code pocket type', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } }); + + expect(isAgentCell(cell)).to.be.false; + }); + + test('returns false for cell with markdown pocket type', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'markdown' } } }); + + expect(isAgentCell(cell)).to.be.false; + }); + + test('returns false for cell without pocket', () => { + const cell = createMockCell({ metadata: {} }); + + expect(isAgentCell(cell)).to.be.false; + }); + + test('returns false for cell without metadata', () => { + const cell = createMockCell({ metadata: undefined }); + + expect(isAgentCell(cell)).to.be.false; + }); + }); + + suite('executeAgentCell', () => { + let clock: sinon.SinonFakeTimers; + let mockExecution: { + clearOutput: sinon.SinonStub; + end: sinon.SinonStub; + replaceOutput: sinon.SinonStub; + replaceOutputItems: sinon.SinonStub; + start: sinon.SinonStub; + }; + let mockController: NotebookController; + + setup(() => { + clock = sinon.useFakeTimers(); + + mockExecution = { + clearOutput: sinon.stub().resolves(), + end: sinon.stub(), + replaceOutput: sinon.stub().resolves(), + replaceOutputItems: sinon.stub().resolves(), + start: sinon.stub() + }; + + mockController = { + createNotebookCellExecution: sinon.stub().returns(mockExecution) + } as unknown as NotebookController; + }); + + teardown(() => { + clock.restore(); + }); + + async function runToCompletion(promise: Promise): Promise { + // Total delay across all chunks: 500 + 1000 + 2000 + 3000 = 6500ms + await clock.tickAsync(7000); + await promise; + } + + test('creates execution and starts it', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Analyze data' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect((mockController.createNotebookCellExecution as sinon.SinonStub).calledOnceWith(cell)).to.be.true; + expect(mockExecution.start.calledOnce).to.be.true; + }); + + test('clears output before streaming', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Analyze data' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.clearOutput.calledOnce).to.be.true; + expect(mockExecution.clearOutput.calledBefore(mockExecution.replaceOutput)).to.be.true; + }); + + test('sets initial output via replaceOutput', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Hello world' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.replaceOutput.calledOnce).to.be.true; + + const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; + expect(outputs).to.have.lengthOf(1); + expect(outputs[0].items).to.have.lengthOf(1); + + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('[Agent] Received prompt (11 chars)'); + }); + + test('streams 4 chunks via replaceOutputItems', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test prompt' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.replaceOutputItems.callCount).to.equal(4); + }); + + test('streaming chunks accumulate text progressively', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + const getChunkText = (callIndex: number): string => { + const item = mockExecution.replaceOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; + + return Buffer.from(item.data).toString('utf-8'); + }; + + const chunk1 = getChunkText(0); + const chunk2 = getChunkText(1); + const chunk3 = getChunkText(2); + const chunk4 = getChunkText(3); + + expect(chunk1).to.include('Analyzing prompt'); + expect(chunk2).to.include('Generating plan'); + expect(chunk2).to.include('Analyzing prompt'); + expect(chunk3).to.include('Executing steps'); + expect(chunk3).to.include('Generating plan'); + expect(chunk4).to.include('Done'); + expect(chunk4).to.include('Prompt: Test'); + }); + + test('streaming chunks fire at correct intervals', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test' + }); + + const promise = executeAgentCell(cell, mockController); + + expect(mockExecution.replaceOutputItems.callCount).to.equal(0); + + await clock.tickAsync(500); + expect(mockExecution.replaceOutputItems.callCount).to.equal(1); + + await clock.tickAsync(1000); + expect(mockExecution.replaceOutputItems.callCount).to.equal(2); + + await clock.tickAsync(2000); + expect(mockExecution.replaceOutputItems.callCount).to.equal(3); + + await clock.tickAsync(3000); + expect(mockExecution.replaceOutputItems.callCount).to.equal(4); + + await promise; + }); + + test('ends execution with success', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.true; + }); + + test('ends execution with failure when error occurs', async () => { + mockExecution.clearOutput.rejects(new Error('Test error')); + + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.false; + }); + + test('writes error message to stderr output on failure', async () => { + mockExecution.clearOutput.rejects(new Error('Something went wrong')); + + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.replaceOutput.calledOnce).to.be.true; + + const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; + expect(outputs).to.have.lengthOf(1); + + const item = outputs[0].items[0]; + expect(item.mime).to.equal('application/vnd.code.notebook.stderr'); + + const text = Buffer.from(item.data).toString('utf-8'); + expect(text).to.equal('Something went wrong'); + }); + + test('handles empty prompt', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: '' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.true; + + const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('(0 chars)'); + }); + }); +}); diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.ts b/src/notebooks/deepnote/converters/agentBlockConverter.ts index 357ab5d3c9..6f9ebbd31c 100644 --- a/src/notebooks/deepnote/converters/agentBlockConverter.ts +++ b/src/notebooks/deepnote/converters/agentBlockConverter.ts @@ -6,9 +6,9 @@ import type { BlockConverter } from './blockConverter'; /** * Converter for agent blocks. * - * Agent blocks are rendered as code cells with markdown language so the natural-language - * prompt gets reasonable syntax highlighting while remaining visually distinct from - * Python code blocks. The prompt text is stored in `block.content`. + * Agent blocks are rendered as code cells with plaintext language so the + * natural-language prompt appears without syntax highlighting while remaining + * executable. The prompt text is stored in `block.content`. * * Agent-specific metadata (model, MCP servers, max iterations, etc.) is preserved * through the generic metadata pass-through in DeepnoteDataConverter. @@ -23,7 +23,7 @@ export class AgentBlockConverter implements BlockConverter { } convertToCell(block: DeepnoteBlock): NotebookCellData { - const cell = new NotebookCellData(NotebookCellKind.Code, block.content || '', 'markdown'); + const cell = new NotebookCellData(NotebookCellKind.Code, block.content || '', 'plaintext'); return cell; } diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts index 43fac425db..a3ce26acf8 100644 --- a/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts +++ b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts @@ -36,7 +36,7 @@ suite('AgentBlockConverter', () => { }); suite('convertToCell', () => { - test('converts agent block to code cell with markdown language', () => { + test('converts agent block to code cell with plaintext language', () => { const block: DeepnoteBlock = { blockGroup: 'test-group', content: 'Analyze the dataset and create a summary report', @@ -50,7 +50,7 @@ suite('AgentBlockConverter', () => { assert.strictEqual(cell.kind, NotebookCellKind.Code); assert.strictEqual(cell.value, 'Analyze the dataset and create a summary report'); - assert.strictEqual(cell.languageId, 'markdown'); + assert.strictEqual(cell.languageId, 'plaintext'); }); test('handles empty content', () => { @@ -67,7 +67,7 @@ suite('AgentBlockConverter', () => { assert.strictEqual(cell.kind, NotebookCellKind.Code); assert.strictEqual(cell.value, ''); - assert.strictEqual(cell.languageId, 'markdown'); + assert.strictEqual(cell.languageId, 'plaintext'); }); test('handles undefined content', () => { @@ -83,7 +83,7 @@ suite('AgentBlockConverter', () => { assert.strictEqual(cell.kind, NotebookCellKind.Code); assert.strictEqual(cell.value, ''); - assert.strictEqual(cell.languageId, 'markdown'); + assert.strictEqual(cell.languageId, 'plaintext'); }); test('preserves multiline prompt', () => { @@ -109,7 +109,7 @@ suite('AgentBlockConverter', () => { assert.strictEqual(cell.kind, NotebookCellKind.Code); assert.strictEqual(cell.value, prompt); - assert.strictEqual(cell.languageId, 'markdown'); + assert.strictEqual(cell.languageId, 'plaintext'); }); test('preserves agent block with metadata', () => { @@ -128,7 +128,7 @@ suite('AgentBlockConverter', () => { assert.strictEqual(cell.kind, NotebookCellKind.Code); assert.strictEqual(cell.value, 'Analyze the data'); - assert.strictEqual(cell.languageId, 'markdown'); + assert.strictEqual(cell.languageId, 'plaintext'); }); }); @@ -145,7 +145,7 @@ suite('AgentBlockConverter', () => { const cell = new NotebookCellData( NotebookCellKind.Code, 'New prompt with updated instructions', - 'markdown' + 'plaintext' ); converter.applyChangesToBlock(block, cell); @@ -162,7 +162,7 @@ suite('AgentBlockConverter', () => { metadata: { deepnote_agent_model: 'auto' }, type: 'agent' }; - const cell = new NotebookCellData(NotebookCellKind.Code, '', 'markdown'); + const cell = new NotebookCellData(NotebookCellKind.Code, '', 'plaintext'); converter.applyChangesToBlock(block, cell); @@ -181,7 +181,7 @@ suite('AgentBlockConverter', () => { sortingKey: 'a2', type: 'agent' }; - const cell = new NotebookCellData(NotebookCellKind.Code, 'New prompt', 'markdown'); + const cell = new NotebookCellData(NotebookCellKind.Code, 'New prompt', 'plaintext'); converter.applyChangesToBlock(block, cell); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 63e7129200..2a5964b6c8 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -56,6 +56,7 @@ import { logger } from '../../platform/logging'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { IControllerRegistration, IVSCodeNotebookController } from '../controllers/types'; import { IDeepnoteNotebookManager } from '../types'; +import { executeAgentCell, isAgentCell } from './agentCellExecutionHandler'; import { IDeepnoteInitNotebookRunner } from './deepnoteInitNotebookRunner.node'; import { computeRequirementsHash } from './deepnoteProjectUtils'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; @@ -1204,7 +1205,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, ); controller.supportsExecutionOrder = true; - controller.supportedLanguages = ['python', 'sql', 'markdown']; + controller.supportedLanguages = ['python', 'sql', 'markdown', 'plaintext']; // Execution handler that shows environment picker when user tries to run without an environment controller.executeHandler = async (cells, doc) => { @@ -1214,6 +1215,28 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } cells` ); + const agentCells = cells.filter((cell) => isAgentCell(cell)); + const kernelCells = cells.filter((cell) => !isAgentCell(cell)); + + // Execute agent cells directly without kernel involvement + if (agentCells.length > 0) { + logger.info( + `Executing ${agentCells.length} agent cell(s) for ${getDisplayPath(doc.uri)} without kernel` + ); + + for (const cell of agentCells) { + try { + await executeAgentCell(cell, controller); + } catch (cellError) { + logger.error(`Error executing agent cell ${cell.index}`, cellError); + } + } + } + + if (kernelCells.length === 0) { + return; + } + // Create a cancellation token that cancels when the notebook is closed const cts = new CancellationTokenSource(); const closeListener = workspace.onDidCloseNotebookDocument((closedDoc) => { @@ -1242,7 +1265,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return; } - logger.info(`Executing ${cells.length} cells through kernel after environment configuration`); + logger.info(`Executing ${kernelCells.length} cells through kernel after environment configuration`); // Get or create a kernel for this notebook with the new connection const kernel = this.kernelProvider.getOrCreate(doc, { @@ -1254,16 +1277,15 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Execute cells through the kernel const kernelExecution = this.kernelProvider.getKernelExecution(kernel); - for (const cell of cells) { + for (const cell of kernelCells) { try { await kernelExecution.executeCell(cell); } catch (cellError) { logger.error(`Error executing cell ${cell.index}`, cellError); - // Continue with remaining cells } } - logger.info(`Finished executing ${cells.length} cells`); + logger.info(`Finished executing ${kernelCells.length} cells`); } catch (error) { if (isCancellationError(error)) { logger.info(`Environment setup cancelled for ${getDisplayPath(doc.uri)}`); From 9fbe622953b843ad4c9b57c6663041188d1f5bf9 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 13 Mar 2026 13:13:55 +0000 Subject: [PATCH 06/19] feat(agent-block): Integrate deepnote runtime-core to execute Agent blocks --- .../deepnote/agentCellExecutionHandler.ts | 296 +++++++++++++++++- .../deepnote/deepnoteDataConverter.ts | 2 +- 2 files changed, 281 insertions(+), 17 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index 51ab4b0ce5..6de5454b4a 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -1,7 +1,32 @@ -import { NotebookCell, NotebookCellOutput, NotebookCellOutputItem, NotebookController } from 'vscode'; +import { + NotebookCell, + NotebookCellOutput, + NotebookCellOutputItem, + NotebookController, + NotebookDocument, + NotebookEdit, + NotebookRange, + WorkspaceEdit, + commands, + workspace +} from 'vscode'; +import { AgentBlock, DeepnoteBlock } from '@deepnote/blocks'; +import { + AgentBlockContext, + AgentStreamEvent, + executeAgentBlock, + serializeNotebookContextFromBlocks +} from '@deepnote/runtime-core'; + +import { translateCellDisplayOutput } from '../../kernels/execution/helpers'; +import { createDeferred } from '../../platform/common/utils/async'; +import { uuidUtils } from '../../platform/common/uuid'; import type { Pocket } from '../../platform/deepnote/pocket'; import { logger } from '../../platform/logging'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; +import { generateBlockId, generateSortingKey } from './dataConversionUtils'; +import { DeepnoteDataConverter } from './deepnoteDataConverter'; export function isAgentCell(cell: NotebookCell): boolean { const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; @@ -9,43 +34,282 @@ export function isAgentCell(cell: NotebookCell): boolean { return pocket?.type === 'agent'; } +export function serializeNotebookContext({ cells }: { cells: NotebookCell[] }): string { + const converter = new DeepnoteDataConverter(); + + const blocks = cells.reduce((acc, cell) => { + try { + const block = converter.convertCellToBlock( + { + kind: cell.kind, + value: cell.document.getText(), + languageId: cell.document.languageId, + metadata: cell.metadata, + outputs: [...(cell.outputs || [])] + }, + cell.index + ); + acc.push(block); + } catch (error) { + logger.error(`Error converting cell to block: ${error}`); + } + return acc; + }, []); + + return serializeNotebookContextFromBlocks({ blocks, notebookName: null }); +} + export async function executeAgentCell(cell: NotebookCell, controller: NotebookController): Promise { const execution = controller.createNotebookCellExecution(cell); execution.start(Date.now()); try { await execution.clearOutput(); + const prompt = cell.document.getText(); - const output = new NotebookCellOutput([ - NotebookCellOutputItem.text(`[Agent] Received prompt (${prompt.length} chars)...\n`) - ]); + let accumulated = `[Agent] Planning next steps...`; + const output = new NotebookCellOutput([NotebookCellOutputItem.text(accumulated)]); await execution.replaceOutput([output]); - const chunks = [ - { delay: 500, text: '[Agent] Analyzing prompt...\n' }, - { delay: 1000, text: '[Agent] Generating plan...\n' }, - { delay: 2000, text: '[Agent] Executing steps...\n' }, - { delay: 3000, text: `[Agent] Done.\n\nPrompt: ${prompt}\n` } - ]; + await removeEphemeralCellsForAgent(cell); - let accumulated = `[Agent] Received prompt (${prompt.length} chars)...\n`; - for (const chunk of chunks) { - await delay(chunk.delay); - accumulated += chunk.text; - await execution.replaceOutputItems(NotebookCellOutputItem.text(accumulated), output); + const dataConverter = new DeepnoteDataConverter(); + const deepnoteBlock = dataConverter.convertCellToBlock( + { + kind: cell.kind, + value: cell.document.getText(), + languageId: cell.document.languageId, + metadata: cell.metadata, + outputs: [...(cell.outputs || [])] + }, + cell.index + ); + const agentBlock: AgentBlock | null = deepnoteBlock.type === 'agent' ? deepnoteBlock : null; + + if (agentBlock == null) { + // TODO: better DX error handling + throw new Error('Cell is not an agent cell'); } + let lastAgentEventType: AgentStreamEvent['type'] | undefined; + + const notebookContext = serializeNotebookContext({ + cells: cell.notebook.getCells().filter((c) => c.index !== cell.index) + }); + + const openAiToken = process.env.OPENAI_API_KEY; + if (openAiToken == null) { + throw new Error('OPENAI_API_KEY is not set'); + } + + const context: AgentBlockContext = { + openAiToken, + mcpServers: [], + notebookContext, + addMarkdownBlock: async ({ content }: { content: string }) => { + await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'markdown', content); + return { success: true }; + }, + addAndExecuteCodeBlock: async ({ code }: { code: string }) => { + const cellIndex = await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'code', code); + + const { success } = await executeEphemeralCell(cell.notebook, cellIndex); + return success ? { success } : { success: false, error: new Error('Ephemeral cell execution failed') }; + }, + onLog: (message: string) => { + logger.info('Agent log', message); + // accumulated += message; + // TODO: replaceOutputItems is Async function + // execution.replaceOutputItems(NotebookCellOutputItem.text(accumulated), output); + }, + onAgentEvent: async (event: AgentStreamEvent) => { + logger.info('Agent event', JSON.stringify(event)); + if (lastAgentEventType != null && lastAgentEventType !== event.type) { + accumulated += `\n\n`; + } + switch (event.type) { + case 'tool_called': + // Ignore calling tool_called events + // accumulated += `[Agent] Tool called: ${event.toolName}`; + break; + case 'tool_output': + accumulated += `[Agent] Tool output: ${event.toolName}`; + accumulated += `[Agent] Tool output length: ${event.output?.length}`; + break; + case 'text_delta': + if (lastAgentEventType !== 'text_delta') { + accumulated += `[Agent] Text:\n`; + } + accumulated += event.text; + break; + case 'reasoning_delta': + if (lastAgentEventType !== 'reasoning_delta') { + accumulated += `[Agent] Reasoning:\n`; + } + accumulated += event.text; + break; + default: + event satisfies never; + } + lastAgentEventType = event.type; + + await execution.replaceOutputItems(NotebookCellOutputItem.text(accumulated), output); + } + }; + + logger.info( + `Agent cell: starting executeAgentBlock, model=${agentBlock.metadata.deepnote_agent_model}, prompt length=${prompt.length}` + ); + const result = await executeAgentBlock(agentBlock, context); + logger.info(`Agent cell: executeAgentBlock completed, finalOutput length=${result.finalOutput.length}`); + execution.end(true, Date.now()); } catch (error) { logger.error('Agent cell execution failed', error); + if (error instanceof Error) { + logger.error(`Agent error name=${error.name}, message=${error.message}`); + if (error.cause) { + logger.error('Agent error cause:', error.cause); + } + if (error.stack) { + logger.error('Agent error stack:', error.stack); + } + } + const message = error instanceof Error ? error.message : String(error); const stderrOutput = new NotebookCellOutput([NotebookCellOutputItem.stderr(message)]); - await execution.replaceOutput([stderrOutput]).then(undefined, () => undefined); + await execution.appendOutput([stderrOutput]).then(undefined, () => undefined); execution.end(false, Date.now()); } } +function getInsertIndexAfterAgentCell( + notebook: NotebookDocument, + agentCellIndex: number, + agentBlockId: string +): number { + let index = agentCellIndex + 1; + + while (index < notebook.cellCount) { + const cell = notebook.cellAt(index); + if (cell.metadata?.is_ephemeral === true && cell.metadata?.agent_source_block_id === agentBlockId) { + index++; + } else { + break; + } + } + + return index; +} + +async function insertEphemeralCell( + notebook: NotebookDocument, + agentCellIndex: number, + agentBlockId: string, + blockType: 'code' | 'markdown', + content: string +): Promise { + const insertIndex = getInsertIndexAfterAgentCell(notebook, agentCellIndex, agentBlockId); + + const block: DeepnoteBlock = { + type: blockType, + id: generateBlockId(), + blockGroup: uuidUtils.generateUuid(), + sortingKey: generateSortingKey(insertIndex), + content, + metadata: { + is_ephemeral: true, + agent_source_block_id: agentBlockId + } + }; + + const converter = new DeepnoteDataConverter(); + const [cellData] = converter.convertBlocksToCells([block]); + + const edit = new WorkspaceEdit(); + edit.set(notebook.uri, [NotebookEdit.insertCells(insertIndex, [cellData])]); + await workspace.applyEdit(edit); + + return insertIndex; +} + +const EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS = 5 * 60 * 1000; + +async function executeEphemeralCell( + notebook: NotebookDocument, + cellIndex: number +): Promise<{ success: boolean; outputs: unknown[]; executionCount: number | null }> { + const cell = notebook.cellAt(cellIndex); + const completionDeferred = createDeferred(); + + const disposable = notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell === cell && e.state === NotebookCellExecutionState.Idle) { + completionDeferred.resolve(); + } + }); + + const timeout = setTimeout(() => { + completionDeferred.reject(new Error('Ephemeral cell execution timed out')); + }, EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS); + + try { + await commands.executeCommand('notebook.cell.execute', { + ranges: [{ start: cellIndex, end: cellIndex + 1 }], + document: notebook.uri + }); + + await completionDeferred.promise; + + return { + success: cell.executionSummary?.success !== false, + outputs: cell.outputs.map(translateCellDisplayOutput), + executionCount: cell.executionSummary?.executionOrder ?? null + }; + } catch (error) { + return { + success: false, + outputs: [], + executionCount: null + }; + } finally { + disposable.dispose(); + clearTimeout(timeout); + } +} + +async function removeEphemeralCellsForAgent(agentCell: NotebookCell): Promise { + const agentBlockId = (agentCell.metadata?.id ?? agentCell.metadata?.__deepnoteBlockId) as string | undefined; + if (!agentBlockId) { + return; + } + + const notebook = agentCell.notebook; + const deletions: NotebookEdit[] = []; + + for (let i = notebook.cellCount - 1; i >= 0; i--) { + const cell = notebook.cellAt(i); + + if (cell.metadata?.is_ephemeral === true && cell.metadata?.agent_source_block_id === agentBlockId) { + deletions.push(NotebookEdit.deleteCells(new NotebookRange(i, i + 1))); + } + } + + if (deletions.length === 0) { + return; + } + + const edit = new WorkspaceEdit(); + edit.set(notebook.uri, deletions); + + const success = await workspace.applyEdit(edit); + if (success) { + logger.info(`Removed ${deletions.length} ephemeral cell(s) for agent block ${agentBlockId}`); + } else { + logger.warn(`Failed to remove ephemeral cells for agent block ${agentBlockId}`); + } +} + function delay(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } diff --git a/src/notebooks/deepnote/deepnoteDataConverter.ts b/src/notebooks/deepnote/deepnoteDataConverter.ts index 51007cae9d..d6ef93b52f 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.ts @@ -1,5 +1,5 @@ import { isExecutableBlock, type DeepnoteBlock } from '@deepnote/blocks'; -import { NotebookCellData, NotebookCellKind, NotebookCellOutput, NotebookCellOutputItem } from 'vscode'; +import { NotebookCell, NotebookCellData, NotebookCellKind, NotebookCellOutput, NotebookCellOutputItem } from 'vscode'; import { generateBlockId, generateSortingKey } from './dataConversionUtils'; import type { DeepnoteOutput } from '../../platform/deepnote/deepnoteTypes'; From 46f9a4c1184229c9d750d17373f91f3cda76c7c9 Mon Sep 17 00:00:00 2001 From: tomas Date: Sat, 14 Mar 2026 09:41:46 +0000 Subject: [PATCH 07/19] feat(agent-cell): Enhance executeAgentCell with options for custom execution functions and improve ephemeral cell handling --- .../deepnote/agentCellExecutionHandler.ts | 46 ++-- .../agentCellExecutionHandler.unit.test.ts | 208 ++++++++++-------- .../deepnote/deepnoteDataConverter.ts | 2 +- src/notebooks/deepnote/deepnoteTestHelpers.ts | 15 +- 4 files changed, 151 insertions(+), 120 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index 6de5454b4a..b84ee63255 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -59,7 +59,16 @@ export function serializeNotebookContext({ cells }: { cells: NotebookCell[] }): return serializeNotebookContextFromBlocks({ blocks, notebookName: null }); } -export async function executeAgentCell(cell: NotebookCell, controller: NotebookController): Promise { +export interface ExecuteAgentCellOptions { + executeAgentBlockFn?: typeof executeAgentBlock; +} + +export async function executeAgentCell( + cell: NotebookCell, + controller: NotebookController, + options?: ExecuteAgentCellOptions +): Promise { + const executeAgentBlockFn = options?.executeAgentBlockFn ?? executeAgentBlock; const execution = controller.createNotebookCellExecution(cell); execution.start(Date.now()); @@ -72,8 +81,6 @@ export async function executeAgentCell(cell: NotebookCell, controller: NotebookC const output = new NotebookCellOutput([NotebookCellOutputItem.text(accumulated)]); await execution.replaceOutput([output]); - await removeEphemeralCellsForAgent(cell); - const dataConverter = new DeepnoteDataConverter(); const deepnoteBlock = dataConverter.convertCellToBlock( { @@ -92,6 +99,8 @@ export async function executeAgentCell(cell: NotebookCell, controller: NotebookC throw new Error('Cell is not an agent cell'); } + await removeEphemeralCellsForAgent(cell.notebook, agentBlock.id); + let lastAgentEventType: AgentStreamEvent['type'] | undefined; const notebookContext = serializeNotebookContext({ @@ -113,8 +122,9 @@ export async function executeAgentCell(cell: NotebookCell, controller: NotebookC }, addAndExecuteCodeBlock: async ({ code }: { code: string }) => { const cellIndex = await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'code', code); + const insertedCell = cell.notebook.cellAt(cellIndex); - const { success } = await executeEphemeralCell(cell.notebook, cellIndex); + const { success } = await executeEphemeralCell(insertedCell); return success ? { success } : { success: false, error: new Error('Ephemeral cell execution failed') }; }, onLog: (message: string) => { @@ -131,10 +141,10 @@ export async function executeAgentCell(cell: NotebookCell, controller: NotebookC switch (event.type) { case 'tool_called': // Ignore calling tool_called events - // accumulated += `[Agent] Tool called: ${event.toolName}`; + accumulated += `[Agent] Tool called: ${event.toolName}`; break; case 'tool_output': - accumulated += `[Agent] Tool output: ${event.toolName}`; + accumulated += `[Agent] Tool output: ${event.toolName}\n`; accumulated += `[Agent] Tool output length: ${event.output?.length}`; break; case 'text_delta': @@ -161,7 +171,7 @@ export async function executeAgentCell(cell: NotebookCell, controller: NotebookC logger.info( `Agent cell: starting executeAgentBlock, model=${agentBlock.metadata.deepnote_agent_model}, prompt length=${prompt.length}` ); - const result = await executeAgentBlock(agentBlock, context); + const result = await executeAgentBlockFn(agentBlock, context); logger.info(`Agent cell: executeAgentBlock completed, finalOutput length=${result.finalOutput.length}`); execution.end(true, Date.now()); @@ -236,11 +246,9 @@ async function insertEphemeralCell( const EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS = 5 * 60 * 1000; -async function executeEphemeralCell( - notebook: NotebookDocument, - cellIndex: number +export async function executeEphemeralCell( + cell: NotebookCell ): Promise<{ success: boolean; outputs: unknown[]; executionCount: number | null }> { - const cell = notebook.cellAt(cellIndex); const completionDeferred = createDeferred(); const disposable = notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { @@ -254,9 +262,11 @@ async function executeEphemeralCell( }, EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS); try { + const cellIndex = cell.index; + await commands.executeCommand('notebook.cell.execute', { ranges: [{ start: cellIndex, end: cellIndex + 1 }], - document: notebook.uri + document: cell.notebook.uri }); await completionDeferred.promise; @@ -278,13 +288,7 @@ async function executeEphemeralCell( } } -async function removeEphemeralCellsForAgent(agentCell: NotebookCell): Promise { - const agentBlockId = (agentCell.metadata?.id ?? agentCell.metadata?.__deepnoteBlockId) as string | undefined; - if (!agentBlockId) { - return; - } - - const notebook = agentCell.notebook; +async function removeEphemeralCellsForAgent(notebook: NotebookDocument, agentBlockId: string): Promise { const deletions: NotebookEdit[] = []; for (let i = notebook.cellCount - 1; i >= 0; i--) { @@ -309,7 +313,3 @@ async function removeEphemeralCellsForAgent(agentCell: NotebookCell): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)); -} diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index 77544aeafd..b84bb6286c 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -1,8 +1,17 @@ import { expect } from 'chai'; import * as sinon from 'sinon'; +import { anything, capture, reset, when } from 'ts-mockito'; import { NotebookCellOutput, NotebookCellOutputItem, NotebookController } from 'vscode'; -import { executeAgentCell, isAgentCell } from './agentCellExecutionHandler'; +import type { AgentBlock } from '@deepnote/blocks'; +import type { AgentBlockContext, AgentBlockResult } from '@deepnote/runtime-core'; + +import { + NotebookCellExecutionState, + notebookCellExecutions +} from '../../platform/notebooks/cellExecutionStateService'; +import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; +import { executeAgentCell, executeEphemeralCell, isAgentCell } from './agentCellExecutionHandler'; import { createMockCell } from './deepnoteTestHelpers'; suite('AgentCellExecutionHandler', () => { @@ -39,8 +48,8 @@ suite('AgentCellExecutionHandler', () => { }); suite('executeAgentCell', () => { - let clock: sinon.SinonFakeTimers; let mockExecution: { + appendOutput: sinon.SinonStub; clearOutput: sinon.SinonStub; end: sinon.SinonStub; replaceOutput: sinon.SinonStub; @@ -48,11 +57,15 @@ suite('AgentCellExecutionHandler', () => { start: sinon.SinonStub; }; let mockController: NotebookController; + let executeAgentBlockStub: sinon.SinonStub; + let savedOpenAiKey: string | undefined; setup(() => { - clock = sinon.useFakeTimers(); + savedOpenAiKey = process.env.OPENAI_API_KEY; + process.env.OPENAI_API_KEY = 'test-key'; mockExecution = { + appendOutput: sinon.stub().resolves(), clearOutput: sinon.stub().resolves(), end: sinon.stub(), replaceOutput: sinon.stub().resolves(), @@ -63,52 +76,47 @@ suite('AgentCellExecutionHandler', () => { mockController = { createNotebookCellExecution: sinon.stub().returns(mockExecution) } as unknown as NotebookController; + + executeAgentBlockStub = sinon.stub().resolves({ finalOutput: 'done' } as AgentBlockResult); }); teardown(() => { - clock.restore(); + if (savedOpenAiKey !== undefined) { + process.env.OPENAI_API_KEY = savedOpenAiKey; + } else { + delete process.env.OPENAI_API_KEY; + } }); - async function runToCompletion(promise: Promise): Promise { - // Total delay across all chunks: 500 + 1000 + 2000 + 3000 = 6500ms - await clock.tickAsync(7000); - await promise; + function createAgentCell(text: string = 'Test prompt') { + return createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text + }); } test('creates execution and starts it', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Analyze data' - }); + const cell = createAgentCell('Analyze data'); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect((mockController.createNotebookCellExecution as sinon.SinonStub).calledOnceWith(cell)).to.be.true; expect(mockExecution.start.calledOnce).to.be.true; }); test('clears output before streaming', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Analyze data' - }); + const cell = createAgentCell('Analyze data'); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect(mockExecution.clearOutput.calledOnce).to.be.true; expect(mockExecution.clearOutput.calledBefore(mockExecution.replaceOutput)).to.be.true; }); test('sets initial output via replaceOutput', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Hello world' - }); + const cell = createAgentCell('Hello world'); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect(mockExecution.replaceOutput.calledOnce).to.be.true; @@ -117,29 +125,35 @@ suite('AgentCellExecutionHandler', () => { expect(outputs[0].items).to.have.lengthOf(1); const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); - expect(text).to.include('[Agent] Received prompt (11 chars)'); + expect(text).to.include('[Agent] Planning next steps...'); }); - test('streams 4 chunks via replaceOutputItems', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Test prompt' + test('streams events via replaceOutputItems using onAgentEvent callback', async () => { + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.onAgentEvent?.({ type: 'text_delta', text: 'Hello ' }); + await context.onAgentEvent?.({ type: 'text_delta', text: 'world' }); + + return { finalOutput: 'Hello world' } as AgentBlockResult; }); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - expect(mockExecution.replaceOutputItems.callCount).to.equal(4); + expect(mockExecution.replaceOutputItems.callCount).to.equal(2); }); test('streaming chunks accumulate text progressively', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Test' + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.onAgentEvent?.({ type: 'text_delta', text: 'first' }); + await context.onAgentEvent?.({ type: 'text_delta', text: ' second' }); + + return { finalOutput: 'first second' } as AgentBlockResult; }); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); const getChunkText = (callIndex: number): string => { const item = mockExecution.replaceOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; @@ -149,51 +163,39 @@ suite('AgentCellExecutionHandler', () => { const chunk1 = getChunkText(0); const chunk2 = getChunkText(1); - const chunk3 = getChunkText(2); - const chunk4 = getChunkText(3); - - expect(chunk1).to.include('Analyzing prompt'); - expect(chunk2).to.include('Generating plan'); - expect(chunk2).to.include('Analyzing prompt'); - expect(chunk3).to.include('Executing steps'); - expect(chunk3).to.include('Generating plan'); - expect(chunk4).to.include('Done'); - expect(chunk4).to.include('Prompt: Test'); - }); - test('streaming chunks fire at correct intervals', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Test' - }); + expect(chunk1).to.include('[Agent] Text:'); + expect(chunk1).to.include('first'); + expect(chunk2).to.include('first second'); + }); - const promise = executeAgentCell(cell, mockController); + test('separates different event types with blank lines', async () => { + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.onAgentEvent?.({ type: 'text_delta', text: 'thinking...' }); + await context.onAgentEvent?.({ type: 'tool_called', toolName: 'search' }); - expect(mockExecution.replaceOutputItems.callCount).to.equal(0); + return { finalOutput: '' } as AgentBlockResult; + }); - await clock.tickAsync(500); - expect(mockExecution.replaceOutputItems.callCount).to.equal(1); + const cell = createAgentCell(); - await clock.tickAsync(1000); - expect(mockExecution.replaceOutputItems.callCount).to.equal(2); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - await clock.tickAsync(2000); - expect(mockExecution.replaceOutputItems.callCount).to.equal(3); + const getChunkText = (callIndex: number): string => { + const item = mockExecution.replaceOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; - await clock.tickAsync(3000); - expect(mockExecution.replaceOutputItems.callCount).to.equal(4); + return Buffer.from(item.data).toString('utf-8'); + }; - await promise; + const chunk2 = getChunkText(1); + expect(chunk2).to.include('\n\n'); + expect(chunk2).to.include('[Agent] Tool called: search'); }); test('ends execution with success', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Test' - }); + const cell = createAgentCell(); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect(mockExecution.end.calledOnce).to.be.true; expect(mockExecution.end.firstCall.args[0]).to.be.true; @@ -202,13 +204,9 @@ suite('AgentCellExecutionHandler', () => { test('ends execution with failure when error occurs', async () => { mockExecution.clearOutput.rejects(new Error('Test error')); - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Test' - }); + const cell = createAgentCell(); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect(mockExecution.end.calledOnce).to.be.true; expect(mockExecution.end.firstCall.args[0]).to.be.false; @@ -217,17 +215,13 @@ suite('AgentCellExecutionHandler', () => { test('writes error message to stderr output on failure', async () => { mockExecution.clearOutput.rejects(new Error('Something went wrong')); - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Test' - }); + const cell = createAgentCell(); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - expect(mockExecution.replaceOutput.calledOnce).to.be.true; + expect(mockExecution.appendOutput.calledOnce).to.be.true; - const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; + const outputs = mockExecution.appendOutput.firstCall.args[0] as NotebookCellOutput[]; expect(outputs).to.have.lengthOf(1); const item = outputs[0].items[0]; @@ -238,20 +232,48 @@ suite('AgentCellExecutionHandler', () => { }); test('handles empty prompt', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: '' - }); + const cell = createAgentCell(''); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect(mockExecution.end.calledOnce).to.be.true; expect(mockExecution.end.firstCall.args[0]).to.be.true; const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); - expect(text).to.include('(0 chars)'); + expect(text).to.include('[Agent] Planning next steps...'); + }); + }); + + suite('executeEphemeralCell', () => { + teardown(() => { + reset(mockedVSCodeNamespaces.commands); + }); + + test('uses current cell index, not stale index from insertion time', async () => { + const staleIndex = 5; + const currentIndex = 6; + + const cell = createMockCell({ index: staleIndex }); + + // Simulate a concurrent insertion shifting the cell's index + (cell as { index: number }).index = currentIndex; + + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenCall(async () => { + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Idle); + }); + + await executeEphemeralCell(cell); + + const [commandName, commandArg] = capture( + mockedVSCodeNamespaces.commands.executeCommand as (cmd: string, arg: unknown) => Thenable + ).last(); + + expect(commandName).to.equal('notebook.cell.execute'); + expect(commandArg).to.deep.equal({ + ranges: [{ start: currentIndex, end: currentIndex + 1 }], + document: cell.notebook.uri + }); }); }); }); diff --git a/src/notebooks/deepnote/deepnoteDataConverter.ts b/src/notebooks/deepnote/deepnoteDataConverter.ts index d6ef93b52f..51007cae9d 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.ts @@ -1,5 +1,5 @@ import { isExecutableBlock, type DeepnoteBlock } from '@deepnote/blocks'; -import { NotebookCell, NotebookCellData, NotebookCellKind, NotebookCellOutput, NotebookCellOutputItem } from 'vscode'; +import { NotebookCellData, NotebookCellKind, NotebookCellOutput, NotebookCellOutputItem } from 'vscode'; import { generateBlockId, generateSortingKey } from './dataConversionUtils'; import type { DeepnoteOutput } from '../../platform/deepnote/deepnoteTypes'; diff --git a/src/notebooks/deepnote/deepnoteTestHelpers.ts b/src/notebooks/deepnote/deepnoteTestHelpers.ts index 18d03e558d..eb0dc26464 100644 --- a/src/notebooks/deepnote/deepnoteTestHelpers.ts +++ b/src/notebooks/deepnote/deepnoteTestHelpers.ts @@ -47,8 +47,16 @@ export function createMockNotebook(options?: CreateMockNotebookOptions): Noteboo return { uri, notebookType, - metadata - } as NotebookDocument; + metadata, + cellCount: 0, + cellAt: () => ({}) as NotebookCell, + getCells: () => [], + version: 1, + isDirty: false, + isUntitled: false, + isClosed: false, + save: async () => true + } satisfies NotebookDocument; } /** @@ -121,7 +129,8 @@ export function createMockCell(options?: CreateMockCellOptions): NotebookCell { positionAt: () => ({}) as unknown, validateRange: () => ({}) as unknown, validatePosition: () => ({}) as unknown, - getWordRangeAtPosition: () => undefined + getWordRangeAtPosition: () => undefined, + encoding: 'utf-8' } as unknown as TextDocument; return { From 58e653d2d5efc12e750c324e5b509f71e1cf6dd0 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 16 Mar 2026 17:10:52 +0000 Subject: [PATCH 08/19] feat(ephemeral-cells): Introduce isEphemeralCell utility and enhance handling of ephemeral cells in serialization and decoration --- build/esbuild/build.ts | 3 +- .../deepnote/agentCellExecutionHandler.ts | 7 +- src/notebooks/deepnote/dataConversionUtils.ts | 9 +++ src/notebooks/deepnote/deepnoteSerializer.ts | 17 +++-- .../deepnote/deepnoteSerializer.unit.test.ts | 65 +++++++++++++++++++ .../ephemeralCellDecorationProvider.ts | 3 +- .../ephemeralCellStatusBarProvider.ts | 7 +- 7 files changed, 96 insertions(+), 15 deletions(-) diff --git a/build/esbuild/build.ts b/build/esbuild/build.ts index c313ce8cf8..40fdd60cc3 100644 --- a/build/esbuild/build.ts +++ b/build/esbuild/build.ts @@ -72,7 +72,8 @@ const commonExternals = [ const webExternals = [ ...commonExternals, 'canvas', // Native module used by vega for server-side rendering, not needed in browser - 'mathjax-electron' // Uses Node.js path module, MathJax rendering handled differently in browser + 'mathjax-electron', // Uses Node.js path module, MathJax rendering handled differently in browser + '@deepnote/runtime-core' // Uses tcp-port-used → net, only needed in desktop for agent block execution ]; const desktopExternals = [...commonExternals, ...deskTopNodeModulesToExternalize]; const bundleConfig = getBundleConfiguration(); diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index b84ee63255..ea8485e8a4 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -25,7 +25,7 @@ import { uuidUtils } from '../../platform/common/uuid'; import type { Pocket } from '../../platform/deepnote/pocket'; import { logger } from '../../platform/logging'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; -import { generateBlockId, generateSortingKey } from './dataConversionUtils'; +import { generateBlockId, generateSortingKey, isEphemeralCell } from './dataConversionUtils'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; export function isAgentCell(cell: NotebookCell): boolean { @@ -107,6 +107,7 @@ export async function executeAgentCell( cells: cell.notebook.getCells().filter((c) => c.index !== cell.index) }); + // eslint-disable-next-line local-rules/dont-use-process const openAiToken = process.env.OPENAI_API_KEY; if (openAiToken == null) { throw new Error('OPENAI_API_KEY is not set'); @@ -203,7 +204,7 @@ function getInsertIndexAfterAgentCell( while (index < notebook.cellCount) { const cell = notebook.cellAt(index); - if (cell.metadata?.is_ephemeral === true && cell.metadata?.agent_source_block_id === agentBlockId) { + if (isEphemeralCell(cell) && cell.metadata?.agent_source_block_id === agentBlockId) { index++; } else { break; @@ -294,7 +295,7 @@ async function removeEphemeralCellsForAgent(notebook: NotebookDocument, agentBlo for (let i = notebook.cellCount - 1; i >= 0; i--) { const cell = notebook.cellAt(i); - if (cell.metadata?.is_ephemeral === true && cell.metadata?.agent_source_block_id === agentBlockId) { + if (isEphemeralCell(cell) && cell.metadata?.agent_source_block_id === agentBlockId) { deletions.push(NotebookEdit.deleteCells(new NotebookRange(i, i + 1))); } } diff --git a/src/notebooks/deepnote/dataConversionUtils.ts b/src/notebooks/deepnote/dataConversionUtils.ts index 1b30484770..8b01da1256 100644 --- a/src/notebooks/deepnote/dataConversionUtils.ts +++ b/src/notebooks/deepnote/dataConversionUtils.ts @@ -2,6 +2,8 @@ * Utility functions for Deepnote block ID and sorting key generation */ +import { NotebookCell, NotebookCellData } from 'vscode'; + export function parseJsonWithFallback(value: string, fallback?: unknown): unknown | null { try { return JSON.parse(value); @@ -22,6 +24,13 @@ export function generateBlockId(): string { return id; } +/** + * Returns true if the cell metadata indicates an ephemeral cell (auto-generated by agent). + */ +export function isEphemeralCell(cell: NotebookCell | NotebookCellData): boolean { + return cell.metadata?.is_ephemeral === true; +} + /** * Generate sorting key based on index (format: a0, a1, ..., a99, b0, b1, ...) */ diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 17443e93b9..e9372f94fc 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -6,6 +6,7 @@ import { l10n, window, workspace, type CancellationToken, type NotebookData, typ import { logger } from '../../platform/logging'; import { IDeepnoteNotebookManager } from '../types'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; +import { isEphemeralCell } from './dataConversionUtils'; import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; import { SnapshotService } from './snapshots/snapshotService'; import { computeHash } from '../../platform/common/crypto'; @@ -273,11 +274,17 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { throw new Error(`Notebook with ID ${notebookId} not found in project`); } - logger.debug(`SerializeNotebook: Found notebook, converting ${data.cells.length} cells to blocks`); + // Exclude ephemeral cells (agent-generated) from persistence + const nonEphemeralCells = data.cells.filter((cell) => !isEphemeralCell(cell)); + + logger.debug( + `SerializeNotebook: Found notebook, converting ${nonEphemeralCells.length} cells to blocks ` + + `(${data.cells.length - nonEphemeralCells.length} ephemeral excluded)` + ); // Log cell metadata IDs before conversion - for (let i = 0; i < data.cells.length; i++) { - const cell = data.cells[i]; + for (let i = 0; i < nonEphemeralCells.length; i++) { + const cell = nonEphemeralCells[i]; logger.trace( `SerializeNotebook: cell[${i}] metadata.id=${cell.metadata?.id}, metadata keys=${ cell.metadata ? Object.keys(cell.metadata).join(',') : 'none' @@ -287,7 +294,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { // Clone blocks while removing circular references that may have been // introduced by VS Code's notebook cell/output handling - const blocks = this.converter.convertCellsToBlocks(data.cells); + const blocks = this.converter.convertCellsToBlocks(nonEphemeralCells); logger.debug(`SerializeNotebook: Converted to ${blocks.length} blocks`); @@ -301,7 +308,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } // Add snapshot metadata to blocks (contentHash and execution timing) - await this.addSnapshotMetadataToBlocks(blocks, data); + await this.addSnapshotMetadataToBlocks(blocks, { ...data, cells: nonEphemeralCells }); // Handle snapshot mode: strip outputs and execution metadata from main file if (this.snapshotService?.isSnapshotsEnabled()) { diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index c3332f974d..2813f4acd0 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -205,6 +205,71 @@ project: assert.include(yamlString, 'project-123'); assert.include(yamlString, 'notebook-1'); }); + + test('should exclude ephemeral cells from serialized output', async () => { + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2023-01-01T00:00:00Z', + modifiedAt: '2023-01-02T00:00:00Z' + }, + project: { + id: 'project-ephemeral-exclude', + name: 'Ephemeral Exclude Test', + notebooks: [ + { + id: 'notebook-1', + name: 'Test Notebook', + blocks: [ + { + id: 'block-1', + content: 'print("persisted")', + blockGroup: 'group-1', + metadata: {}, + sortingKey: 'a0', + type: 'code' + } + ], + executionMode: 'block', + isModule: false + } + ], + settings: {} + } + }; + + manager.storeOriginalProject('project-ephemeral-exclude', projectData, 'notebook-1'); + + const mockNotebookData = { + cells: [ + { + kind: 2, + value: 'print("persisted")', + languageId: 'python', + metadata: { id: 'block-1' } + }, + { + kind: 2, + value: 'print("ephemeral - should not persist")', + languageId: 'python', + metadata: { id: 'ephemeral-block', is_ephemeral: true } + } + ], + metadata: { + deepnoteProjectId: 'project-ephemeral-exclude', + deepnoteNotebookId: 'notebook-1' + } + }; + + const result = await serializer.serializeNotebook(mockNotebookData as any, {} as any); + const yamlString = new TextDecoder().decode(result); + const parsedResult = deserializeDeepnoteFile(yamlString); + + const notebook = parsedResult.project.notebooks.find((nb) => nb.id === 'notebook-1'); + assert.isDefined(notebook); + assert.strictEqual(notebook!.blocks.length, 1, 'Ephemeral cell should be excluded'); + assert.strictEqual(notebook!.blocks[0].content, 'print("persisted")'); + }); }); suite('findCurrentNotebookId', () => { diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts index 5c913700ff..36b5d24053 100644 --- a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts @@ -12,6 +12,7 @@ import { } from 'vscode'; import { injectable } from 'inversify'; +import { isEphemeralCell } from './dataConversionUtils'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; const NOTEBOOK_CELL_SCHEME = 'vscode-notebook-cell'; @@ -106,7 +107,7 @@ export class EphemeralCellDecorationProvider implements IExtensionSyncActivation } const cell = this.findCellForEditor(editor); - if (!cell || cell.metadata?.is_ephemeral !== true) { + if (!cell || !isEphemeralCell(cell)) { editor.setDecorations(this.ephemeralDecorationType, []); continue; } diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts index 7089391e7c..60c0a67fba 100644 --- a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts @@ -11,6 +11,7 @@ import { } from 'vscode'; import { injectable } from 'inversify'; +import { isEphemeralCell } from './dataConversionUtils'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; const EPHEMERAL_INDICATOR_PRIORITY = 1000; @@ -54,7 +55,7 @@ export class EphemeralCellStatusBarProvider return undefined; } - if (!this.isEphemeralCell(cell)) { + if (!isEphemeralCell(cell)) { return undefined; } @@ -76,8 +77,4 @@ export class EphemeralCellStatusBarProvider tooltip: tooltipLines.join('\n') }; } - - private isEphemeralCell(cell: NotebookCell): boolean { - return cell.metadata?.is_ephemeral === true; - } } From 6f7b4aa4af611f0baac555b030cd5200a46f68cd Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 16 Mar 2026 21:28:54 +0000 Subject: [PATCH 09/19] Refactor Deepnote server management to utilize a new mock child process helper - Introduced `createMockChildProcess` in `deepnoteTestHelpers.ts` for consistent mock process creation in tests. - Updated `DeepnoteLspClientManager` and `DeepnoteServerStarter` to include the mock process in server info. - Removed unnecessary `runtimeCoreServerInfo` from `ProjectContext` and adjusted related logic to use the new `serverInfo` structure. - Ensured all relevant tests are updated to reflect these changes, improving test reliability and maintainability. --- ...epnoteLspClientManager.node.vscode.test.ts | 10 +++-- .../deepnote/deepnoteServerStarter.node.ts | 44 +++++++------------ .../deepnoteServerStarter.unit.test.ts | 9 +++- src/kernels/deepnote/deepnoteTestHelpers.ts | 15 +++++++ ...teEnvironmentTreeDataProvider.unit.test.ts | 4 +- src/kernels/deepnote/types.ts | 9 +--- ...epnoteKernelAutoSelector.node.unit.test.ts | 4 +- 7 files changed, 54 insertions(+), 41 deletions(-) create mode 100644 src/kernels/deepnote/deepnoteTestHelpers.ts diff --git a/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts b/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts index 5195d0a28b..adad171fbc 100644 --- a/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts +++ b/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts @@ -2,6 +2,7 @@ import { assert } from 'chai'; import { Uri } from 'vscode'; import { DeepnoteLspClientManager } from './deepnoteLspClientManager.node'; +import { createMockChildProcess } from './deepnoteTestHelpers'; import { IDisposableRegistry } from '../../platform/common/types'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import * as path from '../../platform/vscode-path/path'; @@ -84,7 +85,8 @@ suite('DeepnoteLspClientManager Integration Tests', () => { url: 'http://localhost:8888', jupyterPort: 8888, lspPort: 8889, - token: 'test-token' + token: 'test-token', + process: createMockChildProcess() }; // This will attempt to start LSP clients but may fail if pylsp isn't installed @@ -135,7 +137,8 @@ suite('DeepnoteLspClientManager Integration Tests', () => { url: 'http://localhost:8888', jupyterPort: 8888, lspPort: 8889, - token: 'test-token' + token: 'test-token', + process: createMockChildProcess() }; try { @@ -166,7 +169,8 @@ suite('DeepnoteLspClientManager Integration Tests', () => { url: 'http://localhost:8888', jupyterPort: 8888, lspPort: 8889, - token: 'test-token' + token: 'test-token', + process: createMockChildProcess() }; try { diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 22bd170310..87ceaa9094 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -11,7 +11,7 @@ import { inject, injectable, named, optional } from 'inversify'; import * as os from 'os'; import { CancellationToken, l10n, Uri } from 'vscode'; -import { startServer, stopServer, type ServerInfo as RuntimeCoreServerInfo } from '@deepnote/runtime-core'; +import { startServer, stopServer } from '@deepnote/runtime-core'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { Cancellation } from '../../platform/common/cancellation'; @@ -49,7 +49,6 @@ type PendingOperation = interface ProjectContext { environmentId: string; - runtimeCoreServerInfo: RuntimeCoreServerInfo | null; serverInfo: DeepnoteServerInfo | null; } @@ -143,7 +142,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } else { const newContext: ProjectContext = { environmentId, - runtimeCoreServerInfo: null, serverInfo: null }; @@ -266,9 +264,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension // Gather SQL integration env vars to pass to the server const extraEnv = await this.gatherSqlIntegrationEnvVars(deepnoteFileUri, environmentId, token); - let runtimeCoreInfo: RuntimeCoreServerInfo; + let serverInfo: DeepnoteServerInfo; try { - runtimeCoreInfo = await startServer({ + serverInfo = await startServer({ pythonEnv: venvPath.fsPath, workingDirectory: path.dirname(deepnoteFileUri.fsPath), port, @@ -286,28 +284,21 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension ); } - projectContext.runtimeCoreServerInfo = runtimeCoreInfo; - - const serverInfo: DeepnoteServerInfo = { - url: runtimeCoreInfo.url, - jupyterPort: runtimeCoreInfo.jupyterPort, - lspPort: runtimeCoreInfo.lspPort, - process: runtimeCoreInfo.process - }; + projectContext.serverInfo = serverInfo; // Set up output channel logging from the server process - this.monitorServerOutput(serverKey, runtimeCoreInfo); + this.monitorServerOutput(serverKey, serverInfo); // Write lock file for orphan-cleanup tracking - const serverPid = runtimeCoreInfo.process.pid; + const serverPid = serverInfo.process.pid; if (serverPid) { await this.writeLockFile(serverPid); } else { logger.warn(`Could not get PID for server process for ${serverKey}`); } - logger.info(`Deepnote server started successfully at ${runtimeCoreInfo.url} for ${serverKey}`); - this.outputChannel.appendLine(l10n.t('✓ Deepnote server running at {0}', runtimeCoreInfo.url)); + logger.info(`Deepnote server started successfully at ${serverInfo.url} for ${serverKey}`); + this.outputChannel.appendLine(l10n.t('✓ Deepnote server running at {0}', serverInfo.url)); return serverInfo; } @@ -324,20 +315,19 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); - const runtimeCoreInfo = projectContext?.runtimeCoreServerInfo; + const serverInfo = projectContext?.serverInfo; - if (runtimeCoreInfo) { - const serverPid = runtimeCoreInfo.process.pid; + if (serverInfo) { + const serverPid = serverInfo.process.pid; try { logger.info(`Stopping Deepnote server for ${fileKey}...`); - await stopServer(runtimeCoreInfo); + await stopServer(serverInfo); this.outputChannel.appendLine(l10n.t('Deepnote server stopped for {0}', fileKey)); } catch (ex) { logger.error('Error stopping Deepnote server', ex); } finally { if (projectContext) { - projectContext.runtimeCoreServerInfo = null; projectContext.serverInfo = null; } @@ -439,8 +429,8 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension /** * Stream stdout/stderr from the server process to the VSCode output channel. */ - private monitorServerOutput(serverKey: string, runtimeCoreInfo: RuntimeCoreServerInfo): void { - const proc = runtimeCoreInfo.process; + private monitorServerOutput(serverKey: string, serverInfo: DeepnoteServerInfo): void { + const proc = serverInfo.process; const disposables: IDisposable[] = []; this.disposablesByFile.set(serverKey, disposables); @@ -488,15 +478,15 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const pidsToCleanup: number[] = []; for (const [key, ctx] of this.projectContexts.entries()) { - if (ctx.runtimeCoreServerInfo) { - const pid = ctx.runtimeCoreServerInfo.process.pid; + if (ctx.serverInfo) { + const pid = ctx.serverInfo.process.pid; if (pid) { pidsToCleanup.push(pid); } logger.info(`Stopping Deepnote server for ${key}...`); stopPromises.push( - stopServer(ctx.runtimeCoreServerInfo).catch((ex) => { + stopServer(ctx.serverInfo).catch((ex) => { logger.error(`Error stopping Deepnote server for ${key}`, ex); }) ); diff --git a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts index c63174c83b..b92d7bd8a2 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts @@ -3,6 +3,7 @@ import { anything, instance, mock, when } from 'ts-mockito'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; import { DeepnoteServerStarter } from './deepnoteServerStarter.node'; +import { createMockChildProcess } from './deepnoteTestHelpers'; import { IProcessServiceFactory } from '../../platform/common/process/types.node'; import { IAsyncDisposableRegistry, IOutputChannel } from '../../platform/common/types'; import { IDeepnoteToolkitInstaller } from './types'; @@ -73,8 +74,12 @@ suite('DeepnoteServerStarter', () => { const projectContexts = (serverStarter as any).projectContexts as Map; projectContexts.set('existing-key', { environmentId: 'env1', - runtimeCoreServerInfo: null, - serverInfo: { url: 'http://localhost:8888', jupyterPort: 8888, lspPort: 8889 } + serverInfo: { + url: 'http://localhost:8888', + jupyterPort: 8888, + lspPort: 8889, + process: createMockChildProcess() + } }); const port = await reserveStartPort('test-key-2'); diff --git a/src/kernels/deepnote/deepnoteTestHelpers.ts b/src/kernels/deepnote/deepnoteTestHelpers.ts new file mode 100644 index 0000000000..524c6e0e47 --- /dev/null +++ b/src/kernels/deepnote/deepnoteTestHelpers.ts @@ -0,0 +1,15 @@ +import type { ChildProcess } from 'node:child_process'; + +/** + * Creates a mock ChildProcess for use in Deepnote server info tests. + * Satisfies the ChildProcess interface with minimal stub values. + */ +export function createMockChildProcess(overrides?: Partial): ChildProcess { + return { + pid: undefined, + stdout: null, + stderr: null, + exitCode: null, + ...overrides + } as ChildProcess; +} diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.unit.test.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.unit.test.ts index 4cc6d5df5d..e91c9d3a0b 100644 --- a/src/kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.unit.test.ts +++ b/src/kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.unit.test.ts @@ -1,6 +1,7 @@ import { assert } from 'chai'; import { instance, mock, when } from 'ts-mockito'; import { Uri, EventEmitter } from 'vscode'; +import { createMockChildProcess } from '../deepnoteTestHelpers'; import { DeepnoteEnvironmentTreeDataProvider } from './deepnoteEnvironmentTreeDataProvider.node'; import { IDeepnoteEnvironmentManager } from '../types'; import { DeepnoteEnvironment } from './deepnoteEnvironment'; @@ -40,7 +41,8 @@ suite('DeepnoteEnvironmentTreeDataProvider', () => { url: 'http://localhost:8888', jupyterPort: 8888, lspPort: 8889, - token: 'test-token' + token: 'test-token', + process: createMockChildProcess() } }; diff --git a/src/kernels/deepnote/types.ts b/src/kernels/deepnote/types.ts index 5c63eacd1d..a0c17a31ba 100644 --- a/src/kernels/deepnote/types.ts +++ b/src/kernels/deepnote/types.ts @@ -1,7 +1,7 @@ // Copyright (c) Microsoft Corporation. // Licensed under the MIT License. -import type { ChildProcess } from 'node:child_process'; +import type { ServerInfo as RuntimeCoreServerInfo } from '@deepnote/runtime-core'; import * as vscode from 'vscode'; import { serializePythonEnvironment } from '../../platform/api/pythonApi'; @@ -186,13 +186,8 @@ export interface IDeepnoteServerStarter { dispose(): Promise; } -export interface DeepnoteServerInfo { - url: string; - jupyterPort: number; - lspPort: number; +export interface DeepnoteServerInfo extends RuntimeCoreServerInfo { token?: string; - /** The underlying server process from @deepnote/runtime-core, used for lifecycle management */ - process?: ChildProcess; } export const IDeepnoteServerProvider = Symbol('IDeepnoteServerProvider'); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts index 8141e32093..1872281092 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -2,6 +2,7 @@ import { assert } from 'chai'; import * as sinon from 'sinon'; import { anything, instance, mock, verify, when } from 'ts-mockito'; import { DeepnoteKernelAutoSelector } from './deepnoteKernelAutoSelector.node'; +import { createMockChildProcess } from '../../kernels/deepnote/deepnoteTestHelpers'; import { IDeepnoteEnvironmentManager, IDeepnoteLspClientManager, @@ -1042,7 +1043,8 @@ function createMockEnvironment(id: string, name: string, hasServer: boolean = fa url: `http://localhost:8888`, jupyterPort: 8888, lspPort: 8889, - token: 'test-token' + token: 'test-token', + process: createMockChildProcess() } : undefined }; From 7d8ee860c3bcb409f98ccf63473ae6f1e981cbdb Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 16 Mar 2026 21:55:09 +0000 Subject: [PATCH 10/19] Fix eslint error --- .../deepnote/deepnoteLspClientManager.node.vscode.test.ts | 2 +- src/kernels/deepnote/deepnoteServerStarter.unit.test.ts | 2 +- .../{deepnoteTestHelpers.ts => deepnoteTestHelpers.node.ts} | 0 .../deepnoteEnvironmentTreeDataProvider.unit.test.ts | 2 +- .../deepnote/deepnoteKernelAutoSelector.node.unit.test.ts | 2 +- 5 files changed, 4 insertions(+), 4 deletions(-) rename src/kernels/deepnote/{deepnoteTestHelpers.ts => deepnoteTestHelpers.node.ts} (100%) diff --git a/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts b/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts index adad171fbc..31a6c85845 100644 --- a/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts +++ b/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts @@ -2,7 +2,7 @@ import { assert } from 'chai'; import { Uri } from 'vscode'; import { DeepnoteLspClientManager } from './deepnoteLspClientManager.node'; -import { createMockChildProcess } from './deepnoteTestHelpers'; +import { createMockChildProcess } from './deepnoteTestHelpers.node'; import { IDisposableRegistry } from '../../platform/common/types'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import * as path from '../../platform/vscode-path/path'; diff --git a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts index b92d7bd8a2..abec74328f 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts @@ -3,7 +3,7 @@ import { anything, instance, mock, when } from 'ts-mockito'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; import { DeepnoteServerStarter } from './deepnoteServerStarter.node'; -import { createMockChildProcess } from './deepnoteTestHelpers'; +import { createMockChildProcess } from './deepnoteTestHelpers.node'; import { IProcessServiceFactory } from '../../platform/common/process/types.node'; import { IAsyncDisposableRegistry, IOutputChannel } from '../../platform/common/types'; import { IDeepnoteToolkitInstaller } from './types'; diff --git a/src/kernels/deepnote/deepnoteTestHelpers.ts b/src/kernels/deepnote/deepnoteTestHelpers.node.ts similarity index 100% rename from src/kernels/deepnote/deepnoteTestHelpers.ts rename to src/kernels/deepnote/deepnoteTestHelpers.node.ts diff --git a/src/kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.unit.test.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.unit.test.ts index e91c9d3a0b..05b690f3b7 100644 --- a/src/kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.unit.test.ts +++ b/src/kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.unit.test.ts @@ -1,7 +1,7 @@ import { assert } from 'chai'; import { instance, mock, when } from 'ts-mockito'; import { Uri, EventEmitter } from 'vscode'; -import { createMockChildProcess } from '../deepnoteTestHelpers'; +import { createMockChildProcess } from '../deepnoteTestHelpers.node'; import { DeepnoteEnvironmentTreeDataProvider } from './deepnoteEnvironmentTreeDataProvider.node'; import { IDeepnoteEnvironmentManager } from '../types'; import { DeepnoteEnvironment } from './deepnoteEnvironment'; diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts index 1872281092..824635343c 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -2,7 +2,7 @@ import { assert } from 'chai'; import * as sinon from 'sinon'; import { anything, instance, mock, verify, when } from 'ts-mockito'; import { DeepnoteKernelAutoSelector } from './deepnoteKernelAutoSelector.node'; -import { createMockChildProcess } from '../../kernels/deepnote/deepnoteTestHelpers'; +import { createMockChildProcess } from '../../kernels/deepnote/deepnoteTestHelpers.node'; import { IDeepnoteEnvironmentManager, IDeepnoteLspClientManager, From 9c4151ac873ca2e85f2aabb3ed7e1961fb7ee7aa Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 17 Mar 2026 08:54:53 +0000 Subject: [PATCH 11/19] Enhance Deepnote server stopping logic to handle missing project context - Added a warning log when no project context is found, preventing server stop attempts. - Updated the `stopServerForEnvironment` method to require a non-null project context, ensuring safer operation handling. --- src/kernels/deepnote/deepnoteServerStarter.node.ts | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 87ceaa9094..daef7e2de5 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -185,6 +185,11 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const fileKey = deepnoteFileUri.fsPath; const projectContext = this.projectContexts.get(fileKey) ?? null; + if (projectContext == null) { + logger.warn(`No project context found for ${fileKey}, skipping stop server...`); + return; + } + const pendingOp = this.pendingOperations.get(fileKey); if (pendingOp) { logger.info(`Waiting for pending operation on ${fileKey} before stopping...`); @@ -307,7 +312,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension * Stop the server using @deepnote/runtime-core's `stopServer` (SIGTERM -> wait -> SIGKILL). */ private async stopServerForEnvironment( - projectContext: ProjectContext | null, + projectContext: ProjectContext, deepnoteFileUri: Uri, token?: CancellationToken ): Promise { @@ -315,7 +320,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); - const serverInfo = projectContext?.serverInfo; + const { serverInfo } = projectContext; if (serverInfo) { const serverPid = serverInfo.process.pid; From 347a8097391298bdd85b78d19d5b34e6962e4611 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 17 Mar 2026 12:03:09 +0000 Subject: [PATCH 12/19] Refactor Deepnote server management to use fileKey instead of serverKey - Updated the DeepnoteServerStarter class to consistently use fileKey for managing pending operations and project contexts, improving clarity and reducing potential errors. - Adjusted logging messages to reflect the change, ensuring accurate information is logged during server operations. --- .../deepnote/deepnoteServerStarter.node.ts | 53 +++++++++---------- 1 file changed, 25 insertions(+), 28 deletions(-) diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index daef7e2de5..557876d86e 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -106,11 +106,10 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension token?: CancellationToken ): Promise { const fileKey = deepnoteFileUri.fsPath; - const serverKey = `${fileKey}-${environmentId}`; - let pendingOp = this.pendingOperations.get(serverKey); + let pendingOp = this.pendingOperations.get(fileKey); if (pendingOp) { - logger.info(`Waiting for pending operation on ${serverKey} to complete...`); + logger.info(`Waiting for pending operation on ${fileKey} to complete...`); try { await pendingOp.promise; } catch { @@ -118,26 +117,29 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } - let existingContext = this.projectContexts.get(serverKey); + let existingContext = this.projectContexts.get(fileKey); if (existingContext != null) { const { environmentId: existingEnvironmentId, serverInfo: existingServerInfo } = existingContext; if (existingEnvironmentId === environmentId) { if (existingServerInfo != null && (await this.isServerRunning(existingServerInfo))) { - logger.info(`Deepnote server already running at ${existingServerInfo.url} for ${serverKey}`); + logger.info( + `Deepnote server already running at ${existingServerInfo.url} for ${fileKey} (environmentId ${environmentId})` + ); return existingServerInfo; } - pendingOp = this.pendingOperations.get(serverKey); + pendingOp = this.pendingOperations.get(fileKey); if (pendingOp && pendingOp.type === 'start') { return await pendingOp.promise; } } else { logger.info( - `Stopping existing server for ${serverKey} with environmentId ${existingEnvironmentId} to start new one with environmentId ${environmentId}...` + `Stopping existing server for ${fileKey} with environmentId ${existingEnvironmentId} to start new one with environmentId ${environmentId}...` ); await this.stopServerForEnvironment(existingContext, deepnoteFileUri, token); + existingContext.environmentId = environmentId; } } else { const newContext: ProjectContext = { @@ -145,7 +147,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension serverInfo: null }; - this.projectContexts.set(serverKey, newContext); + this.projectContexts.set(fileKey, newContext); existingContext = newContext; } @@ -162,7 +164,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension token ) }; - this.pendingOperations.set(serverKey, operation); + this.pendingOperations.set(fileKey, operation); try { const result = await operation.promise; @@ -170,8 +172,8 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension existingContext.serverInfo = result; return result; } finally { - if (this.pendingOperations.get(serverKey) === operation) { - this.pendingOperations.delete(serverKey); + if (this.pendingOperations.get(fileKey) === operation) { + this.pendingOperations.delete(fileKey); } } } @@ -238,7 +240,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension token?: CancellationToken ): Promise { const fileKey = deepnoteFileUri.fsPath; - const serverKey = `${fileKey}-${environmentId}`; Cancellation.throwIfCanceled(token); @@ -259,11 +260,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); // Serialize port allocation across concurrent server starts - const port = await this.reserveStartPort(serverKey); + const port = await this.reserveStartPort(fileKey); - logger.info( - `Starting deepnote-toolkit server on port ${port} for ${serverKey} with environmentId ${environmentId}` - ); + logger.info(`Starting deepnote-toolkit server on port ${port} for ${fileKey} (environmentId ${environmentId})`); this.outputChannel.appendLine(l10n.t('Starting Deepnote server on port {0}...', port)); // Gather SQL integration env vars to pass to the server @@ -292,17 +291,17 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension projectContext.serverInfo = serverInfo; // Set up output channel logging from the server process - this.monitorServerOutput(serverKey, serverInfo); + this.monitorServerOutput(fileKey, serverInfo); // Write lock file for orphan-cleanup tracking const serverPid = serverInfo.process.pid; if (serverPid) { await this.writeLockFile(serverPid); } else { - logger.warn(`Could not get PID for server process for ${serverKey}`); + logger.warn(`Could not get PID for server process for ${fileKey}`); } - logger.info(`Deepnote server started successfully at ${serverInfo.url} for ${serverKey}`); + logger.info(`Deepnote server started successfully at ${serverInfo.url} for ${fileKey}`); this.outputChannel.appendLine(l10n.t('✓ Deepnote server running at {0}', serverInfo.url)); return serverInfo; @@ -332,9 +331,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } catch (ex) { logger.error('Error stopping Deepnote server', ex); } finally { - if (projectContext) { - projectContext.serverInfo = null; - } + projectContext.serverInfo = null; if (serverPid) { await this.deleteLockFile(serverPid); @@ -370,7 +367,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension * servers start concurrently in the extension, they can race. This lock serializes * the starts so each `startServer` call sees the ports bound by previous calls. */ - private async reserveStartPort(serverKey: string): Promise { + private async reserveStartPort(fileKey: string): Promise { const previousLock = this.portAllocationLock; let releaseLock: () => void; const currentLock = new Promise((resolve) => { @@ -389,7 +386,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } - logger.info(`Reserved start port ${maxPort} for ${serverKey}`); + logger.info(`Reserved start port ${maxPort} for ${fileKey}`); return maxPort; } finally { releaseLock!(); @@ -434,16 +431,16 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension /** * Stream stdout/stderr from the server process to the VSCode output channel. */ - private monitorServerOutput(serverKey: string, serverInfo: DeepnoteServerInfo): void { + private monitorServerOutput(fileKey: string, serverInfo: DeepnoteServerInfo): void { const proc = serverInfo.process; const disposables: IDisposable[] = []; - this.disposablesByFile.set(serverKey, disposables); + this.disposablesByFile.set(fileKey, disposables); if (proc.stdout) { const stdout = proc.stdout; const onData = (data: Buffer) => { const text = data.toString(); - logger.trace(`Deepnote server (${serverKey}): ${text}`); + logger.trace(`Deepnote server (${fileKey}): ${text}`); this.outputChannel.appendLine(text); }; stdout.on('data', onData); @@ -458,7 +455,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const stderr = proc.stderr; const onData = (data: Buffer) => { const text = data.toString(); - logger.warn(`Deepnote server stderr (${serverKey}): ${text}`); + logger.warn(`Deepnote server stderr (${fileKey}): ${text}`); this.outputChannel.appendLine(text); }; stderr.on('data', onData); From 5eb2c1719bf6ce70a561a59a2d7af0540af583c8 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 17 Mar 2026 12:39:46 +0000 Subject: [PATCH 13/19] Refactor DeepnoteServerStarter to remove port allocation logic - Eliminated the port allocation serialization logic from the DeepnoteServerStarter class, as it is now handled by the @deepnote/runtime-core's startServer method. - Updated related logging messages to reflect the changes in server startup processes. - Adjusted unit tests to focus on SQL environment variable gathering and lifecycle orchestration, removing tests related to port reservation. --- .../deepnote/deepnoteServerStarter.node.ts | 49 ++--------------- .../deepnoteServerStarter.unit.test.ts | 52 +------------------ 2 files changed, 6 insertions(+), 95 deletions(-) diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 557876d86e..8da9fb2325 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -1,6 +1,5 @@ /** * @deepnote/runtime-core functions not currently exported that would be useful: - * - findConsecutiveAvailablePorts(startPort) — duplicated logic for multi-server port reservation * - waitForServer(info, timeoutMs) — health-check polling on /api * - createJsonWebSocketFactory() — forces JSON-only Jupyter WS protocol, potential stability improvement * - ExecutionEngine.toPythonLiteral(value) — JS-to-Python literal conversion @@ -65,7 +64,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension private readonly disposablesByFile: Map = new Map(); private readonly projectContexts: Map = new Map(); private readonly pendingOperations: Map = new Map(); - private portAllocationLock: Promise = Promise.resolve(); private readonly sessionId: string = generateUuid(); private readonly lockFileDir: string = path.join(os.tmpdir(), 'vscode-deepnote-locks'); @@ -227,7 +225,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension * - SQL integration env var injection (via ServerOptions.env) * - Lock file creation (after start, using returned PID) * - Output channel logging (via process stdout/stderr streams) - * - Port allocation serialization across concurrent starts */ private async startServerForEnvironment( projectContext: ProjectContext, @@ -259,28 +256,23 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); - // Serialize port allocation across concurrent server starts - const port = await this.reserveStartPort(fileKey); + logger.info(`Starting deepnote-toolkit server for ${fileKey} (environmentId ${environmentId})`); + this.outputChannel.appendLine(l10n.t('Starting Deepnote server...')); - logger.info(`Starting deepnote-toolkit server on port ${port} for ${fileKey} (environmentId ${environmentId})`); - this.outputChannel.appendLine(l10n.t('Starting Deepnote server on port {0}...', port)); - - // Gather SQL integration env vars to pass to the server const extraEnv = await this.gatherSqlIntegrationEnvVars(deepnoteFileUri, environmentId, token); - let serverInfo: DeepnoteServerInfo; + let serverInfo: DeepnoteServerInfo | undefined; try { serverInfo = await startServer({ pythonEnv: venvPath.fsPath, workingDirectory: path.dirname(deepnoteFileUri.fsPath), - port, startupTimeoutMs: SERVER_STARTUP_TIMEOUT_MS, env: extraEnv }); } catch (error) { throw new DeepnoteServerStartupError( interpreter.uri.fsPath, - port, + serverInfo?.jupyterPort ?? 0, 'unknown', '', error instanceof Error ? error.message : String(error), @@ -360,39 +352,6 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } - /** - * Serialize port reservation across concurrent server starts. - * - * runtime-core's `startServer` finds its own consecutive ports, but when multiple - * servers start concurrently in the extension, they can race. This lock serializes - * the starts so each `startServer` call sees the ports bound by previous calls. - */ - private async reserveStartPort(fileKey: string): Promise { - const previousLock = this.portAllocationLock; - let releaseLock: () => void; - const currentLock = new Promise((resolve) => { - releaseLock = resolve; - }); - this.portAllocationLock = previousLock.then(() => currentLock); - - await previousLock; - - try { - // Collect ports already in use by running servers to pick a non-conflicting start port - let maxPort = 8888; - for (const ctx of this.projectContexts.values()) { - if (ctx.serverInfo) { - maxPort = Math.max(maxPort, ctx.serverInfo.jupyterPort + 2, ctx.serverInfo.lspPort + 1); - } - } - - logger.info(`Reserved start port ${maxPort} for ${fileKey}`); - return maxPort; - } finally { - releaseLock!(); - } - } - /** * Gather SQL integration environment variables for the deepnote-toolkit server. */ diff --git a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts index abec74328f..5792e8c39d 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts @@ -3,7 +3,6 @@ import { anything, instance, mock, when } from 'ts-mockito'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; import { DeepnoteServerStarter } from './deepnoteServerStarter.node'; -import { createMockChildProcess } from './deepnoteTestHelpers.node'; import { IProcessServiceFactory } from '../../platform/common/process/types.node'; import { IAsyncDisposableRegistry, IOutputChannel } from '../../platform/common/types'; import { IDeepnoteToolkitInstaller } from './types'; @@ -12,10 +11,9 @@ import { ISqlIntegrationEnvVarsProvider } from '../../platform/notebooks/deepnot /** * Unit tests for DeepnoteServerStarter. * - * Port allocation, server spawning, and health checks are now delegated to + * Port allocation, server spawning, and health checks are delegated to * @deepnote/runtime-core's startServer/stopServer. These tests focus on the - * extension-specific layers: port reservation serialization, SQL env var - * gathering, and lifecycle orchestration. + * extension-specific layers: SQL env var gathering and lifecycle orchestration. */ suite('DeepnoteServerStarter', () => { let serverStarter: DeepnoteServerStarter; @@ -58,52 +56,6 @@ suite('DeepnoteServerStarter', () => { } }); - suite('reserveStartPort - Port Serialization', () => { - test('should return default port when no servers are running', async () => { - const reserveStartPort = getPrivateMethod(serverStarter, 'reserveStartPort'); - const port = await reserveStartPort('test-key'); - - assert.strictEqual(port, 8888); - }); - - test('should return ports beyond existing servers', async () => { - const reserveStartPort = getPrivateMethod(serverStarter, 'reserveStartPort'); - - // Simulate a running server context by directly setting projectContexts - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const projectContexts = (serverStarter as any).projectContexts as Map; - projectContexts.set('existing-key', { - environmentId: 'env1', - serverInfo: { - url: 'http://localhost:8888', - jupyterPort: 8888, - lspPort: 8889, - process: createMockChildProcess() - } - }); - - const port = await reserveStartPort('test-key-2'); - - assert.isAtLeast(port, 8890, 'Should skip ports used by existing servers'); - }); - - test('should serialize concurrent calls', async () => { - const reserveStartPort = getPrivateMethod(serverStarter, 'reserveStartPort'); - - // Launch concurrent port reservations - const [port1, port2, port3] = await Promise.all([ - reserveStartPort('key-1'), - reserveStartPort('key-2'), - reserveStartPort('key-3') - ]); - - // All should return valid numbers (even if same, since no server info is stored between calls) - assert.isNumber(port1); - assert.isNumber(port2); - assert.isNumber(port3); - }); - }); - suite('gatherSqlIntegrationEnvVars', () => { test('should return empty object when no provider is available', async () => { // Create a starter without SQL provider From 3740df1162a66da94d052959ac75b18a8acd08ab Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 17 Mar 2026 15:02:33 +0000 Subject: [PATCH 14/19] Enhance DeepnoteServerStarter with output tracking and error reporting improvements - Introduced a new `serverOutputByFile` map to track stdout and stderr outputs for each server instance, limiting the output length to improve performance and manageability. - Updated error handling in the server startup process to capture and report both stdout and stderr in case of failures, providing better diagnostics. - Adjusted the `dispose` method to ensure all internal states, including the new output tracking, are cleared appropriately. - Enhanced unit tests to validate the new output tracking functionality and ensure proper handling of cancellation errors. --- .../deepnote/deepnoteServerStarter.node.ts | 37 +++++++-- .../deepnoteServerStarter.unit.test.ts | 77 ++++++++++++++++++- 2 files changed, 106 insertions(+), 8 deletions(-) diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 8da9fb2325..08c19a33eb 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -27,6 +27,7 @@ import * as path from '../../platform/vscode-path/path'; import { DeepnoteServerInfo, IDeepnoteServerStarter, IDeepnoteToolkitInstaller } from './types'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; +const MAX_OUTPUT_TRACKING_LENGTH = 5000; const SERVER_STARTUP_TIMEOUT_MS = 120_000; const GRACEFUL_SHUTDOWN_TIMEOUT_MS = 3000; @@ -62,8 +63,9 @@ interface ProjectContext { @injectable() export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtensionSyncActivationService { private readonly disposablesByFile: Map = new Map(); - private readonly projectContexts: Map = new Map(); private readonly pendingOperations: Map = new Map(); + private readonly projectContexts: Map = new Map(); + private readonly serverOutputByFile: Map = new Map(); private readonly sessionId: string = generateUuid(); private readonly lockFileDir: string = path.join(os.tmpdir(), 'vscode-deepnote-locks'); @@ -261,6 +263,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const extraEnv = await this.gatherSqlIntegrationEnvVars(deepnoteFileUri, environmentId, token); + // Initialize output tracking for error reporting + this.serverOutputByFile.set(fileKey, { stdout: '', stderr: '' }); + let serverInfo: DeepnoteServerInfo | undefined; try { serverInfo = await startServer({ @@ -270,13 +275,16 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension env: extraEnv }); } catch (error) { + const capturedOutput = this.serverOutputByFile.get(fileKey); + this.serverOutputByFile.delete(fileKey); + throw new DeepnoteServerStartupError( interpreter.uri.fsPath, serverInfo?.jupyterPort ?? 0, 'unknown', - '', - error instanceof Error ? error.message : String(error), - error instanceof Error ? error : undefined + capturedOutput?.stdout || '', + capturedOutput?.stderr || (error instanceof Error ? error.message : String(error)), + error instanceof Error ? error : new Error(`${error}`) ); } @@ -333,6 +341,8 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); + this.serverOutputByFile.delete(fileKey); + const disposables = this.disposablesByFile.get(fileKey); if (disposables) { disposables.forEach((d) => d.dispose()); @@ -345,7 +355,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension */ private async isServerRunning(serverInfo: DeepnoteServerInfo): Promise { try { - const response = await fetch(`${serverInfo.url}/api`); + const response = await fetch(`${serverInfo.url}/api`, { signal: AbortSignal.timeout(5000) }); return response.ok; } catch { return false; @@ -401,6 +411,11 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const text = data.toString(); logger.trace(`Deepnote server (${fileKey}): ${text}`); this.outputChannel.appendLine(text); + + const outputTracking = this.serverOutputByFile.get(fileKey); + if (outputTracking) { + outputTracking.stdout = (outputTracking.stdout + text).slice(-MAX_OUTPUT_TRACKING_LENGTH); + } }; stdout.on('data', onData); disposables.push({ @@ -416,6 +431,11 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const text = data.toString(); logger.warn(`Deepnote server stderr (${fileKey}): ${text}`); this.outputChannel.appendLine(text); + + const outputTracking = this.serverOutputByFile.get(fileKey); + if (outputTracking) { + outputTracking.stderr = (outputTracking.stderr + text).slice(-MAX_OUTPUT_TRACKING_LENGTH); + } }; stderr.on('data', onData); disposables.push({ @@ -432,7 +452,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension const pendingOps = Array.from(this.pendingOperations.values()); if (pendingOps.length > 0) { logger.info(`Waiting for ${pendingOps.length} pending operations to complete...`); - await Promise.allSettled(pendingOps.map((op) => Promise.race([op, sleep(GRACEFUL_SHUTDOWN_TIMEOUT_MS)]))); + await Promise.allSettled( + pendingOps.map((op) => Promise.race([op.promise, sleep(GRACEFUL_SHUTDOWN_TIMEOUT_MS)])) + ); } const stopPromises: Promise[] = []; @@ -471,9 +493,10 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } - this.projectContexts.clear(); this.disposablesByFile.clear(); this.pendingOperations.clear(); + this.projectContexts.clear(); + this.serverOutputByFile.clear(); logger.info('DeepnoteServerStarter disposed successfully'); } diff --git a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts index 5792e8c39d..11207ec342 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts @@ -1,4 +1,5 @@ import { assert } from 'chai'; +import * as fakeTimers from '@sinonjs/fake-timers'; import { anything, instance, mock, when } from 'ts-mockito'; import { DeepnoteAgentSkillsManager } from './deepnoteAgentSkillsManager.node'; @@ -75,17 +76,91 @@ suite('DeepnoteServerStarter', () => { await starterWithoutSql.dispose(); }); + + test('should return empty object when provider rejects with cancellation error', async () => { + const { CancellationError, Uri } = await import('vscode'); + + const cancelledProvider = mock(); + when(cancelledProvider.getEnvironmentVariables(anything(), anything())).thenReject( + new CancellationError() + ); + + const starterWithCancelledSql = new DeepnoteServerStarter( + instance(mockProcessServiceFactory), + instance(mockToolkitInstaller), + instance(mockAgentSkillsManager), + instance(mockOutputChannel), + instance(mockAsyncRegistry), + instance(cancelledProvider) + ); + + const gatherEnvVars = getPrivateMethod(starterWithCancelledSql, 'gatherSqlIntegrationEnvVars'); + const result = await gatherEnvVars(Uri.file('/test/file.deepnote'), 'env1'); + + assert.deepStrictEqual(result, {}); + + await starterWithCancelledSql.dispose(); + }); }); suite('dispose', () => { + let clock: fakeTimers.InstalledClock; + + setup(() => { + clock = fakeTimers.install(); + }); + + teardown(() => { + clock.uninstall(); + }); + test('should clear all internal state', async () => { await serverStarter.dispose(); // eslint-disable-next-line @typescript-eslint/no-explicit-any const starter = serverStarter as any; - assert.strictEqual(starter.projectContexts.size, 0); assert.strictEqual(starter.disposablesByFile.size, 0); assert.strictEqual(starter.pendingOperations.size, 0); + assert.strictEqual(starter.projectContexts.size, 0); + assert.strictEqual(starter.serverOutputByFile.size, 0); + }); + + test('should wait for in-flight pending operations before completing', async () => { + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const starter = serverStarter as any; + + let resolveDeferred!: () => void; + const deferred = new Promise((resolve) => { + resolveDeferred = resolve; + }); + + starter.pendingOperations.set('/test/inflight.deepnote', { + type: 'stop', + promise: deferred + }); + + let disposeResolved = false; + const disposePromise = serverStarter.dispose().then(() => { + disposeResolved = true; + }); + + await clock.tickAsync(0); + assert.strictEqual( + disposeResolved, + false, + 'dispose() should not resolve while a pending operation is in flight' + ); + + resolveDeferred(); + await clock.tickAsync(0); + await disposePromise; + + assert.strictEqual( + disposeResolved, + true, + 'dispose() should resolve after pending operation completes' + ); + assert.strictEqual(starter.pendingOperations.size, 0); }); }); }); From 06a03672f4ec7f17c25236f77e41e442b03746b9 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 17 Mar 2026 15:04:33 +0000 Subject: [PATCH 15/19] Update error handling in DeepnoteServerStarter to improve diagnostics - Modified the error reporting logic to ensure that stderr output is captured only when available, enhancing clarity in error messages. - This change aims to streamline the error handling process during server startup, providing more accurate feedback in case of failures. --- src/kernels/deepnote/deepnoteServerStarter.node.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/kernels/deepnote/deepnoteServerStarter.node.ts b/src/kernels/deepnote/deepnoteServerStarter.node.ts index 08c19a33eb..112f7f0e09 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.node.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.node.ts @@ -283,7 +283,7 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension serverInfo?.jupyterPort ?? 0, 'unknown', capturedOutput?.stdout || '', - capturedOutput?.stderr || (error instanceof Error ? error.message : String(error)), + capturedOutput?.stderr || '', error instanceof Error ? error : new Error(`${error}`) ); } From 1e7cb073bdc65e6fa009bcffe3de3138088f9cbe Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 17 Mar 2026 15:37:08 +0000 Subject: [PATCH 16/19] Reformat code --- .../deepnote/deepnoteServerStarter.unit.test.ts | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts index 11207ec342..f90e3a8ec1 100644 --- a/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts +++ b/src/kernels/deepnote/deepnoteServerStarter.unit.test.ts @@ -81,9 +81,7 @@ suite('DeepnoteServerStarter', () => { const { CancellationError, Uri } = await import('vscode'); const cancelledProvider = mock(); - when(cancelledProvider.getEnvironmentVariables(anything(), anything())).thenReject( - new CancellationError() - ); + when(cancelledProvider.getEnvironmentVariables(anything(), anything())).thenReject(new CancellationError()); const starterWithCancelledSql = new DeepnoteServerStarter( instance(mockProcessServiceFactory), @@ -155,11 +153,7 @@ suite('DeepnoteServerStarter', () => { await clock.tickAsync(0); await disposePromise; - assert.strictEqual( - disposeResolved, - true, - 'dispose() should resolve after pending operation completes' - ); + assert.strictEqual(disposeResolved, true, 'dispose() should resolve after pending operation completes'); assert.strictEqual(starter.pendingOperations.size, 0); }); }); From 75d02207abc9a1c2dfabe43798918f414bce9ef1 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 18 Mar 2026 10:16:58 +0000 Subject: [PATCH 17/19] Enhance agent cell execution handling and status bar provider - Introduced `getOpenAiApiKey` function to retrieve the OpenAI API key from configuration, improving error handling when the key is not set. - Updated `executeAgentCell` and `executeEphemeralCell` functions to utilize the new API key retrieval method and handle cancellation tokens. - Enhanced `AgentCellStatusBarProvider` to validate max iterations using Zod schema, ensuring robust input handling and defaulting to safe values. - Added unit tests for new functionality and edge cases in both execution handling and status bar provider. --- .../deepnote/agentCellExecutionHandler.ts | 49 ++++++++--- .../agentCellExecutionHandler.unit.test.ts | 83 +++++++++++++++---- .../deepnote/agentCellStatusBarProvider.ts | 19 ++++- .../agentCellStatusBarProvider.unit.test.ts | 60 ++++++++++++++ 4 files changed, 177 insertions(+), 34 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index ea8485e8a4..e8a1268194 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -1,4 +1,6 @@ import { + CancellationError, + CancellationToken, NotebookCell, NotebookCellOutput, NotebookCellOutputItem, @@ -20,7 +22,9 @@ import { } from '@deepnote/runtime-core'; import { translateCellDisplayOutput } from '../../kernels/execution/helpers'; +import type { IDisposable } from '../../platform/common/types'; import { createDeferred } from '../../platform/common/utils/async'; +import { dispose } from '../../platform/common/utils/lifecycle'; import { uuidUtils } from '../../platform/common/uuid'; import type { Pocket } from '../../platform/deepnote/pocket'; import { logger } from '../../platform/logging'; @@ -59,6 +63,17 @@ export function serializeNotebookContext({ cells }: { cells: NotebookCell[] }): return serializeNotebookContextFromBlocks({ blocks, notebookName: null }); } +export function getOpenAiApiKey(): string { + const config = workspace.getConfiguration('deepnote'); + const key = config.get('agent.openAiApiKey', ''); + + if (!key) { + throw new Error('deepnote.agent.openAiApiKey is not set. Configure it in VS Code settings.'); + } + + return key; +} + export interface ExecuteAgentCellOptions { executeAgentBlockFn?: typeof executeAgentBlock; } @@ -107,11 +122,7 @@ export async function executeAgentCell( cells: cell.notebook.getCells().filter((c) => c.index !== cell.index) }); - // eslint-disable-next-line local-rules/dont-use-process - const openAiToken = process.env.OPENAI_API_KEY; - if (openAiToken == null) { - throw new Error('OPENAI_API_KEY is not set'); - } + const openAiToken = getOpenAiApiKey(); const context: AgentBlockContext = { openAiToken, @@ -125,7 +136,7 @@ export async function executeAgentCell( const cellIndex = await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'code', code); const insertedCell = cell.notebook.cellAt(cellIndex); - const { success } = await executeEphemeralCell(insertedCell); + const { success } = await executeEphemeralCell(insertedCell, execution.token); return success ? { success } : { success: false, error: new Error('Ephemeral cell execution failed') }; }, onLog: (message: string) => { @@ -248,15 +259,27 @@ async function insertEphemeralCell( const EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS = 5 * 60 * 1000; export async function executeEphemeralCell( - cell: NotebookCell + cell: NotebookCell, + token?: CancellationToken ): Promise<{ success: boolean; outputs: unknown[]; executionCount: number | null }> { const completionDeferred = createDeferred(); + const disposables: IDisposable[] = []; + + disposables.push( + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell === cell && e.state === NotebookCellExecutionState.Idle) { + completionDeferred.resolve(); + } + }) + ); - const disposable = notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { - if (e.cell === cell && e.state === NotebookCellExecutionState.Idle) { - completionDeferred.resolve(); + if (token) { + if (token.isCancellationRequested) { + completionDeferred.reject(new CancellationError()); + } else { + disposables.push(token.onCancellationRequested(() => completionDeferred.reject(new CancellationError()))); } - }); + } const timeout = setTimeout(() => { completionDeferred.reject(new Error('Ephemeral cell execution timed out')); @@ -273,7 +296,7 @@ export async function executeEphemeralCell( await completionDeferred.promise; return { - success: cell.executionSummary?.success !== false, + success: cell.executionSummary?.success === true, outputs: cell.outputs.map(translateCellDisplayOutput), executionCount: cell.executionSummary?.executionOrder ?? null }; @@ -284,7 +307,7 @@ export async function executeEphemeralCell( executionCount: null }; } finally { - disposable.dispose(); + dispose(disposables); clearTimeout(timeout); } } diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index b84bb6286c..51b9d089b2 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -1,17 +1,20 @@ import { expect } from 'chai'; import * as sinon from 'sinon'; -import { anything, capture, reset, when } from 'ts-mockito'; -import { NotebookCellOutput, NotebookCellOutputItem, NotebookController } from 'vscode'; +import { anything, capture, instance, mock, reset, when } from 'ts-mockito'; +import { + CancellationTokenSource, + NotebookCellOutput, + NotebookCellOutputItem, + NotebookController, + WorkspaceConfiguration +} from 'vscode'; import type { AgentBlock } from '@deepnote/blocks'; import type { AgentBlockContext, AgentBlockResult } from '@deepnote/runtime-core'; -import { - NotebookCellExecutionState, - notebookCellExecutions -} from '../../platform/notebooks/cellExecutionStateService'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; -import { executeAgentCell, executeEphemeralCell, isAgentCell } from './agentCellExecutionHandler'; +import { executeAgentCell, executeEphemeralCell, getOpenAiApiKey, isAgentCell } from './agentCellExecutionHandler'; import { createMockCell } from './deepnoteTestHelpers'; suite('AgentCellExecutionHandler', () => { @@ -47,6 +50,24 @@ suite('AgentCellExecutionHandler', () => { }); }); + suite('getOpenAiApiKey', () => { + test('returns key when configured', () => { + const mockConfig = mock(); + when(mockConfig.get('agent.openAiApiKey', '')).thenReturn('test-key'); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + + expect(getOpenAiApiKey()).to.equal('test-key'); + }); + + test('throws when key is not set', () => { + const mockConfig = mock(); + when(mockConfig.get('agent.openAiApiKey', '')).thenReturn(''); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + + expect(() => getOpenAiApiKey()).to.throw('deepnote.agent.openAiApiKey is not set'); + }); + }); + suite('executeAgentCell', () => { let mockExecution: { appendOutput: sinon.SinonStub; @@ -58,11 +79,11 @@ suite('AgentCellExecutionHandler', () => { }; let mockController: NotebookController; let executeAgentBlockStub: sinon.SinonStub; - let savedOpenAiKey: string | undefined; setup(() => { - savedOpenAiKey = process.env.OPENAI_API_KEY; - process.env.OPENAI_API_KEY = 'test-key'; + const mockConfig = mock(); + when(mockConfig.get('agent.openAiApiKey', '')).thenReturn('test-key'); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); mockExecution = { appendOutput: sinon.stub().resolves(), @@ -80,14 +101,6 @@ suite('AgentCellExecutionHandler', () => { executeAgentBlockStub = sinon.stub().resolves({ finalOutput: 'done' } as AgentBlockResult); }); - teardown(() => { - if (savedOpenAiKey !== undefined) { - process.env.OPENAI_API_KEY = savedOpenAiKey; - } else { - delete process.env.OPENAI_API_KEY; - } - }); - function createAgentCell(text: string = 'Test prompt') { return createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } }, @@ -243,6 +256,24 @@ suite('AgentCellExecutionHandler', () => { const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); expect(text).to.include('[Agent] Planning next steps...'); }); + + test('ends with failure and writes error when API key is not set', async () => { + const mockConfig = mock(); + when(mockConfig.get('agent.openAiApiKey', '')).thenReturn(''); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.false; + expect(mockExecution.appendOutput.calledOnce).to.be.true; + + const outputs = mockExecution.appendOutput.firstCall.args[0] as NotebookCellOutput[]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('deepnote.agent.openAiApiKey is not set'); + }); }); suite('executeEphemeralCell', () => { @@ -275,5 +306,21 @@ suite('AgentCellExecutionHandler', () => { document: cell.notebook.uri }); }); + + test('returns success false immediately when token is pre-cancelled', async () => { + const cell = createMockCell({ index: 0 }); + const tokenSource = new CancellationTokenSource(); + tokenSource.cancel(); + + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); + + const result = await executeEphemeralCell(cell, tokenSource.token); + + expect(result).to.deep.equal({ + success: false, + outputs: [], + executionCount: null + }); + }); }); }); diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index 8d4ba8eba4..75b3cc1e4e 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -14,14 +14,17 @@ import { workspace } from 'vscode'; import { injectable } from 'inversify'; +import { z } from 'zod'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import type { Pocket } from '../../platform/deepnote/pocket'; +import { logger } from '../../platform/logging'; const DEFAULT_MAX_ITERATIONS = 20; const MIN_ITERATIONS = 1; const MAX_ITERATIONS = 100; const AGENT_MODEL_OPTIONS = ['auto', 'gpt-4o', 'sonnet']; +const MaxIterationsSchema = z.coerce.number().int().min(MIN_ITERATIONS); /** * Provides status bar items for agent cells showing the block type indicator, @@ -67,7 +70,9 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv } public dispose(): void { - this.disposables.forEach((d) => d.dispose()); + for (const disposable of this.disposables) { + disposable.dispose(); + } } public provideCellStatusBarItems( @@ -141,8 +146,16 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv private getMaxIterations(metadata: Record | undefined): number { const value = metadata?.deepnote_max_iterations; - if (typeof value === 'number' && Number.isInteger(value) && value >= MIN_ITERATIONS) { - return value; + const result = MaxIterationsSchema.safeParse(value); + + if (result.success) { + return result.data; + } + + if (value !== undefined) { + logger.debug( + `getMaxIterations: invalid value ${JSON.stringify(value)}, using default ${DEFAULT_MAX_ITERATIONS}` + ); } return DEFAULT_MAX_ITERATIONS; diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts index e5397a1948..62adc562cc 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -214,6 +214,66 @@ suite('AgentCellStatusBarProvider', () => { expect(items[2].text).to.include('Max iterations: 20'); }); + test('Should display default when max iterations is negative', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: -5 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + }); + + test('Should display 1 when max iterations is MIN_ITERATIONS boundary', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: 1 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 1'); + }); + + test('Should display 100 when max iterations is at upper bound', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: 100 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 100'); + }); + + test('Should display default when max iterations is null', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: null + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + }); + + test('Should display default when max iterations is boolean', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: true + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + }); + test('Should have set max iterations command', () => { const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); const items = provider.provideCellStatusBarItems(cell, mockToken)!; From f7bec65bf12b7a41abcdafaa9f3eb55d5d9f7437 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 18 Mar 2026 10:27:20 +0000 Subject: [PATCH 18/19] Add Agent OpenAI API key extension configuration --- package.json | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/package.json b/package.json index 1b31582e7b..dab239bb98 100644 --- a/package.json +++ b/package.json @@ -1638,6 +1638,12 @@ "type": "object", "title": "Deepnote", "properties": { + "deepnote.agent.openAiApiKey": { + "type": "string", + "default": "", + "description": "OpenAI API key for agent cell execution", + "scope": "application" + }, "deepnote.domain": { "type": "string", "default": "deepnote.com", From ea715e798b02090841f6f0a64b7ab23b43e52215 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 18 Mar 2026 17:32:49 +0000 Subject: [PATCH 19/19] Implement OpenAI API key management in Deepnote - Added commands to set and clear the OpenAI API key, enhancing user interaction. - Introduced a new `deepnoteSecretStore` module for managing secrets, including functions to get, set, and clear the OpenAI API key. - Updated `agentCellExecutionHandler` to utilize the new secret management functions, improving error handling when the API key is not set. - Enhanced unit tests to cover the new secret management functionality and ensure robust error handling. --- package.json | 16 +- package.nls.json | 2 + .../deepnote/agentCellExecutionHandler.ts | 24 +- .../agentCellExecutionHandler.unit.test.ts | 110 ++++++-- .../deepnote/agentCellStatusBarProvider.ts | 19 +- src/notebooks/deepnote/deepnoteSecretStore.ts | 122 +++++++++ .../deepnote/deepnoteSecretStore.unit.test.ts | 241 ++++++++++++++++++ .../ephemeralCellDecorationProvider.ts | 4 +- 8 files changed, 487 insertions(+), 51 deletions(-) create mode 100644 src/notebooks/deepnote/deepnoteSecretStore.ts create mode 100644 src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts diff --git a/package.json b/package.json index dab239bb98..d800875f53 100644 --- a/package.json +++ b/package.json @@ -341,6 +341,16 @@ "title": "%deepnote.command.manageAccessToKernels%", "category": "Jupyter" }, + { + "command": "deepnote.setOpenAiApiKey", + "title": "%deepnote.command.setOpenAiApiKey%", + "category": "Deepnote" + }, + { + "command": "deepnote.clearOpenAiApiKey", + "title": "%deepnote.command.clearOpenAiApiKey%", + "category": "Deepnote" + }, { "command": "dataScience.ClearUserProviderJupyterServerCache", "title": "%deepnote.command.dataScience.clearUserProviderJupyterServerCache.title%", @@ -1638,12 +1648,6 @@ "type": "object", "title": "Deepnote", "properties": { - "deepnote.agent.openAiApiKey": { - "type": "string", - "default": "", - "description": "OpenAI API key for agent cell execution", - "scope": "application" - }, "deepnote.domain": { "type": "string", "default": "deepnote.com", diff --git a/package.nls.json b/package.nls.json index 35ee95ae66..07f3b2ba4a 100644 --- a/package.nls.json +++ b/package.nls.json @@ -116,6 +116,8 @@ "deepnote.command.deepnote.openOutlineView.title": "Show Table Of Contents (Outline View)", "deepnote.command.deepnote.openOutlineView.shorttitle": "Outline", "deepnote.command.manageAccessToKernels": "Manage Access To Jupyter Kernels", + "deepnote.command.setOpenAiApiKey": "Set OpenAI API Key", + "deepnote.command.clearOpenAiApiKey": "Clear OpenAI API Key", "deepnote.commandPalette.deepnote.replayPylanceLog.title": "Replay Pylance Log", "deepnote.notebookRenderer.IPyWidget.displayName": "Jupyter IPyWidget Renderer", "deepnote.notebookRenderer.Error.displayName": "Jupyter Error Renderer", diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index e8a1268194..3ed36cf0e2 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -31,6 +31,11 @@ import { logger } from '../../platform/logging'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; import { generateBlockId, generateSortingKey, isEphemeralCell } from './dataConversionUtils'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; +import { getOrPromptOpenAiApiKey } from './deepnoteSecretStore'; + +export async function getOpenAiApiKey(): Promise { + return getOrPromptOpenAiApiKey(); +} export function isAgentCell(cell: NotebookCell): boolean { const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; @@ -63,17 +68,6 @@ export function serializeNotebookContext({ cells }: { cells: NotebookCell[] }): return serializeNotebookContextFromBlocks({ blocks, notebookName: null }); } -export function getOpenAiApiKey(): string { - const config = workspace.getConfiguration('deepnote'); - const key = config.get('agent.openAiApiKey', ''); - - if (!key) { - throw new Error('deepnote.agent.openAiApiKey is not set. Configure it in VS Code settings.'); - } - - return key; -} - export interface ExecuteAgentCellOptions { executeAgentBlockFn?: typeof executeAgentBlock; } @@ -122,7 +116,7 @@ export async function executeAgentCell( cells: cell.notebook.getCells().filter((c) => c.index !== cell.index) }); - const openAiToken = getOpenAiApiKey(); + const openAiToken = await getOpenAiApiKey(); const context: AgentBlockContext = { openAiToken, @@ -139,12 +133,6 @@ export async function executeAgentCell( const { success } = await executeEphemeralCell(insertedCell, execution.token); return success ? { success } : { success: false, error: new Error('Ephemeral cell execution failed') }; }, - onLog: (message: string) => { - logger.info('Agent log', message); - // accumulated += message; - // TODO: replaceOutputItems is Async function - // execution.replaceOutputItems(NotebookCellOutputItem.text(accumulated), output); - }, onAgentEvent: async (event: AgentStreamEvent) => { logger.info('Agent event', JSON.stringify(event)); if (lastAgentEventType != null && lastAgentEventType !== event.type) { diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index 51b9d089b2..b124742244 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -3,21 +3,32 @@ import * as sinon from 'sinon'; import { anything, capture, instance, mock, reset, when } from 'ts-mockito'; import { CancellationTokenSource, + Disposable, + EventEmitter, + ExtensionMode, NotebookCellOutput, NotebookCellOutputItem, NotebookController, - WorkspaceConfiguration + SecretStorage, + SecretStorageChangeEvent } from 'vscode'; import type { AgentBlock } from '@deepnote/blocks'; import type { AgentBlockContext, AgentBlockResult } from '@deepnote/runtime-core'; +import type { IDisposable } from '../../platform/common/types'; +import { IExtensionContext } from '../../platform/common/types'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; +import { dispose } from '../../platform/common/utils/lifecycle'; import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; +import { ServiceContainer } from '../../platform/ioc/container'; import { executeAgentCell, executeEphemeralCell, getOpenAiApiKey, isAgentCell } from './agentCellExecutionHandler'; import { createMockCell } from './deepnoteTestHelpers'; suite('AgentCellExecutionHandler', () => { + const secretStorage = new Map(); + let disposables: IDisposable[] = []; + suite('isAgentCell', () => { test('returns true for cell with agent pocket type', () => { const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); @@ -51,20 +62,47 @@ suite('AgentCellExecutionHandler', () => { }); suite('getOpenAiApiKey', () => { - test('returns key when configured', () => { - const mockConfig = mock(); - when(mockConfig.get('agent.openAiApiKey', '')).thenReturn('test-key'); - when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + setup(() => { + secretStorage.clear(); + const context = mock(); + const secrets = mock(); + const onDidChangeSecrets = new EventEmitter(); + const serviceContainer = mock(); + sinon.stub(ServiceContainer, 'instance').get(() => instance(serviceContainer)); + when(serviceContainer.get(IExtensionContext)).thenReturn(instance(context)); + when(context.extensionMode).thenReturn(ExtensionMode.Production); + when(context.secrets).thenReturn(instance(secrets)); + when(secrets.onDidChange).thenReturn(onDidChangeSecrets.event); + when(secrets.get(anything())).thenCall((key: string) => Promise.resolve(secretStorage.get(key))); + when(secrets.store(anything(), anything())).thenCall((key: string, value: string) => { + secretStorage.set(key, value); + + return Promise.resolve(); + }); + disposables.push(new Disposable(() => sinon.restore())); + }); - expect(getOpenAiApiKey()).to.equal('test-key'); + teardown(() => { + disposables = dispose(disposables); }); - test('throws when key is not set', () => { - const mockConfig = mock(); - when(mockConfig.get('agent.openAiApiKey', '')).thenReturn(''); - when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + test('returns key when configured', async () => { + secretStorage.set('openAiApiKey', 'test-key'); + + const key = await getOpenAiApiKey(); - expect(() => getOpenAiApiKey()).to.throw('deepnote.agent.openAiApiKey is not set'); + expect(key).to.equal('test-key'); + }); + + test('throws when key is not set', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + try { + await getOpenAiApiKey(); + expect.fail('Should have thrown'); + } catch (e) { + expect((e as Error).message).to.include('OpenAI API key is not set'); + } }); }); @@ -81,9 +119,24 @@ suite('AgentCellExecutionHandler', () => { let executeAgentBlockStub: sinon.SinonStub; setup(() => { - const mockConfig = mock(); - when(mockConfig.get('agent.openAiApiKey', '')).thenReturn('test-key'); - when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + secretStorage.clear(); + secretStorage.set('openAiApiKey', 'test-key'); + const context = mock(); + const secrets = mock(); + const onDidChangeSecrets = new EventEmitter(); + const serviceContainer = mock(); + sinon.stub(ServiceContainer, 'instance').get(() => instance(serviceContainer)); + when(serviceContainer.get(IExtensionContext)).thenReturn(instance(context)); + when(context.extensionMode).thenReturn(ExtensionMode.Production); + when(context.secrets).thenReturn(instance(secrets)); + when(secrets.onDidChange).thenReturn(onDidChangeSecrets.event); + when(secrets.get(anything())).thenCall((key: string) => Promise.resolve(secretStorage.get(key))); + when(secrets.store(anything(), anything())).thenCall((key: string, value: string) => { + secretStorage.set(key, value); + + return Promise.resolve(); + }); + disposables.push(new Disposable(() => sinon.restore())); mockExecution = { appendOutput: sinon.stub().resolves(), @@ -101,6 +154,10 @@ suite('AgentCellExecutionHandler', () => { executeAgentBlockStub = sinon.stub().resolves({ finalOutput: 'done' } as AgentBlockResult); }); + teardown(() => { + disposables = dispose(disposables); + }); + function createAgentCell(text: string = 'Test prompt') { return createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } }, @@ -258,9 +315,8 @@ suite('AgentCellExecutionHandler', () => { }); test('ends with failure and writes error when API key is not set', async () => { - const mockConfig = mock(); - when(mockConfig.get('agent.openAiApiKey', '')).thenReturn(''); - when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + secretStorage.clear(); + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); const cell = createAgentCell(); @@ -272,7 +328,7 @@ suite('AgentCellExecutionHandler', () => { const outputs = mockExecution.appendOutput.firstCall.args[0] as NotebookCellOutput[]; const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); - expect(text).to.include('deepnote.agent.openAiApiKey is not set'); + expect(text).to.include('OpenAI API key is not set'); }); }); @@ -314,13 +370,17 @@ suite('AgentCellExecutionHandler', () => { when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); - const result = await executeEphemeralCell(cell, tokenSource.token); - - expect(result).to.deep.equal({ - success: false, - outputs: [], - executionCount: null - }); + try { + const result = await executeEphemeralCell(cell, tokenSource.token); + + expect(result).to.deep.equal({ + success: false, + outputs: [], + executionCount: null + }); + } finally { + tokenSource.dispose(); + } }); }); }); diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index 75b3cc1e4e..0bbc73f3b7 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -19,12 +19,13 @@ import { z } from 'zod'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import type { Pocket } from '../../platform/deepnote/pocket'; import { logger } from '../../platform/logging'; +import { clearOpenAiApiKey, promptForOpenAiApiKey } from './deepnoteSecretStore'; const DEFAULT_MAX_ITERATIONS = 20; const MIN_ITERATIONS = 1; const MAX_ITERATIONS = 100; const AGENT_MODEL_OPTIONS = ['auto', 'gpt-4o', 'sonnet']; -const MaxIterationsSchema = z.coerce.number().int().min(MIN_ITERATIONS); +const MaxIterationsSchema = z.coerce.number().int().min(MIN_ITERATIONS).max(MAX_ITERATIONS); /** * Provides status bar items for agent cells showing the block type indicator, @@ -66,6 +67,22 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv }) ); + this.disposables.push( + commands.registerCommand('deepnote.setOpenAiApiKey', async () => { + const key = await promptForOpenAiApiKey(); + if (key) { + void window.showInformationMessage(l10n.t('OpenAI API key has been saved.')); + } + }) + ); + + this.disposables.push( + commands.registerCommand('deepnote.clearOpenAiApiKey', async () => { + await clearOpenAiApiKey(); + void window.showInformationMessage(l10n.t('OpenAI API key has been cleared.')); + }) + ); + this.disposables.push(this._onDidChangeCellStatusBarItems); } diff --git a/src/notebooks/deepnote/deepnoteSecretStore.ts b/src/notebooks/deepnote/deepnoteSecretStore.ts new file mode 100644 index 0000000000..fadf11bd42 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteSecretStore.ts @@ -0,0 +1,122 @@ +import { ExtensionMode, l10n, window } from 'vscode'; + +import { ServiceContainer } from '../../platform/ioc/container'; +import { IExtensionContext } from '../../platform/common/types'; + +export interface SecretPromptOptions { + prompt: string; + placeHolder?: string; + password?: boolean; +} + +function getContext(): IExtensionContext | null { + const context = ServiceContainer.instance.get(IExtensionContext); + + if (context.extensionMode === ExtensionMode.Test) { + return null; + } + + return context; +} + +export async function getSecret(key: string): Promise { + const context = getContext(); + + if (!context) { + return undefined; + } + + const value = await context.secrets.get(key); + + return value && value.length > 0 ? value : undefined; +} + +export async function setSecret(key: string, value: string): Promise { + const context = getContext(); + + if (!context) { + return; + } + + await context.secrets.store(key, value); +} + +export async function clearSecret(key: string): Promise { + const context = getContext(); + + if (!context) { + return; + } + + await context.secrets.delete(key); +} + +export async function promptForSecret(key: string, options: SecretPromptOptions): Promise { + const input = await window.showInputBox({ + prompt: options.prompt, + placeHolder: options.placeHolder, + password: options.password ?? true, + ignoreFocusOut: true + }); + + if (!input || input.trim().length === 0) { + return undefined; + } + + const trimmed = input.trim(); + await setSecret(key, trimmed); + + return trimmed; +} + +export async function getOrPromptSecret( + key: string, + options: SecretPromptOptions, + errorMessage: string +): Promise { + let value = await getSecret(key); + + if (!value) { + value = await promptForSecret(key, options); + } + + if (!value) { + throw new Error(errorMessage); + } + + return value; +} + +// OpenAI API key - specific wrappers + +const OPENAI_API_KEY = 'openAiApiKey'; + +const OPENAI_PROMPT_OPTIONS: SecretPromptOptions = { + prompt: l10n.t('Enter your OpenAI API key'), + placeHolder: l10n.t('sk-...'), + password: true +}; + +export async function getOpenAiApiKey(): Promise { + return getSecret(OPENAI_API_KEY); +} + +export async function setOpenAiApiKey(key: string): Promise { + return setSecret(OPENAI_API_KEY, key); +} + +export async function clearOpenAiApiKey(): Promise { + return clearSecret(OPENAI_API_KEY); +} + +export async function promptForOpenAiApiKey(): Promise { + return promptForSecret(OPENAI_API_KEY, OPENAI_PROMPT_OPTIONS); +} + +export async function getOrPromptOpenAiApiKey(): Promise { + return getOrPromptSecret( + OPENAI_API_KEY, + OPENAI_PROMPT_OPTIONS, + l10n.t('OpenAI API key is not set. Use the command "Deepnote: Set OpenAI API Key" to configure it.') + ); +} diff --git a/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts new file mode 100644 index 0000000000..e99bafc638 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts @@ -0,0 +1,241 @@ +import { expect } from 'chai'; +import * as sinon from 'sinon'; +import { anything, instance, mock, when } from 'ts-mockito'; +import { EventEmitter, ExtensionMode, SecretStorage, SecretStorageChangeEvent } from 'vscode'; + +import { IExtensionContext } from '../../platform/common/types'; +import { ServiceContainer } from '../../platform/ioc/container'; +import { + clearOpenAiApiKey, + clearSecret, + getOpenAiApiKey, + getOrPromptOpenAiApiKey, + getOrPromptSecret, + getSecret, + promptForOpenAiApiKey, + promptForSecret, + setOpenAiApiKey, + setSecret +} from './deepnoteSecretStore'; +import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; + +suite('deepnoteSecretStore', () => { + const secretStorage = new Map(); + let context: IExtensionContext; + let secrets: SecretStorage; + let onDidChangeSecrets: EventEmitter; + + setup(() => { + secretStorage.clear(); + context = mock(); + secrets = mock(); + onDidChangeSecrets = new EventEmitter(); + + const serviceContainer = mock(); + sinon.stub(ServiceContainer, 'instance').get(() => instance(serviceContainer)); + when(serviceContainer.get(IExtensionContext)).thenReturn(instance(context)); + when(context.extensionMode).thenReturn(ExtensionMode.Production); + when(context.secrets).thenReturn(instance(secrets)); + when(secrets.onDidChange).thenReturn(onDidChangeSecrets.event); + when(secrets.get(anything())).thenCall((key: string) => Promise.resolve(secretStorage.get(key))); + when(secrets.store(anything(), anything())).thenCall((key: string, value: string) => { + secretStorage.set(key, value); + onDidChangeSecrets.fire({ key }); + + return Promise.resolve(); + }); + when(secrets.delete(anything())).thenCall((key: string) => { + secretStorage.delete(key); + + return Promise.resolve(); + }); + }); + + teardown(() => { + sinon.restore(); + }); + + suite('generic getSecret', () => { + test('returns value when stored', async () => { + secretStorage.set('customKey', 'custom-value'); + + const value = await getSecret('customKey'); + + expect(value).to.equal('custom-value'); + }); + + test('returns undefined when not set', async () => { + const value = await getSecret('customKey'); + + expect(value).to.be.undefined; + }); + + test('returns undefined when value is empty string', async () => { + secretStorage.set('customKey', ''); + + const value = await getSecret('customKey'); + + expect(value).to.be.undefined; + }); + }); + + suite('generic setSecret', () => { + test('stores value in secrets', async () => { + await setSecret('customKey', 'custom-value'); + + expect(secretStorage.get('customKey')).to.equal('custom-value'); + }); + }); + + suite('generic clearSecret', () => { + test('deletes value from secrets', async () => { + secretStorage.set('customKey', 'custom-value'); + + await clearSecret('customKey'); + + expect(secretStorage.has('customKey')).to.be.false; + }); + }); + + suite('generic promptForSecret', () => { + test('stores and returns value when user enters input', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('user-input')); + + const value = await promptForSecret('customKey', { + prompt: 'Enter value', + placeHolder: 'placeholder', + password: false + }); + + expect(value).to.equal('user-input'); + expect(secretStorage.get('customKey')).to.equal('user-input'); + }); + + test('returns undefined when user cancels', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const value = await promptForSecret('customKey', { prompt: 'Enter value' }); + + expect(value).to.be.undefined; + }); + }); + + suite('generic getOrPromptSecret', () => { + test('returns value when present in store', async () => { + secretStorage.set('customKey', 'stored-value'); + + const value = await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); + + expect(value).to.equal('stored-value'); + }); + + test('throws when value missing and user cancels prompt', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + try { + await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); + expect.fail('Should have thrown'); + } catch (e) { + expect((e as Error).message).to.equal('Value is required'); + } + }); + }); + + suite('getOpenAiApiKey', () => { + test('returns key when stored', async () => { + secretStorage.set('openAiApiKey', 'test-key'); + + const key = await getOpenAiApiKey(); + + expect(key).to.equal('test-key'); + }); + + test('returns undefined when not set', async () => { + const key = await getOpenAiApiKey(); + + expect(key).to.be.undefined; + }); + + test('returns undefined when key is empty string', async () => { + secretStorage.set('openAiApiKey', ''); + + const key = await getOpenAiApiKey(); + + expect(key).to.be.undefined; + }); + }); + + suite('setOpenAiApiKey', () => { + test('stores key in secrets', async () => { + await setOpenAiApiKey('my-api-key'); + + expect(secretStorage.get('openAiApiKey')).to.equal('my-api-key'); + }); + }); + + suite('clearOpenAiApiKey', () => { + test('deletes key from secrets', async () => { + secretStorage.set('openAiApiKey', 'test-key'); + + await clearOpenAiApiKey(); + + expect(secretStorage.has('openAiApiKey')).to.be.false; + }); + }); + + suite('promptForOpenAiApiKey', () => { + test('stores and returns key when user enters value', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('sk-abc123')); + + const key = await promptForOpenAiApiKey(); + + expect(key).to.equal('sk-abc123'); + expect(secretStorage.get('openAiApiKey')).to.equal('sk-abc123'); + }); + + test('returns undefined when user cancels', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const key = await promptForOpenAiApiKey(); + + expect(key).to.be.undefined; + }); + + test('returns undefined when user enters empty string', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(' ')); + + const key = await promptForOpenAiApiKey(); + + expect(key).to.be.undefined; + }); + }); + + suite('getOrPromptOpenAiApiKey', () => { + test('returns key when present in store', async () => { + secretStorage.set('openAiApiKey', 'stored-key'); + + const key = await getOrPromptOpenAiApiKey(); + + expect(key).to.equal('stored-key'); + }); + + test('prompts and returns key when missing', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('prompted-key')); + + const key = await getOrPromptOpenAiApiKey(); + + expect(key).to.equal('prompted-key'); + }); + + test('throws when key missing and user cancels prompt', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + try { + await getOrPromptOpenAiApiKey(); + expect.fail('Should have thrown'); + } catch (e) { + expect((e as Error).message).to.include('OpenAI API key is not set'); + } + }); + }); +}); diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts index 36b5d24053..19ca8b76fc 100644 --- a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts @@ -67,7 +67,9 @@ export class EphemeralCellDecorationProvider implements IExtensionSyncActivation } public dispose(): void { - this.disposables.forEach((d) => d.dispose()); + for (const disposable of this.disposables) { + disposable.dispose(); + } } private findCellForEditor(editor: TextEditor): NotebookCell | undefined {