From ffe09588ac299d7ad88f852d707ee91eb5041181 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 9 Mar 2026 08:30:26 +0000 Subject: [PATCH 01/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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/24] 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 { From 3b2b95043ab7a6868fe768d49487600e1c78839d Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 31 Jul 2026 13:02:43 +0000 Subject: [PATCH 20/24] fix(agent-block): adapt agent cell execution to runtime-core 0.4.0 API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PR #358 pinned @deepnote/runtime-core ^0.2.0, which exports no agent API at all, so the branch never compiled on its own — it was written against an unreleased build. Merging main moves the dependency to ^0.4.0, which does ship the agent API but with a tightened contract: - serializeNotebookContextFromBlocks() no longer accepts a null notebookName, so pass the document's deepnoteNotebookName (empty string when absent). - The addMarkdownBlock / addAndExecuteCodeBlock tool callbacks now return a string rather than a {success} object. That string is the tool result fed back to the model, so mirror the wording runtime-core uses in its own ExecutionEngine implementation of the same tools: the agent now sees the executed cell's real output instead of only whether it succeeded. extractOutputsText() reads a stream output's `text` only when it is a string, but translateCellDisplayOutput() emits nbformat's multiline array form, so normalize before extracting — otherwise every print() from an ephemeral cell would be dropped from the tool result. Also teach the runtime-core test mock about the agent exports. The mock is a main-only file the branch never had, and the ESM loader swaps it in for the whole module, so without them the import binding fails and no unit test in the suite can load. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Bz32QuAABW7Xk9wKXc7Sk8 --- .../deepnote/agentCellExecutionHandler.ts | 61 ++++++++++++++++--- src/test/mocks/deepnoteRuntimeCore.ts | 33 +++++++++- 2 files changed, 83 insertions(+), 11 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index 3ed36cf0e2..50057a2eec 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -13,7 +13,7 @@ import { workspace } from 'vscode'; -import { AgentBlock, DeepnoteBlock } from '@deepnote/blocks'; +import { AgentBlock, DeepnoteBlock, extractOutputsText } from '@deepnote/blocks'; import { AgentBlockContext, AgentStreamEvent, @@ -33,6 +33,12 @@ import { generateBlockId, generateSortingKey, isEphemeralCell } from './dataConv import { DeepnoteDataConverter } from './deepnoteDataConverter'; import { getOrPromptOpenAiApiKey } from './deepnoteSecretStore'; +// Tool results reported back to the agent. These mirror the wording @deepnote/runtime-core uses in +// its own ExecutionEngine implementation of the same tools, so the agent sees identical phrasing +// whether a block runs in the extension or on the backend. +const MARKDOWN_BLOCK_ADDED_TEXT = 'Markdown block added.'; +const NO_OUTPUT_TEXT = '(no output)'; + export async function getOpenAiApiKey(): Promise { return getOrPromptOpenAiApiKey(); } @@ -43,7 +49,13 @@ export function isAgentCell(cell: NotebookCell): boolean { return pocket?.type === 'agent'; } -export function serializeNotebookContext({ cells }: { cells: NotebookCell[] }): string { +export function serializeNotebookContext({ + cells, + notebookName +}: { + cells: NotebookCell[]; + notebookName: string; +}): string { const converter = new DeepnoteDataConverter(); const blocks = cells.reduce((acc, cell) => { @@ -65,7 +77,28 @@ export function serializeNotebookContext({ cells }: { cells: NotebookCell[] }): return acc; }, []); - return serializeNotebookContextFromBlocks({ blocks, notebookName: null }); + return serializeNotebookContextFromBlocks({ blocks, notebookName }); +} + +/** + * `translateCellDisplayOutput` follows nbformat's multiline convention and emits stream `text` as an + * array of lines. `extractOutputsText` only reads `text` when it is a string, so join it first — + * otherwise every `print()` an ephemeral cell produces would be dropped from the agent's tool result. + */ +function normalizeOutputsForTextExtraction(outputs: unknown[]): unknown[] { + return outputs.map((output) => { + const candidate = output as { output_type?: unknown; text?: unknown } | null; + + if (candidate?.output_type === 'stream' && Array.isArray(candidate.text)) { + return { ...candidate, text: candidate.text.join('') }; + } + + return output; + }); +} + +function describeExecutionOutputs(outputs: unknown[]): string { + return extractOutputsText(normalizeOutputsForTextExtraction(outputs), { includeTraceback: true }) || NO_OUTPUT_TEXT; } export interface ExecuteAgentCellOptions { @@ -113,7 +146,8 @@ export async function executeAgentCell( let lastAgentEventType: AgentStreamEvent['type'] | undefined; const notebookContext = serializeNotebookContext({ - cells: cell.notebook.getCells().filter((c) => c.index !== cell.index) + cells: cell.notebook.getCells().filter((c) => c.index !== cell.index), + notebookName: (cell.notebook.metadata?.deepnoteNotebookName as string | undefined) ?? '' }); const openAiToken = await getOpenAiApiKey(); @@ -124,14 +158,23 @@ export async function executeAgentCell( notebookContext, addMarkdownBlock: async ({ content }: { content: string }) => { await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'markdown', content); - return { success: true }; + + return MARKDOWN_BLOCK_ADDED_TEXT; }, addAndExecuteCodeBlock: async ({ code }: { code: string }) => { - const cellIndex = await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'code', code); - const insertedCell = cell.notebook.cellAt(cellIndex); + try { + const cellIndex = await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'code', code); + const insertedCell = cell.notebook.cellAt(cellIndex); + + const { success, outputs } = await executeEphemeralCell(insertedCell, execution.token); + const outputText = describeExecutionOutputs(outputs); - const { success } = await executeEphemeralCell(insertedCell, execution.token); - return success ? { success } : { success: false, error: new Error('Ephemeral cell execution failed') }; + return success ? `Output:\n${outputText}` : `Execution failed:\n${outputText}`; + } catch (error) { + const executionError = error instanceof Error ? error : new Error(String(error)); + + return `Execution error: ${executionError.message}`; + } }, onAgentEvent: async (event: AgentStreamEvent) => { logger.info('Agent event', JSON.stringify(event)); diff --git a/src/test/mocks/deepnoteRuntimeCore.ts b/src/test/mocks/deepnoteRuntimeCore.ts index 26e25d6797..9470a8a8ce 100644 --- a/src/test/mocks/deepnoteRuntimeCore.ts +++ b/src/test/mocks/deepnoteRuntimeCore.ts @@ -1,9 +1,11 @@ -import type { ServerInfo, ServerOptions } from '@deepnote/runtime-core'; +import type { AgentBlockContext, ServerInfo, ServerOptions } from '@deepnote/runtime-core'; +import type { AgentBlock } from '@deepnote/blocks'; import type { ChildProcess } from 'child_process'; /** * Mock of @deepnote/runtime-core for unit tests: the real startServer/stopServer spawn and - * kill Python processes, so this records calls and returns fake server info instead. + * kill Python processes, and the real executeAgentBlock calls the OpenAI API, so this records + * calls and returns fake results instead. * * build/mocha-esm-loader.js resolves the '@deepnote/runtime-core' specifier to this module, * so code under test and tests importing the __ helpers below share one module instance. @@ -12,6 +14,8 @@ import type { ChildProcess } from 'child_process'; type RuntimeCore = typeof import('@deepnote/runtime-core'); +const executeAgentBlockCalls: { block: AgentBlock; context: AgentBlockContext }[] = []; +const serializeNotebookContextFromBlocksCalls: { blockCount: number; notebookName: string }[] = []; const startServerCalls: ServerOptions[] = []; const stopServerCalls: ServerInfo[] = []; let nextServerId = 0; @@ -27,6 +31,21 @@ function makeFakeProcess(id: number): ChildProcess { } as unknown as ChildProcess; } +export const executeAgentBlock: RuntimeCore['executeAgentBlock'] = async (block, context) => { + executeAgentBlockCalls.push({ block, context }); + + return { finalOutput: '' }; +}; + +export const serializeNotebookContextFromBlocks: RuntimeCore['serializeNotebookContextFromBlocks'] = ({ + blocks, + notebookName +}) => { + serializeNotebookContextFromBlocksCalls.push({ blockCount: blocks.length, notebookName }); + + return `notebook:${notebookName} blocks:${blocks.length}`; +}; + export const startServer: RuntimeCore['startServer'] = async (options) => { startServerCalls.push(options); @@ -49,6 +68,14 @@ export const stopServer: RuntimeCore['stopServer'] = async (info) => { }; // Test-only helpers (prefixed with __ to signal they are not part of the real API). +export function __getExecuteAgentBlockCalls(): { block: AgentBlock; context: AgentBlockContext }[] { + return executeAgentBlockCalls; +} + +export function __getSerializeNotebookContextFromBlocksCalls(): { blockCount: number; notebookName: string }[] { + return serializeNotebookContextFromBlocksCalls; +} + export function __getStartServerCalls(): ServerOptions[] { return startServerCalls; } @@ -62,6 +89,8 @@ export function __setStartServerImpl(impl: RuntimeCore['startServer'] | null): v } export function __resetRuntimeCoreMock(): void { + executeAgentBlockCalls.length = 0; + serializeNotebookContextFromBlocksCalls.length = 0; startServerCalls.length = 0; stopServerCalls.length = 0; nextServerId = 0; From a010e0778891931e4d47f48019d553005d44e67b Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 31 Jul 2026 13:05:38 +0000 Subject: [PATCH 21/24] fix(agent-block): reject boolean max-iteration metadata MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit z.coerce.number() turns true into 1, which then satisfies .int().min(1).max(100), so a boolean deepnote_max_iterations was accepted as an iteration count of 1 instead of falling back to the default of 20 — contradicting the test that already documented the intended behaviour. Pre-existing on the branch rather than a merge regression; it only surfaced now because the branch compiles for the first time, so its tests could finally run. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Bz32QuAABW7Xk9wKXc7Sk8 --- src/notebooks/deepnote/agentCellStatusBarProvider.ts | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index 0bbc73f3b7..a168ea6a1f 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -163,9 +163,11 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv private getMaxIterations(metadata: Record | undefined): number { const value = metadata?.deepnote_max_iterations; - const result = MaxIterationsSchema.safeParse(value); + // z.coerce.number() turns true into 1, which then satisfies the range check, so booleans + // would be accepted as an iteration count instead of falling back to the default. + const result = typeof value === 'boolean' ? undefined : MaxIterationsSchema.safeParse(value); - if (result.success) { + if (result?.success) { return result.data; } From cbbf91c98767a5e2154d745b0c7c295a78522db6 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 31 Jul 2026 16:21:53 +0000 Subject: [PATCH 22/24] Add hubot to cspell config --- cspell.json | 1 + 1 file changed, 1 insertion(+) diff --git a/cspell.json b/cspell.json index 8aa54f8759..91b895eca3 100644 --- a/cspell.json +++ b/cspell.json @@ -49,6 +49,7 @@ "evalue", "findstr", "getsitepackages", + "hubot", "IMAGENAME", "ipykernel", "ipynb", From cdc2a46b82ccda97e8d69527e1da1637b051a76b Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 31 Jul 2026 16:29:57 +0000 Subject: [PATCH 23/24] fix(agent-block): type the markdown-it renderer hook instead of using any MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ephemeral-cell markdown wrapper reached for `any` in four places, which @typescript-eslint/no-explicit-any rejects — the rule is 'error' repo-wide and is only relaxed for tests and *.d.ts, so this failed CI lint. None of the casts were necessary. RendererApi exposes extension hooks through an index signature, so extendMarkdownIt already arrives as `unknown` and just needs narrowing to a call signature. markdown-it itself ships no types and is only a transitive dependency, so describe the small surface this renderer actually touches rather than pulling in @types/markdown-it. Narrowing on `typeof extendMarkdownIt === 'function'` also replaces a bare truthiness check on the renderer, so a markdown renderer without the hook no longer throws. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01Bz32QuAABW7Xk9wKXc7Sk8 --- src/renderers/client/markdown.ts | 38 ++++++++++++++++++++++++++++---- 1 file changed, 34 insertions(+), 4 deletions(-) diff --git a/src/renderers/client/markdown.ts b/src/renderers/client/markdown.ts index 8c69bfd618..f18a4392df 100644 --- a/src/renderers/client/markdown.ts +++ b/src/renderers/client/markdown.ts @@ -1,5 +1,31 @@ import type { ActivationFunction } from 'vscode-notebook-renderer'; +// markdown-it ships no type declarations and is only a transitive dependency, so describe the +// small surface this renderer touches rather than depending on its internals wholesale. +interface MarkdownItToken { + content: string; +} + +interface MarkdownItRuleState { + Token: new (type: string, tag: string, nesting: number) => MarkdownItToken; + env?: { + outputItem?: { + metadata?: Record; + }; + }; + tokens: MarkdownItToken[]; +} + +interface MarkdownIt { + core: { + ruler: { + push(name: string, rule: (state: MarkdownItRuleState) => void): void; + }; + }; +} + +type ExtendMarkdownIt = (callback: (md: MarkdownIt) => void) => void; + const styleContent = ` .alert { width: auto; @@ -61,8 +87,12 @@ export const activate: ActivationFunction = async (ctx) => { document.head.appendChild(template); const markdownRenderer = await ctx.getRenderer('vscode.markdown-it-renderer'); - if (markdownRenderer) { - (markdownRenderer as any).extendMarkdownIt((md: any) => { + // RendererApi exposes extension hooks through an index signature, so extendMarkdownIt arrives + // as unknown and has to be narrowed before it can be called. + const extendMarkdownIt = markdownRenderer?.extendMarkdownIt as ExtendMarkdownIt | undefined; + + if (typeof extendMarkdownIt === 'function') { + extendMarkdownIt((md) => { addEphemeralCellWrapper(md); }); } @@ -70,8 +100,8 @@ export const activate: ActivationFunction = async (ctx) => { return undefined; }; -function addEphemeralCellWrapper(md: any): void { - md.core.ruler.push('ephemeral_wrapper', (state: any) => { +function addEphemeralCellWrapper(md: MarkdownIt): void { + md.core.ruler.push('ephemeral_wrapper', (state) => { const metadata = state.env?.outputItem?.metadata; if (!metadata || metadata.is_ephemeral !== true) { return; From 2567e383ce787b8e57a38e03114b0de0b851d13c Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 3 Aug 2026 12:47:27 +0000 Subject: [PATCH 24/24] fix(agent-block): address pre-merge review findings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Remediation for 13 findings from a merged, adversarially verified review of this branch. Grouped here because the changes interleave across the same files; each is independently described below. Correctness - Model picker read/wrote `deepnote_model`, but execution and the file schema use `deepnote_agent_model` — the picker was inert in both directions and persisted a key no consumer reads into the .deepnote file. Route both sides through one constant, and write the literal 'auto' rather than deleting the key: convertCellToBlock does not re-run the zod schema, so a missing key reaches runtime-core as undefined and is passed to openai() as a model name. Drop 'sonnet' from the options — createOpenAI gets no baseURL here, so an Anthropic model name 404s. - Remove the max-iterations control. Nothing consumes `deepnote_max_iterations`: executeAgentBlock hardcodes maxTurns = 10 and AgentBlockContext exposes no turn limit, so the UI advertised a default of 20 and a 1-100 range over a setting that changed nothing. - Rich outputs reached the agent comma-mangled. translateCellDisplayOutput splits `text/*` into nbformat line arrays for execute_result/display_data too, and extractOutputText stringifies them with String(...), which joins with commas — so every df.head() the agent read had a comma glued to the start of each line after the first. Only stream text was being joined. - insertEphemeralCell ignored applyEdit's result and returned a bare index. cellAt clamps rather than throwing, so a rejected edit or a concurrent structural change handed back a pre-existing user cell, which the agent then executed and reported as its own result. Check the boolean and resolve the inserted cell by __deepnoteBlockId instead. - Both execute handlers ran every agent cell before any kernel cell, regardless of document order, so agent-generated code executed against a kernel that had not run the setup cells above it. Walk in document order. Reliability - executeEphemeralCell rejected the completion deferred on an already cancelled token but still dispatched notebook.cell.execute, so the kernel ran the generated code after the user cancelled. Throw before dispatch. Also propagate the failure reason instead of collapsing cancellation, timeout and command failure alike into "(no output)", which invited the agent to retry. - Run All aborted silently when an agent cell had already deleted the ephemeral cells queued alongside it: createNotebookCellExecution throws for a removed cell, and nothing caught it. Filter out cells whose index is -1. - Acquire the OpenAI key before the destructive ephemeral cleanup. It prompts and throws on dismissal, so cancelling the prompt destroyed the previous run's results for a run that never started. Security - The placeholder execute handler had no workspace-trust check, while the real controller did. Agent blocks can spawn MCP servers declared in the project file, so gate both paths — the manifest already promises cell execution is unsupported in untrusted workspaces. - Pass project-level `mcpServers` from project.settings, matching what the CLI's ExecutionEngine provides. The empty array was a stub from the initial implementation; it dropped the project-level tier (declared in the file, intended to be configured) while runtime-core still merged in the block-level tier, which is the invisible one. Performance - Stream agent output as incremental stdout items rather than re-encoding and re-sending the whole transcript on every token: that was O(n^2) bytes across the extension-host boundary, and since runtime-core awaits onAgentEvent inside its stream loop the cost was added to the run's wall clock. Maintainability and tests - Move isAgentCell next to isEphemeralCell in dataConversionUtils so the status bar provider stops duplicating it and no longer needs the runtime-core-backed handler module in its import graph. Delete the getOpenAiApiKey wrapper, whose name collided with a differently-behaving export one import away. - Log stream events at trace with type only. At info they wrote model text and reasoning into the user-visible output channel on every token. Keep the explicit stack log: logger.error(msg, error) only renders the stack for errors branded isJupyterError. - createMockNotebook now reads through a caller-supplied cells array, so tests can exercise the insert/remove/ordering logic that previously had no coverage at all. Adds regression tests for the output mangling, the cancelled key prompt, failed inserts, cross-agent deletion and delta streaming; each was run against the unfixed code and observed failing. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- .../controllers/vscodeNotebookController.ts | 39 +- .../deepnote/agentCellExecutionHandler.ts | 205 +++++++--- .../agentCellExecutionHandler.unit.test.ts | 351 ++++++++++++++---- .../deepnote/agentCellStatusBarProvider.ts | 126 +------ .../agentCellStatusBarProvider.unit.test.ts | 153 +------- src/notebooks/deepnote/dataConversionUtils.ts | 14 + .../deepnoteKernelAutoSelector.node.ts | 35 +- src/notebooks/deepnote/deepnoteTestHelpers.ts | 21 +- 8 files changed, 538 insertions(+), 406 deletions(-) diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index e716ee4179..ae64367887 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -625,21 +625,42 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont if (!this.cellQueue.has(doc)) { return; } - // 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 allCells = this.cellQueue.get(doc) || []; + // Cleared before any await so the re-entrant execute request an agent cell issues for its + // generated code starts from an empty queue. this.cellQueue.delete(doc); - const agentCells = allCells.filter((cell) => isAgentCell(cell)); - const kernelCells = allCells.filter((cell) => !isAgentCell(cell)); + // Walk in document order rather than running every agent cell first: an agent executes the + // code it generates against the kernel immediately, so it must not overtake the cells above + // it that set up the state it reads. + let pendingKernelCells: NotebookCell[] = []; - // 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); + for (const cell of allCells) { + if (!isAgentCell(cell)) { + pendingKernelCells.push(cell); + continue; + } + + await this.executeKernelCells(doc, pendingKernelCells); + pendingKernelCells = []; + + logger.trace(`Executing agent cell ${cell.index} for ${getDisplayPath(doc.uri)} without kernel`); + await executeAgentCell(cell, this.controller).catch(noop); } + await this.executeKernelCells(doc, pendingKernelCells); + } + + private async executeKernelCells(doc: NotebookDocument, cells: NotebookCell[]) { + // 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 }; + + // An agent run deletes the ephemeral cells it produced last time, and those are ordinary code + // cells that Run All queues. createNotebookCellExecution throws for a cell that has since been + // removed, which would abort the rest of the batch. + const kernelCells = cells.filter((cell) => cell.index >= 0); + if (kernelCells.length === 0) { return; } diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index 50057a2eec..d6d2706357 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -26,29 +26,50 @@ 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 { ServiceContainer } from '../../platform/ioc/container'; import { logger } from '../../platform/logging'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; -import { generateBlockId, generateSortingKey, isEphemeralCell } from './dataConversionUtils'; +import { IDeepnoteNotebookManager } from '../types'; +import { generateBlockId, generateSortingKey, isAgentCell, isEphemeralCell } from './dataConversionUtils'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; import { getOrPromptOpenAiApiKey } from './deepnoteSecretStore'; +export { isAgentCell }; + +/** + * Project-level MCP servers declared in the `.deepnote` file, matching what the CLI's ExecutionEngine + * passes. `executeAgentBlock` merges these with any block-level `deepnote_mcp_servers` (block wins on + * name), so leaving this empty silently drops the project-level half of that contract. + * + * Spawning these is arbitrary local command execution declared by a workspace file, so every caller + * must already be behind a `workspace.isTrusted` check. + */ +function getProjectMcpServers(notebook: NotebookDocument): AgentBlockContext['mcpServers'] { + const projectId = notebook.metadata?.deepnoteProjectId as string | undefined; + const notebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; + + if (!projectId || !notebookId) { + return []; + } + + const manager = ServiceContainer.instance.tryGet(IDeepnoteNotebookManager); + const servers = manager?.getProjectForNotebook(projectId, notebookId)?.project.settings?.mcpServers ?? []; + + if (servers.length > 0) { + logger.info( + `Agent cell: using ${servers.length} project MCP server(s): ${servers.map((s) => s.name).join(', ')}` + ); + } + + return servers; +} + // Tool results reported back to the agent. These mirror the wording @deepnote/runtime-core uses in // its own ExecutionEngine implementation of the same tools, so the agent sees identical phrasing // whether a block runs in the extension or on the backend. const MARKDOWN_BLOCK_ADDED_TEXT = 'Markdown block added.'; const NO_OUTPUT_TEXT = '(no output)'; -export async function getOpenAiApiKey(): Promise { - return getOrPromptOpenAiApiKey(); -} - -export function isAgentCell(cell: NotebookCell): boolean { - const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; - - return pocket?.type === 'agent'; -} - export function serializeNotebookContext({ cells, notebookName @@ -80,24 +101,46 @@ export function serializeNotebookContext({ return serializeNotebookContextFromBlocks({ blocks, notebookName }); } +function joinMultilineString(value: unknown): unknown { + return Array.isArray(value) && value.every((entry) => typeof entry === 'string') ? value.join('') : value; +} + /** - * `translateCellDisplayOutput` follows nbformat's multiline convention and emits stream `text` as an - * array of lines. `extractOutputsText` only reads `text` when it is a string, so join it first — - * otherwise every `print()` an ephemeral cell produces would be dropped from the agent's tool result. + * `translateCellDisplayOutput` follows nbformat's multiline convention and emits text as an array of + * lines — both for stream `text` and for the `text/*` entries of `execute_result`/`display_data` + * `data`. `extractOutputsText` reads stream text only when it is a string, and stringifies + * `data['text/plain']` with `String(...)`, which joins an array with commas. Join the lines first so + * `print()` output isn't dropped and a `df.head()` repr doesn't reach the agent with a comma glued to + * the start of every line. */ function normalizeOutputsForTextExtraction(outputs: unknown[]): unknown[] { return outputs.map((output) => { - const candidate = output as { output_type?: unknown; text?: unknown } | null; + const candidate = output as { output_type?: unknown; text?: unknown; data?: unknown } | null; + + if (candidate?.output_type === 'stream') { + return { ...candidate, text: joinMultilineString(candidate.text) }; + } + + if ( + (candidate?.output_type === 'execute_result' || candidate?.output_type === 'display_data') && + candidate.data != null && + typeof candidate.data === 'object' + ) { + const data = Object.fromEntries( + Object.entries(candidate.data).map(([mime, value]) => [ + mime, + mime.startsWith('text/') ? joinMultilineString(value) : value + ]) + ); - if (candidate?.output_type === 'stream' && Array.isArray(candidate.text)) { - return { ...candidate, text: candidate.text.join('') }; + return { ...candidate, data }; } return output; }); } -function describeExecutionOutputs(outputs: unknown[]): string { +export function describeExecutionOutputs(outputs: unknown[]): string { return extractOutputsText(normalizeOutputsForTextExtraction(outputs), { includeTraceback: true }) || NO_OUTPUT_TEXT; } @@ -119,8 +162,12 @@ export async function executeAgentCell( const prompt = cell.document.getText(); - let accumulated = `[Agent] Planning next steps...`; - const output = new NotebookCellOutput([NotebookCellOutputItem.text(accumulated)]); + // Streamed as stdout items so each event can be appended rather than re-sending the whole + // transcript: `NotebookCellOutputItem.text` re-encodes the full buffer on every token, which + // is O(n²) bytes across the extension-host boundary — and since runtime-core awaits + // `onAgentEvent` inside its stream loop, that cost is added to the run's wall clock. + // The stdout mime is the one the renderer concatenates, matching how kernel output streams. + const output = new NotebookCellOutput([NotebookCellOutputItem.stdout(`[Agent] Planning next steps...`)]); await execution.replaceOutput([output]); const dataConverter = new DeepnoteDataConverter(); @@ -141,33 +188,48 @@ export async function executeAgentCell( throw new Error('Cell is not an agent cell'); } + // Acquire the key before the destructive cleanup below: it prompts, and throws when the user + // dismisses the prompt, which would otherwise leave the previous run's cells already deleted. + const openAiToken = await getOrPromptOpenAiApiKey(); + await removeEphemeralCellsForAgent(cell.notebook, agentBlock.id); let lastAgentEventType: AgentStreamEvent['type'] | undefined; + // Must run after the removal — serializeNotebookContextFromBlocks does no ephemeral + // filtering, so the agent would otherwise be handed its own previous scratch cells. const notebookContext = serializeNotebookContext({ cells: cell.notebook.getCells().filter((c) => c.index !== cell.index), notebookName: (cell.notebook.metadata?.deepnoteNotebookName as string | undefined) ?? '' }); - const openAiToken = await getOpenAiApiKey(); - const context: AgentBlockContext = { openAiToken, - mcpServers: [], + mcpServers: getProjectMcpServers(cell.notebook), notebookContext, addMarkdownBlock: async ({ content }: { content: string }) => { - await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'markdown', content); + try { + await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'markdown', content); - return MARKDOWN_BLOCK_ADDED_TEXT; + return MARKDOWN_BLOCK_ADDED_TEXT; + } catch (error) { + const insertError = error instanceof Error ? error : new Error(String(error)); + + return `Failed to add markdown block: ${insertError.message}`; + } }, addAndExecuteCodeBlock: async ({ code }: { code: string }) => { try { - const cellIndex = await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'code', code); - const insertedCell = cell.notebook.cellAt(cellIndex); + const insertedCell = await insertEphemeralCell( + cell.notebook, + cell.index, + agentBlock.id, + 'code', + code + ); - const { success, outputs } = await executeEphemeralCell(insertedCell, execution.token); - const outputText = describeExecutionOutputs(outputs); + const { success, outputs, error } = await executeEphemeralCell(insertedCell, execution.token); + const outputText = error ?? describeExecutionOutputs(outputs); return success ? `Output:\n${outputText}` : `Execution failed:\n${outputText}`; } catch (error) { @@ -177,37 +239,36 @@ export async function executeAgentCell( } }, onAgentEvent: async (event: AgentStreamEvent) => { - logger.info('Agent event', JSON.stringify(event)); - if (lastAgentEventType != null && lastAgentEventType !== event.type) { - accumulated += `\n\n`; - } + logger.trace(`Agent event: ${event.type}`); + + let delta = lastAgentEventType != null && lastAgentEventType !== event.type ? `\n\n` : ''; + switch (event.type) { case 'tool_called': - // Ignore calling tool_called events - accumulated += `[Agent] Tool called: ${event.toolName}`; + delta += `[Agent] Tool called: ${event.toolName}`; break; case 'tool_output': - accumulated += `[Agent] Tool output: ${event.toolName}\n`; - accumulated += `[Agent] Tool output length: ${event.output?.length}`; + delta += `[Agent] Tool output: ${event.toolName}\n`; + delta += `[Agent] Tool output length: ${event.output?.length}`; break; case 'text_delta': if (lastAgentEventType !== 'text_delta') { - accumulated += `[Agent] Text:\n`; + delta += `[Agent] Text:\n`; } - accumulated += event.text; + delta += event.text; break; case 'reasoning_delta': if (lastAgentEventType !== 'reasoning_delta') { - accumulated += `[Agent] Reasoning:\n`; + delta += `[Agent] Reasoning:\n`; } - accumulated += event.text; + delta += event.text; break; default: event satisfies never; } lastAgentEventType = event.type; - await execution.replaceOutputItems(NotebookCellOutputItem.text(accumulated), output); + await execution.appendOutputItems(NotebookCellOutputItem.stdout(delta), output); } }; @@ -219,9 +280,10 @@ export async function executeAgentCell( execution.end(true, Date.now()); } catch (error) { + // `logger.error(msg, error)` only renders `Error.prototype.toString()` unless the error is + // branded with `isJupyterError`, so the stack has to be logged explicitly. 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); } @@ -256,13 +318,20 @@ function getInsertIndexAfterAgentCell( return index; } +/** + * Inserts an ephemeral cell after the agent cell and returns the cell that was actually created. + * + * Resolving by block id rather than by index matters: `cellAt` clamps out-of-range indices instead of + * throwing, so a rejected edit or a concurrent structural change would otherwise hand the caller a + * pre-existing user cell — which `addAndExecuteCodeBlock` would then run. + */ async function insertEphemeralCell( notebook: NotebookDocument, agentCellIndex: number, agentBlockId: string, blockType: 'code' | 'markdown', content: string -): Promise { +): Promise { const insertIndex = getInsertIndexAfterAgentCell(notebook, agentCellIndex, agentBlockId); const block: DeepnoteBlock = { @@ -282,17 +351,42 @@ async function insertEphemeralCell( const edit = new WorkspaceEdit(); edit.set(notebook.uri, [NotebookEdit.insertCells(insertIndex, [cellData])]); - await workspace.applyEdit(edit); - return insertIndex; + if (!(await workspace.applyEdit(edit))) { + throw new Error(`Failed to insert ephemeral ${blockType} cell for agent block ${agentBlockId}`); + } + + // The converter mirrors the block id into `__deepnoteBlockId` precisely because VS Code may + // rewrite `id`, so match on that. + const insertedCell = notebook.getCells().find((c) => c.metadata?.__deepnoteBlockId === block.id); + + if (!insertedCell) { + throw new Error(`Inserted ephemeral ${blockType} cell ${block.id} not found in notebook`); + } + + return insertedCell; } const EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS = 5 * 60 * 1000; +export interface EphemeralCellExecutionResult { + success: boolean; + outputs: unknown[]; + executionCount: number | null; + /** Why the run failed, when the failure wasn't the cell's own output (cancellation, timeout). */ + error?: string; +} + export async function executeEphemeralCell( cell: NotebookCell, token?: CancellationToken -): Promise<{ success: boolean; outputs: unknown[]; executionCount: number | null }> { +): Promise { + // Bail before dispatching: rejecting the deferred alone would abandon the wait but still hand the + // generated code to the kernel. + if (token?.isCancellationRequested) { + throw new CancellationError(); + } + const completionDeferred = createDeferred(); const disposables: IDisposable[] = []; @@ -305,11 +399,7 @@ export async function executeEphemeralCell( ); if (token) { - if (token.isCancellationRequested) { - completionDeferred.reject(new CancellationError()); - } else { - disposables.push(token.onCancellationRequested(() => completionDeferred.reject(new CancellationError()))); - } + disposables.push(token.onCancellationRequested(() => completionDeferred.reject(new CancellationError()))); } const timeout = setTimeout(() => { @@ -332,10 +422,17 @@ export async function executeEphemeralCell( executionCount: cell.executionSummary?.executionOrder ?? null }; } catch (error) { + if (error instanceof CancellationError) { + throw error; + } + + // Report the reason rather than collapsing everything into "(no output)" — a timed-out cell + // is still running, and telling the agent it produced nothing invites an immediate retry. return { success: false, outputs: [], - executionCount: null + executionCount: null, + error: error instanceof Error ? error.message : String(error) }; } finally { dispose(disposables); diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index b124742244..ac377cadb3 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -1,16 +1,21 @@ import { expect } from 'chai'; import * as sinon from 'sinon'; -import { anything, capture, instance, mock, reset, when } from 'ts-mockito'; +import { anything, capture, instance, mock, reset, verify, when } from 'ts-mockito'; import { + CancellationError, CancellationTokenSource, Disposable, EventEmitter, ExtensionMode, + NotebookCell, + NotebookCellData, NotebookCellOutput, NotebookCellOutputItem, NotebookController, SecretStorage, - SecretStorageChangeEvent + SecretStorageChangeEvent, + Uri, + WorkspaceEdit } from 'vscode'; import type { AgentBlock } from '@deepnote/blocks'; @@ -22,8 +27,36 @@ import { NotebookCellExecutionState, notebookCellExecutions } from '../../platfo 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'; +import { + describeExecutionOutputs, + executeAgentCell, + executeEphemeralCell, + isAgentCell +} from './agentCellExecutionHandler'; +import { createMockCell, createMockNotebook } from './deepnoteTestHelpers'; + +/** + * Wires up a ServiceContainer whose IExtensionContext exposes an in-memory SecretStorage, so the + * secret-store helpers take their real code paths instead of the ExtensionMode.Test no-op branch. + */ +function stubSecretStorage(secretStorage: Map): void { + 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(); + }); +} suite('AgentCellExecutionHandler', () => { const secretStorage = new Map(); @@ -61,48 +94,54 @@ suite('AgentCellExecutionHandler', () => { }); }); - suite('getOpenAiApiKey', () => { - 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())); - }); + suite('describeExecutionOutputs', () => { + test('joins nbformat line arrays in stream text', () => { + const output = { + output_type: 'stream', + name: 'stdout', + text: ['hello\n', 'world\n'] + }; - teardown(() => { - disposables = dispose(disposables); + expect(describeExecutionOutputs([output])).to.equal('hello\nworld\n'); }); - test('returns key when configured', async () => { - secretStorage.set('openAiApiKey', 'test-key'); + // translateCellDisplayOutput splits `text/plain` into a line array, and @deepnote/blocks + // stringifies it with String(...) — which joins with commas. Without the fix the agent reads + // its own DataFrame output with a comma glued to the start of every line but the first. + test('joins nbformat line arrays in execute_result text/plain', () => { + const output = { + output_type: 'execute_result', + data: { 'text/plain': [' a b\n', '0 1 4\n', '1 2 5'] }, + metadata: {}, + execution_count: 1 + }; - const key = await getOpenAiApiKey(); + expect(describeExecutionOutputs([output])).to.equal(' a b\n0 1 4\n1 2 5'); + }); - expect(key).to.equal('test-key'); + test('joins nbformat line arrays in display_data text/plain', () => { + const output = { + output_type: 'display_data', + data: { 'text/plain': ['line one\n', 'line two'] }, + metadata: {} + }; + + expect(describeExecutionOutputs([output])).to.equal('line one\nline two'); }); - test('throws when key is not set', async () => { - when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + test('leaves single-line text/plain untouched', () => { + const output = { + output_type: 'execute_result', + data: { 'text/plain': ['42'] }, + metadata: {}, + execution_count: 1 + }; - try { - await getOpenAiApiKey(); - expect.fail('Should have thrown'); - } catch (e) { - expect((e as Error).message).to.include('OpenAI API key is not set'); - } + expect(describeExecutionOutputs([output])).to.equal('42'); + }); + + test('reports no output for an empty output list', () => { + expect(describeExecutionOutputs([])).to.equal('(no output)'); }); }); @@ -112,7 +151,7 @@ suite('AgentCellExecutionHandler', () => { clearOutput: sinon.SinonStub; end: sinon.SinonStub; replaceOutput: sinon.SinonStub; - replaceOutputItems: sinon.SinonStub; + appendOutputItems: sinon.SinonStub; start: sinon.SinonStub; }; let mockController: NotebookController; @@ -121,21 +160,7 @@ suite('AgentCellExecutionHandler', () => { setup(() => { 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(); - }); + stubSecretStorage(secretStorage); disposables.push(new Disposable(() => sinon.restore())); mockExecution = { @@ -143,7 +168,7 @@ suite('AgentCellExecutionHandler', () => { clearOutput: sinon.stub().resolves(), end: sinon.stub(), replaceOutput: sinon.stub().resolves(), - replaceOutputItems: sinon.stub().resolves(), + appendOutputItems: sinon.stub().resolves(), start: sinon.stub() }; @@ -156,6 +181,10 @@ suite('AgentCellExecutionHandler', () => { teardown(() => { disposables = dispose(disposables); + reset(mockedVSCodeNamespaces.commands); + // Restore the default from vscode-mock rather than reset()ing the whole workspace + // namespace, which other suites rely on. + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(true)); }); function createAgentCell(text: string = 'Test prompt') { @@ -165,6 +194,55 @@ suite('AgentCellExecutionHandler', () => { }); } + /** + * Builds an agent cell inside a notebook whose cell list the test can mutate, and applies + * insert/delete notebook edits to that list so the handler observes its own mutations. + * + * The mocked `WorkspaceEdit.set` discards the edits it is given, so record them off the + * prototype rather than reading them back off the edit object. + */ + function createAgentCellInMutableNotebook(cells: NotebookCell[] = [], agentBlockId = 'agent-block-1') { + const notebook = createMockNotebook({ cells }); + const agentCell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' }, id: agentBlockId }, + text: 'Test prompt' + }); + + (agentCell as { notebook: typeof notebook }).notebook = notebook; + (agentCell as { index: number }).index = 0; + cells.unshift(agentCell); + + type RecordedEdit = { range: { start: number; end: number }; newCells: NotebookCellData[] }; + let recordedEdits: RecordedEdit[] = []; + + sinon.stub(WorkspaceEdit.prototype, 'set').callsFake((_uri, edits) => { + recordedEdits = edits as unknown as RecordedEdit[]; + }); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { + for (const notebookEdit of recordedEdits) { + const { start, end } = notebookEdit.range; + const inserted = notebookEdit.newCells.map((cellData) => { + const created = createMockCell({ + text: cellData.value, + metadata: cellData.metadata + }); + (created as { notebook: typeof notebook }).notebook = notebook; + + return created; + }); + + cells.splice(start, end - start, ...inserted); + } + cells.forEach((cell, index) => ((cell as { index: number }).index = index)); + recordedEdits = []; + + return Promise.resolve(true); + }); + + return { agentCell, cells, notebook }; + } + test('creates execution and starts it', async () => { const cell = createAgentCell('Analyze data'); @@ -198,7 +276,7 @@ suite('AgentCellExecutionHandler', () => { expect(text).to.include('[Agent] Planning next steps...'); }); - test('streams events via replaceOutputItems using onAgentEvent callback', async () => { + test('streams events via appendOutputItems 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' }); @@ -210,10 +288,15 @@ suite('AgentCellExecutionHandler', () => { await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - expect(mockExecution.replaceOutputItems.callCount).to.equal(2); + expect(mockExecution.appendOutputItems.callCount).to.equal(2); + + const item = mockExecution.appendOutputItems.firstCall.args[0] as NotebookCellOutputItem; + expect(item.mime).to.equal('application/vnd.code.notebook.stdout'); }); - test('streaming chunks accumulate text progressively', async () => { + // Each event must ship only its own delta: re-sending the whole transcript per token is + // O(n²) bytes over the extension-host boundary, and runtime-core awaits this callback. + test('streaming sends only the incremental text per event', async () => { executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { await context.onAgentEvent?.({ type: 'text_delta', text: 'first' }); await context.onAgentEvent?.({ type: 'text_delta', text: ' second' }); @@ -226,17 +309,13 @@ suite('AgentCellExecutionHandler', () => { await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); const getChunkText = (callIndex: number): string => { - const item = mockExecution.replaceOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; + const item = mockExecution.appendOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; return Buffer.from(item.data).toString('utf-8'); }; - const chunk1 = getChunkText(0); - const chunk2 = getChunkText(1); - - expect(chunk1).to.include('[Agent] Text:'); - expect(chunk1).to.include('first'); - expect(chunk2).to.include('first second'); + expect(getChunkText(0)).to.equal('[Agent] Text:\nfirst'); + expect(getChunkText(1)).to.equal(' second'); }); test('separates different event types with blank lines', async () => { @@ -252,7 +331,7 @@ suite('AgentCellExecutionHandler', () => { await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); const getChunkText = (callIndex: number): string => { - const item = mockExecution.replaceOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; + const item = mockExecution.appendOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; return Buffer.from(item.data).toString('utf-8'); }; @@ -330,6 +409,101 @@ suite('AgentCellExecutionHandler', () => { const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); expect(text).to.include('OpenAI API key is not set'); }); + + // The key prompt is the last fallible step before the run starts, so it has to come before + // the cleanup that throws away the previous run's generated cells. + test('keeps previous ephemeral cells when the API key prompt is cancelled', async () => { + secretStorage.clear(); + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const previousResult = createMockCell({ + text: 'print("previous run")', + metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }, + index: 1 + }); + const { agentCell, cells } = createAgentCellInMutableNotebook([previousResult]); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.end.firstCall.args[0]).to.be.false; + expect(cells).to.include(previousResult); + }); + + test('inserts a markdown cell after the agent cell with ephemeral metadata', async () => { + const { agentCell, cells } = createAgentCellInMutableNotebook(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.addMarkdownBlock({ content: '## Findings' }); + + return { finalOutput: '' } as AgentBlockResult; + }); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(cells).to.have.lengthOf(2); + expect(cells[1].document.getText()).to.equal('## Findings'); + expect(cells[1].metadata?.is_ephemeral).to.be.true; + expect(cells[1].metadata?.agent_source_block_id).to.equal(agentCell.metadata?.id); + }); + + test('inserts successive cells after the ones it already added', async () => { + const { agentCell, cells } = createAgentCellInMutableNotebook(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.addMarkdownBlock({ content: 'first' }); + await context.addMarkdownBlock({ content: 'second' }); + + return { finalOutput: '' } as AgentBlockResult; + }); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(cells.map((cell) => cell.document.getText())).to.deep.equal(['Test prompt', 'first', 'second']); + }); + + // cellAt clamps rather than throwing, so resolving the inserted cell by index would hand the + // agent a pre-existing user cell and execute it. + test('fails the tool call without executing anything when the insert edit is rejected', async () => { + const { agentCell } = createAgentCellInMutableNotebook(); + let toolResult: string | undefined; + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(false)); + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + toolResult = await context.addAndExecuteCodeBlock({ code: 'print(1)' }); + + return { finalOutput: '' } as AgentBlockResult; + }); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(toolResult).to.include('Execution error'); + expect(toolResult).to.include('Failed to insert ephemeral code cell'); + verify(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).never(); + }); + + test('removes only the ephemeral cells belonging to this agent', async () => { + const ownResult = createMockCell({ + text: 'own', + metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }, + index: 1 + }); + const otherAgentResult = createMockCell({ + text: 'other agent', + metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-2' }, + index: 2 + }); + const userCell = createMockCell({ text: 'user code', metadata: {}, index: 3 }); + + const { agentCell, cells } = createAgentCellInMutableNotebook([ownResult, otherAgentResult, userCell]); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(cells).to.not.include(ownResult); + expect(cells).to.include(otherAgentResult); + expect(cells).to.include(userCell); + }); }); suite('executeEphemeralCell', () => { @@ -363,7 +537,9 @@ suite('AgentCellExecutionHandler', () => { }); }); - test('returns success false immediately when token is pre-cancelled', async () => { + // Rejecting the deferred alone abandons only the wait — the generated code would still reach + // the kernel after the user cancelled. + test('throws without dispatching to the kernel when the token is pre-cancelled', async () => { const cell = createMockCell({ index: 0 }); const tokenSource = new CancellationTokenSource(); tokenSource.cancel(); @@ -371,16 +547,43 @@ suite('AgentCellExecutionHandler', () => { when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); try { - const result = await executeEphemeralCell(cell, tokenSource.token); - - expect(result).to.deep.equal({ - success: false, - outputs: [], - executionCount: null - }); + await executeEphemeralCell(cell, tokenSource.token); + expect.fail('Should have thrown'); + } catch (e) { + expect(e).to.be.instanceOf(CancellationError); } finally { tokenSource.dispose(); } + + verify(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).never(); }); + + test('reports the failure reason instead of swallowing it', async () => { + const cell = createMockCell({ index: 0 }); + + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenReject( + new Error('kernel is dead') + ); + + const result = await executeEphemeralCell(cell); + + expect(result.success).to.be.false; + expect(result.error).to.equal('kernel is dead'); + }); + }); +}); + +suite('createMockNotebook', () => { + test('reads through to the backing cell array', () => { + const cells: NotebookCell[] = [createMockCell({ text: 'first' })]; + const notebook = createMockNotebook({ cells, uri: Uri.file('/test/mutable.deepnote') }); + + expect(notebook.cellCount).to.equal(1); + + cells.push(createMockCell({ text: 'second', index: 1 })); + + expect(notebook.cellCount).to.equal(2); + expect(notebook.cellAt(1).document.getText()).to.equal('second'); + expect(notebook.getCells()).to.have.lengthOf(2); }); }); diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index a168ea6a1f..6ccd26fe6a 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -14,22 +14,22 @@ 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'; +import { isAgentCell } from './dataConversionUtils'; 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).max(MAX_ITERATIONS); +/** The key `agentBlockSchema` defines and `executeAgentBlock` reads. */ +const AGENT_MODEL_METADATA_KEY = 'deepnote_agent_model'; + +/** The schema default, and the sentinel runtime-core compares against to fall back to its own choice. */ +const AGENT_MODEL_AUTO = 'auto'; + +const AGENT_MODEL_OPTIONS = [AGENT_MODEL_AUTO, 'gpt-4o', 'gpt-5']; /** - * Provides status bar items for agent cells showing the block type indicator, - * AI model picker, and max iterations setting. + * Provides status bar items for agent cells showing the block type indicator + * and the AI model picker. */ @injectable() export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService { @@ -58,15 +58,6 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv }) ); - this.disposables.push( - commands.registerCommand('deepnote.setAgentMaxIterations', async (cell?: NotebookCell) => { - const activeCell = cell || this.getActiveCell(); - if (activeCell) { - await this.setMaxIterations(activeCell); - } - }) - ); - this.disposables.push( commands.registerCommand('deepnote.setOpenAiApiKey', async () => { const key = await promptForOpenAiApiKey(); @@ -100,19 +91,14 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv return undefined; } - if (!this.isAgentCell(cell)) { + if (!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) - ]; + return [this.createAgentIndicatorItem(), this.createModelPickerItem(cell, model)]; } private createAgentIndicatorItem(): NotebookCellStatusBarItem { @@ -124,20 +110,6 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv }; } - 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)}`, @@ -161,78 +133,17 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv return undefined; } - private getMaxIterations(metadata: Record | undefined): number { - const value = metadata?.deepnote_max_iterations; - // z.coerce.number() turns true into 1, which then satisfies the range check, so booleans - // would be accepted as an iteration count instead of falling back to the default. - const result = typeof value === 'boolean' ? undefined : 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; - } - private getModel(metadata: Record | undefined): string { - const value = metadata?.deepnote_model; + const value = metadata?.[AGENT_MODEL_METADATA_KEY]; 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 }); + return AGENT_MODEL_AUTO; } private async switchModel(cell: NotebookCell): Promise { - if (!this.isAgentCell(cell)) { + if (!isAgentCell(cell)) { return; } @@ -252,9 +163,10 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv return; } - const newModel = selected.label === 'auto' ? undefined : selected.label; - - await this.updateCellMetadata(cell, { deepnote_model: newModel }); + // Write 'auto' rather than deleting the key: `convertCellToBlock` doesn't re-run the zod + // schema, so a missing key reaches runtime-core as `undefined` — which fails its + // `!== "auto"` check and gets passed to `openai()` as the model name. + await this.updateCellMetadata(cell, { [AGENT_MODEL_METADATA_KEY]: selected.label }); } private async updateCellMetadata(cell: NotebookCell, updates: Record): Promise { diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts index 62adc562cc..c8463c4a34 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -26,7 +26,7 @@ suite('AgentCellStatusBarProvider', () => { const items = provider.provideCellStatusBarItems(cell, mockToken); expect(items).to.not.be.undefined; - expect(items).to.have.lengthOf(3); + expect(items).to.have.lengthOf(2); }); test('Should return undefined for code cell', () => { @@ -108,7 +108,7 @@ suite('AgentCellStatusBarProvider', () => { const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' }, - deepnote_model: 'gpt-4o' + deepnote_agent_model: 'gpt-4o' } }); const items = provider.provideCellStatusBarItems(cell, mockToken)!; @@ -116,23 +116,23 @@ suite('AgentCellStatusBarProvider', () => { expect(items[1].text).to.include('Model: gpt-4o'); }); - test('Should display sonnet model', () => { + test('Should display gpt-5 model', () => { const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' }, - deepnote_model: 'sonnet' + deepnote_agent_model: 'gpt-5' } }); const items = provider.provideCellStatusBarItems(cell, mockToken)!; - expect(items[1].text).to.include('Model: sonnet'); + expect(items[1].text).to.include('Model: gpt-5'); }); test('Should display "auto" when model is empty string', () => { const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' }, - deepnote_model: '' + deepnote_agent_model: '' } }); const items = provider.provideCellStatusBarItems(cell, mockToken)!; @@ -157,155 +157,20 @@ suite('AgentCellStatusBarProvider', () => { }); }); - 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 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)!; - - 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', () => { + test('Should ignore metadata keys the runtime does not consume', () => { const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' }, - deepnote_model: 'gpt-4o', + deepnote_agent_model: 'gpt-4o', deepnote_max_iterations: 50 } }); const items = provider.provideCellStatusBarItems(cell, mockToken)!; - expect(items).to.have.lengthOf(3); + expect(items).to.have.lengthOf(2); 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/dataConversionUtils.ts b/src/notebooks/deepnote/dataConversionUtils.ts index 8b01da1256..fc9d227ccd 100644 --- a/src/notebooks/deepnote/dataConversionUtils.ts +++ b/src/notebooks/deepnote/dataConversionUtils.ts @@ -4,6 +4,8 @@ import { NotebookCell, NotebookCellData } from 'vscode'; +import type { Pocket } from '../../platform/deepnote/pocket'; + export function parseJsonWithFallback(value: string, fallback?: unknown): unknown | null { try { return JSON.parse(value); @@ -24,6 +26,18 @@ export function generateBlockId(): string { return id; } +/** + * Returns true if the cell is backed by an agent block. + * + * Lives here rather than next to the execution handler so callers that only need the predicate + * don't pull `@deepnote/runtime-core` into their module graph. + */ +export function isAgentCell(cell: NotebookCell): boolean { + const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; + + return pocket?.type === 'agent'; +} + /** * Returns true if the cell metadata indicates an ephemeral cell (auto-generated by agent). */ diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 45e1e6bcae..f63fc28855 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -1100,25 +1100,26 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } cells` ); - const agentCells = cells.filter((cell) => isAgentCell(cell)); - const kernelCells = cells.filter((cell) => !isAgentCell(cell)); + // Agent blocks can run arbitrary commands declared in the project file (MCP servers), so + // gate this path the same way VSCodeNotebookController gates its own execute handler. + if (!workspace.isTrusted) { + logger.info(`Workspace is not trusted, skipping execution for ${getDisplayPath(doc.uri)}`); - // 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` - ); + return; + } - for (const cell of agentCells) { + const kernelCells = cells.filter((cell) => !isAgentCell(cell)); + + if (kernelCells.length === 0) { + // Nothing needs the kernel, so don't make the user configure an environment first. + for (const cell of cells) { try { await executeAgentCell(cell, controller); } catch (cellError) { logger.error(`Error executing agent cell ${cell.index}`, cellError); } } - } - if (kernelCells.length === 0) { return; } @@ -1150,7 +1151,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return; } - logger.info(`Executing ${kernelCells.length} cells through kernel after environment configuration`); + logger.info(`Executing ${cells.length} cells after environment configuration`); // Get or create a kernel for this notebook with the new connection const kernel = this.kernelProvider.getOrCreate(doc, { @@ -1162,15 +1163,21 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Execute cells through the kernel const kernelExecution = this.kernelProvider.getKernelExecution(kernel); - for (const cell of kernelCells) { + // Document order matters: an agent executes the code it generates immediately, so it + // must not overtake the cells above it that set up the state it reads. + for (const cell of cells) { try { - await kernelExecution.executeCell(cell); + if (isAgentCell(cell)) { + await executeAgentCell(cell, controller); + } else { + await kernelExecution.executeCell(cell); + } } catch (cellError) { logger.error(`Error executing cell ${cell.index}`, cellError); } } - logger.info(`Finished executing ${kernelCells.length} cells`); + logger.info(`Finished executing ${cells.length} cells`); } catch (error) { if (isCancellationError(error)) { logger.info(`Environment setup cancelled for ${getDisplayPath(doc.uri)}`); diff --git a/src/notebooks/deepnote/deepnoteTestHelpers.ts b/src/notebooks/deepnote/deepnoteTestHelpers.ts index 5e1cfeac27..e00762286e 100644 --- a/src/notebooks/deepnote/deepnoteTestHelpers.ts +++ b/src/notebooks/deepnote/deepnoteTestHelpers.ts @@ -47,6 +47,11 @@ export interface CreateMockNotebookOptions { notebookType?: string; uri?: Uri; metadata?: Record; + /** + * Backing cells. Pass the same array you mutate in the test — `cellAt`/`getCells`/`cellCount` + * read through to it, so edits applied during the test are visible to the code under test. + */ + cells?: NotebookCell[]; } /** @@ -56,15 +61,23 @@ export interface CreateMockNotebookOptions { * @returns A mock NotebookDocument */ export function createMockNotebook(options?: CreateMockNotebookOptions): NotebookDocument { - const { notebookType = 'deepnote', uri = Uri.file('/test/notebook.deepnote'), metadata = {} } = options ?? {}; + const { + notebookType = 'deepnote', + uri = Uri.file('/test/notebook.deepnote'), + metadata = {}, + cells = [] + } = options ?? {}; return { uri, notebookType, metadata, - cellCount: 0, - cellAt: () => ({}) as NotebookCell, - getCells: () => [], + get cellCount() { + return cells.length; + }, + // Mirrors VS Code: the index is clamped to the notebook rather than throwing. + cellAt: (index: number) => cells[Math.min(Math.max(index, 0), cells.length - 1)] ?? ({} as NotebookCell), + getCells: () => cells, version: 1, isDirty: false, isUntitled: false,