From ffe09588ac299d7ad88f852d707ee91eb5041181 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 9 Mar 2026 08:30:26 +0000 Subject: [PATCH 01/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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/37] 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, From 33f95f4c50cdbbc31ac5a4de1b06f8690de06c3f Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 3 Aug 2026 14:10:56 +0000 Subject: [PATCH 25/37] Update code comment --- src/notebooks/deepnote/agentCellExecutionHandler.ts | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index d6d2706357..f38a4a8503 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -110,8 +110,8 @@ function joinMultilineString(value: unknown): unknown { * 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. + * `print()` output isn't dropped and a `df.head()` string representation doesn't reach the agent with a + * comma glued to the start of every line. */ function normalizeOutputsForTextExtraction(outputs: unknown[]): unknown[] { return outputs.map((output) => { From 20273415780626461b3d8d554a64cb28d345239c Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 3 Aug 2026 16:22:38 +0000 Subject: [PATCH 26/37] fix(agent-block): correct controller, cleanup, timeout and integration gaps Four defects from a review pass over this branch, each with a regression test written against the unfixed code first. - The placeholder controller's execute handler passed its own captured controller to executeAgentCell after environment setup had already disposed and deselected it, so the agent cell was skipped with nothing but a log line. Use the real controller that owns the notebook by then. - executeEphemeralCell awaited the dispatch before the completion deferred, so the timeout could not end a run whose command never resolved, and a rejection arriving in between was reported as unhandled. Wait on both together. - A rejected ephemeral cleanup edit was only logged, leaving the previous run's cells in the notebook context sent to the model and in Run All's kernel batch. Fail the run instead. - Project integrations were never passed to executeAgentBlock, so runtime-core dropped the integration IDs and dntk.execute_sql instructions from its system prompt entirely. Also drop the isAgentCell re-export from the handler. Both consumers imported executeAgentCell alongside it, so it bought nothing and only made it easy to pull @deepnote/runtime-core into a module graph that has no need for it. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- .../controllers/vscodeNotebookController.ts | 3 +- .../deepnote/agentCellExecutionHandler.ts | 68 +++++++++------ .../agentCellExecutionHandler.unit.test.ts | 84 +++++++++++++++++-- .../deepnoteKernelAutoSelector.node.ts | 7 +- ...epnoteKernelAutoSelector.node.unit.test.ts | 80 +++++++++++++++++- 5 files changed, 203 insertions(+), 39 deletions(-) diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index ae64367887..b472167ded 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -91,7 +91,8 @@ 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'; +import { executeAgentCell } from '../deepnote/agentCellExecutionHandler'; +import { isAgentCell } from '../deepnote/dataConversionUtils'; /** * Our implementation of the VSCode Notebook Controller. Called by VS code to execute cells in a notebook. Also displayed diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index f38a4a8503..c79742919c 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -30,38 +30,48 @@ import { ServiceContainer } from '../../platform/ioc/container'; import { logger } from '../../platform/logging'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; import { IDeepnoteNotebookManager } from '../types'; -import { generateBlockId, generateSortingKey, isAgentCell, isEphemeralCell } from './dataConversionUtils'; +import { generateBlockId, generateSortingKey, 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. + * Project-level MCP servers and database integrations declared in the `.deepnote` file, matching what + * the CLI's ExecutionEngine passes. `executeAgentBlock` merges the servers with any block-level + * `deepnote_mcp_servers` (block wins on name), and only names the integrations — along with the + * `dntk.execute_sql` instructions — in its system prompt when that list is non-empty, so leaving + * either 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. + * Spawning MCP servers 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'] { +function getProjectAgentContext(notebook: NotebookDocument): Pick { const projectId = notebook.metadata?.deepnoteProjectId as string | undefined; const notebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; if (!projectId || !notebookId) { - return []; + return { mcpServers: [] }; } const manager = ServiceContainer.instance.tryGet(IDeepnoteNotebookManager); - const servers = manager?.getProjectForNotebook(projectId, notebookId)?.project.settings?.mcpServers ?? []; + const project = manager?.getProjectForNotebook(projectId, notebookId)?.project; + const mcpServers = project?.settings?.mcpServers ?? []; + const integrations = project?.integrations ?? []; + + if (mcpServers.length > 0) { + logger.info( + `Agent cell: using ${mcpServers.length} project MCP server(s): ${mcpServers.map((s) => s.name).join(', ')}` + ); + } - if (servers.length > 0) { + if (integrations.length > 0) { logger.info( - `Agent cell: using ${servers.length} project MCP server(s): ${servers.map((s) => s.name).join(', ')}` + `Agent cell: using ${integrations.length} project integration(s): ${integrations + .map((i) => i.name) + .join(', ')}` ); } - return servers; + return { mcpServers, integrations }; } // Tool results reported back to the agent. These mirror the wording @deepnote/runtime-core uses in @@ -205,7 +215,7 @@ export async function executeAgentCell( const context: AgentBlockContext = { openAiToken, - mcpServers: getProjectMcpServers(cell.notebook), + ...getProjectAgentContext(cell.notebook), notebookContext, addMarkdownBlock: async ({ content }: { content: string }) => { try { @@ -409,12 +419,16 @@ export async function executeEphemeralCell( try { const cellIndex = cell.index; - await commands.executeCommand('notebook.cell.execute', { - ranges: [{ start: cellIndex, end: cellIndex + 1 }], - document: cell.notebook.uri - }); - - await completionDeferred.promise; + // The dispatch settles independently of the cell reaching Idle, so both waits have to start + // together — otherwise the timeout cannot end a run whose command never resolves, and a + // rejection arriving before the second await is reported as unhandled. + await Promise.all([ + commands.executeCommand('notebook.cell.execute', { + ranges: [{ start: cellIndex, end: cellIndex + 1 }], + document: cell.notebook.uri + }), + completionDeferred.promise + ]); return { success: cell.executionSummary?.success === true, @@ -458,10 +472,12 @@ async function removeEphemeralCellsForAgent(notebook: NotebookDocument, agentBlo 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}`); + // Fatal rather than a warning: the notebook context the agent receives is read off the live + // document, and Run All keeps these cells out of its kernel batch only by their index going + // negative once they are deleted. + if (!(await workspace.applyEdit(edit))) { + throw new Error(`Failed to remove ephemeral cells for agent block ${agentBlockId}`); } + + logger.info(`Removed ${deletions.length} ephemeral cell(s) for agent block ${agentBlockId}`); } diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index ac377cadb3..910ecd7d47 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -27,19 +27,16 @@ 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 { - describeExecutionOutputs, - executeAgentCell, - executeEphemeralCell, - isAgentCell -} from './agentCellExecutionHandler'; -import { createMockCell, createMockNotebook } from './deepnoteTestHelpers'; +import { describeExecutionOutputs, executeAgentCell, executeEphemeralCell } from './agentCellExecutionHandler'; +import { isAgentCell } from './dataConversionUtils'; +import { IDeepnoteNotebookManager } from '../types'; +import { createDeepnoteFile, createDeepnoteProject, 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 { +function stubSecretStorage(secretStorage: Map): ServiceContainer { const context = mock(); const secrets = mock(); const onDidChangeSecrets = new EventEmitter(); @@ -56,6 +53,8 @@ function stubSecretStorage(secretStorage: Map): void { return Promise.resolve(); }); + + return serviceContainer; } suite('AgentCellExecutionHandler', () => { @@ -156,11 +155,12 @@ suite('AgentCellExecutionHandler', () => { }; let mockController: NotebookController; let executeAgentBlockStub: sinon.SinonStub; + let mockServiceContainer: ServiceContainer; setup(() => { secretStorage.clear(); secretStorage.set('openAiApiKey', 'test-key'); - stubSecretStorage(secretStorage); + mockServiceContainer = stubSecretStorage(secretStorage); disposables.push(new Disposable(() => sinon.restore())); mockExecution = { @@ -504,6 +504,49 @@ suite('AgentCellExecutionHandler', () => { expect(cells).to.include(otherAgentResult); expect(cells).to.include(userCell); }); + + // The notebook context is read off the live document, and Run All keeps stale ephemeral cells + // out of the kernel batch only by their index going negative on deletion. + test('fails the run without calling the agent when the cleanup edit is rejected', async () => { + const previousResult = createMockCell({ + text: 'print("previous run")', + metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }, + index: 1 + }); + const { agentCell } = createAgentCellInMutableNotebook([previousResult]); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(false)); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(executeAgentBlockStub.called).to.be.false; + expect(mockExecution.end.firstCall.args[0]).to.be.false; + }); + + test('passes project MCP servers and integrations to the agent', async () => { + const integrations = [{ id: 'warehouse', name: 'Warehouse', type: 'postgres' }]; + const mcpServers = [{ name: 'files', command: 'mcp-files', args: [] }]; + const notebookManager = mock(); + + when(mockServiceContainer.tryGet(IDeepnoteNotebookManager)).thenReturn( + instance(notebookManager) + ); + when(notebookManager.getProjectForNotebook('project-1', 'notebook-1')).thenReturn( + createDeepnoteFile({ project: createDeepnoteProject({ integrations, settings: { mcpServers } }) }) + ); + + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test prompt', + notebookMetadata: { deepnoteProjectId: 'project-1', deepnoteNotebookId: 'notebook-1' } + }); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + const context = executeAgentBlockStub.firstCall.args[1] as AgentBlockContext; + expect(context.mcpServers).to.deep.equal(mcpServers); + expect(context.integrations).to.deep.equal(integrations); + }); }); suite('executeEphemeralCell', () => { @@ -570,6 +613,29 @@ suite('AgentCellExecutionHandler', () => { expect(result.success).to.be.false; expect(result.error).to.equal('kernel is dead'); }); + + // The dispatch settles independently of the cell reaching Idle, so waiting on it first would + // leave the timeout unable to end a run whose command never resolves. + test('times out while the dispatch is still pending', async () => { + const cell = createMockCell({ index: 0 }); + const clock = sinon.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }); + + try { + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenCall( + () => new Promise(() => undefined) + ); + + const resultPromise = executeEphemeralCell(cell); + await clock.tickAsync(5 * 60 * 1000); + + const result = await resultPromise; + + expect(result.success).to.be.false; + expect(result.error).to.equal('Ephemeral cell execution timed out'); + } finally { + clock.restore(); + } + }); }); }); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index f63fc28855..aed1e61bba 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -57,7 +57,8 @@ import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { IControllerRegistration, IVSCodeNotebookController } from '../controllers/types'; import { IDeepnoteNotebookManager } from '../types'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; -import { executeAgentCell, isAgentCell } from './agentCellExecutionHandler'; +import { executeAgentCell } from './agentCellExecutionHandler'; +import { isAgentCell } from './dataConversionUtils'; import { computeRequirementsHash } from './deepnoteProjectUtils'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; @@ -1168,7 +1169,9 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, for (const cell of cells) { try { if (isAgentCell(cell)) { - await executeAgentCell(cell, controller); + // Configuring the environment disposed this placeholder and handed the + // notebook to the real controller, which now owns its executions. + await executeAgentCell(cell, realController.controller); } else { await kernelExecution.executeCell(cell); } diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts index d5f71bc77c..2f6c4cae23 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -20,7 +20,7 @@ import { IConfigurationService } from '../../platform/common/types'; import { IDeepnoteNotebookManager } from '../types'; import { IKernelProvider, IKernel, IJupyterKernelSpec } from '../../kernels/types'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; -import { NotebookDocument, Uri, NotebookController, CancellationToken } from 'vscode'; +import { EventEmitter, NotebookCell, NotebookDocument, Uri, NotebookController, CancellationToken } from 'vscode'; import { DeepnoteEnvironment } from '../../kernels/deepnote/environments/deepnoteEnvironment'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; @@ -1034,6 +1034,84 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { }); }); }); + + suite('Placeholder controller execution', () => { + function createExecutionStub() { + return { + start: sandbox.stub(), + end: sandbox.stub(), + clearOutput: sandbox.stub().resolves(), + replaceOutput: sandbox.stub().resolves(), + appendOutput: sandbox.stub().resolves(), + appendOutputItems: sandbox.stub().resolves() + }; + } + + // Configuring the environment disposes and deselects the placeholder, and VS Code throws for + // executions created on either a disposed or an unassociated controller. + test('runs agent cells on the real controller after configuring the environment', async () => { + const placeholderExecution = createExecutionStub(); + const placeholder = { + supportsExecutionOrder: false, + supportedLanguages: [] as string[], + updateNotebookAffinity: sandbox.stub(), + dispose: sandbox.stub(), + createNotebookCellExecution: sandbox.stub().returns(placeholderExecution) + } as unknown as NotebookController; + + when( + mockedVSCodeNamespaces.notebooks!.createNotebookController(anything(), anything(), anything()) + ).thenReturn(placeholder); + when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(true); + + const onDidCloseNotebookDocument = new EventEmitter(); + when(mockedVSCodeNamespaces.workspace.onDidCloseNotebookDocument).thenReturn( + onDidCloseNotebookDocument.event + ); + + const realExecution = createExecutionStub(); + const realNotebookController = { + createNotebookCellExecution: sandbox.stub().returns(realExecution) + } as unknown as NotebookController; + const realController = mock(); + when(realController.controller).thenReturn(realNotebookController); + + const internals = selector as unknown as { + createPlaceholderController(notebook: NotebookDocument): NotebookController; + notebookControllers: Map; + }; + + internals.createPlaceholderController(mockNotebook); + internals.notebookControllers.set(getNotebookKey(mockNotebook.uri), instance(realController)); + + sandbox.stub(selector, 'ensureEnvironmentConfiguredBeforeExecution').resolves(true); + + const kernelExecution = { executeCell: sandbox.stub().resolves() }; + when(mockKernelProvider.getOrCreate(anything(), anything())).thenReturn(instance(mock())); + when(mockKernelProvider.getKernelExecution(anything())).thenReturn( + kernelExecution as unknown as ReturnType + ); + + const agentCell = { + index: 0, + metadata: { __deepnotePocket: { type: 'agent' } } + } as unknown as NotebookCell; + const codeCell = { index: 1, metadata: {} } as unknown as NotebookCell; + + await placeholder.executeHandler!([agentCell, codeCell], mockNotebook, placeholder); + + assert.isTrue( + (realNotebookController.createNotebookCellExecution as sinon.SinonStub).calledOnceWithExactly( + agentCell + ), + 'agent cell execution should be created on the real controller' + ); + assert.isTrue( + (placeholder.createNotebookCellExecution as sinon.SinonStub).notCalled, + 'no execution should be created on the disposed placeholder' + ); + }); + }); }); /** From 4b55d475c376aca63fd339c843288fb9eed1f5e7 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 4 Aug 2026 05:42:16 +0000 Subject: [PATCH 27/37] test(agent-block): move isAgentCell tests alongside their source isAgentCell lives in dataConversionUtils, not the execution handler; its tests only sat in the handler's suite because the handler used to re-export it. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- .../agentCellExecutionHandler.unit.test.ts | 33 ---------------- .../deepnote/dataConversionUtils.unit.test.ts | 38 +++++++++++++++++++ 2 files changed, 38 insertions(+), 33 deletions(-) create mode 100644 src/notebooks/deepnote/dataConversionUtils.unit.test.ts diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index 910ecd7d47..8d2a1290c0 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -28,7 +28,6 @@ import { dispose } from '../../platform/common/utils/lifecycle'; import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; import { ServiceContainer } from '../../platform/ioc/container'; import { describeExecutionOutputs, executeAgentCell, executeEphemeralCell } from './agentCellExecutionHandler'; -import { isAgentCell } from './dataConversionUtils'; import { IDeepnoteNotebookManager } from '../types'; import { createDeepnoteFile, createDeepnoteProject, createMockCell, createMockNotebook } from './deepnoteTestHelpers'; @@ -61,38 +60,6 @@ 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' } } }); - - 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('describeExecutionOutputs', () => { test('joins nbformat line arrays in stream text', () => { const output = { diff --git a/src/notebooks/deepnote/dataConversionUtils.unit.test.ts b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts new file mode 100644 index 0000000000..ab2f94efc5 --- /dev/null +++ b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts @@ -0,0 +1,38 @@ +import { expect } from 'chai'; + +import { isAgentCell } from './dataConversionUtils'; +import { createMockCell } from './deepnoteTestHelpers'; + +suite('DataConversionUtils', () => { + 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; + }); + }); +}); From c1c9db295011cad605c49f46be962875b319780f Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 4 Aug 2026 13:39:47 +0000 Subject: [PATCH 28/37] test(agent-block): add an E2E test for the agent tool loop MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Runs an agent block end to end against a stand-in OpenAI API, covering the path no test reached before: the agent generating Python via add_code_block, the extension executing it on a real kernel, and the kernel's output going back to the agent. The mock is @copilotkit/aimock, fetched with a pinned npx rather than added to package.json — it declares jest and vitest as peers, and resolving those against this tree forces overrides that would outlive the test. The scripted legs match on toolResultContains rather than a request counter, so the agent can only advance if the extension really ran the generated code and fed the real output back; under --strict a broken round-trip fails loudly instead of taking a different path. Being pure request-shape predicates, they also replay correctly on a Mocha retry, which does not re-run `before`. OPENAI_BASE_URL is set at the spec's module scope: ExTester launches VS Code from a root beforeAll and the extension host inherits its environment at spawn time, so a hook is too late, while Mocha loads spec files before running any hook. Without it runtime-core falls back to the real api.openai.com, so startMockOpenAiServer refuses to run when it is unset. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- .github/workflows/e2e.yml | 7 + package.json | 3 +- test/e2e/fixtures/agent-block.deepnote | 22 ++ test/e2e/helpers/index.ts | 1 + test/e2e/helpers/mockOpenAiServer.ts | 265 ++++++++++++++++++++++++ test/e2e/helpers/notebook.ts | 85 ++++---- test/e2e/suite/agentBlock.e2e.test.ts | 267 +++++++++++++++++++++++++ 7 files changed, 609 insertions(+), 41 deletions(-) create mode 100644 test/e2e/fixtures/agent-block.deepnote create mode 100644 test/e2e/helpers/mockOpenAiServer.ts create mode 100644 test/e2e/suite/agentBlock.e2e.test.ts diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index a28c014d37..3c5b674fc8 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -71,6 +71,13 @@ jobs: - name: Install the Python extension into the test instance run: npm run setup:e2e:deps + - name: Pre-download the mock LLM server + # The agent-block suite starts this mid-run via npx. setup-node's cache is keyed on the + # lockfile, which aimock is deliberately absent from, so ~/.npm/_npx is never restored and the + # fetch would otherwise happen inside the test — where a registry blip surfaces as an opaque + # start-up timeout. Failing here instead points straight at the cause. + run: npm run setup:e2e:mock + - name: Cache pip wheel downloads # Provisioning the Deepnote environment pip-installs the toolkit dependency tree into a # fresh venv on first kernel connect — the bulk of the E2E runtime. Caching pip's wheel diff --git a/package.json b/package.json index dcbc0e6a77..8b89c5523d 100644 --- a/package.json +++ b/package.json @@ -2670,7 +2670,8 @@ "compile-e2e-watch": "tsc -p ./test/e2e/tsconfig.json --watch", "setup:e2e:vscode": "extest get-vscode -c max && extest get-chromedriver -c max", "setup:e2e:deps": "extest install-from-marketplace ms-python.python -e .test-extensions", - "setup:e2e": "npm run setup:e2e:vscode && npm run setup:e2e:deps", + "setup:e2e:mock": "npx -y -p @copilotkit/aimock@1.37.4 llmock --help", + "setup:e2e": "npm run setup:e2e:vscode && npm run setup:e2e:deps && npm run setup:e2e:mock", "test:e2e": "extest setup-and-run \"./out/e2e/suite/*.e2e.test.js\" -c max -o ./test/e2e/settings.json -e .test-extensions -m ./test/e2e/.mocharc.js -i", "test:e2e:prebuilt": "extest run-tests \"./out/e2e/suite/*.e2e.test.js\" -c max -o ./test/e2e/settings.json -e .test-extensions -m ./test/e2e/.mocharc.js", "test:unittests": "mocha --config ./build/.mocha.unittests.js.json ./out/**/*.unit.test.js", diff --git a/test/e2e/fixtures/agent-block.deepnote b/test/e2e/fixtures/agent-block.deepnote new file mode 100644 index 0000000000..334562ffec --- /dev/null +++ b/test/e2e/fixtures/agent-block.deepnote @@ -0,0 +1,22 @@ +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00.000Z' + modifiedAt: '2025-01-01T00:00:00.000Z' +project: + id: e2e-agent-block-project + name: E2E Agent Block + notebooks: + - id: e2e-agent-block-notebook + name: Agent Block + blocks: + - id: e2e-agent-block + blockGroup: e2e-agent-group + type: agent + content: |- + Run some Python, then add a markdown block summarising this notebook. + sortingKey: a0 + metadata: + deepnote_agent_model: gpt-5 + executionMode: block + isModule: false + settings: {} diff --git a/test/e2e/helpers/index.ts b/test/e2e/helpers/index.ts index 3ed57bb00d..45a939edc4 100644 --- a/test/e2e/helpers/index.ts +++ b/test/e2e/helpers/index.ts @@ -4,6 +4,7 @@ export * from './constants'; export * from './deepnoteEnvironment'; export * from './deepnoteTree'; export * from './fixtures'; +export * from './mockOpenAiServer'; export * from './modals'; export * from './notebook'; export * from './notifications'; diff --git a/test/e2e/helpers/mockOpenAiServer.ts b/test/e2e/helpers/mockOpenAiServer.ts new file mode 100644 index 0000000000..8d7833ffaf --- /dev/null +++ b/test/e2e/helpers/mockOpenAiServer.ts @@ -0,0 +1,265 @@ +import { spawn } from 'child_process'; +import * as fs from 'fs'; +import { connect } from 'net'; +import * as os from 'os'; +import * as path from 'path'; +import { setTimeout as delay } from 'timers/promises'; + +// Fetched by npx rather than installed: aimock declares `jest` and `vitest` as peers, and resolving +// those against this repo's tree forces overrides that would outlive the test. npx resolves in its +// own cache, so the dependency graph here is untouched. +// +// Pinned exactly — a range would let the mock the suite asserts against change underneath it. +const AIMOCK_VERSION = '1.37.4'; +// `llmock` is the bin that takes `-f`/`-p`; the package's `aimock` bin takes `--config` instead. +const AIMOCK_BIN = 'llmock'; + +// Deliberately below `ip_local_port_range` (32768-60999 here and on GitHub runners): inside it an +// unrelated outbound connection can hold the number as its source port, which the pre-flight check +// below would not see (a client socket does not accept) and `listen` would then fail with EADDRINUSE. +const MOCK_OPENAI_PORT = 18_937; + +/** + * Points the extension host at the mock server instead of the real OpenAI API. + * + * MUST be called at a spec file's module scope, never from `before`. ExTester launches VS Code from a + * root `beforeAll` (`vscode-extension-tester/out/suite/runner.js`) and the extension host inherits its + * environment at spawn time, so a hook runs too late — while Mocha loads spec files before it runs any + * hook, which is what makes module scope early enough. + * + * `rootHooks.ts` would be the tidier home, but ExTester builds Mocha through `new Mocha(config)`, and + * the programmatic API ignores the `require` option that file is wired up with. + */ +export function pointExtensionHostAtMockServer(): void { + process.env.OPENAI_BASE_URL = `http://127.0.0.1:${MOCK_OPENAI_PORT}/v1`; +} + +// `npm run setup:e2e:mock` primes `~/.npm/_npx` with this exact spec, so a warm start resolves from +// cache without a registry round-trip. The ceiling still covers a cold fetch: the setup step is not +// enforced, and if the two specs ever drift the run silently falls back to downloading here. +const START_TIMEOUT = 90_000; +const POLL_INTERVAL = 200; + +// How long to wait for the tree to go down after each of SIGTERM and SIGKILL. Short because the +// graceful signal is not what we rely on: `server.close()` releases the listening socket before the +// process is gone, so a freed port arrives long before the shutdown it appears to signal. +const STOP_TIMEOUT = 2_000; + +export interface MockOpenAiServer { + /** Stops the server and removes its fixtures. Idempotent; safe to call more than once. */ + stop: () => Promise; +} + +export interface MockToolCall { + arguments: string; + id: string; + name: string; +} + +/** + * Which request a scripted leg answers. Both alternatives are predicates over the request's own + * messages, with no server-side counter — unlike aimock's `sequenceIndex`, which would run past the + * end of the script on a Mocha retry (`.mocharc.js` sets `retries: 1` and `before` does not re-run + * between attempts) and fail the retry for a different reason than the original. + * + * `toolResultContains` additionally requires the last message to be a tool result, so it is what + * proves a round-trip: the leg is only reachable if the extension really ran the previous tool and + * fed its real output back. + */ +export type MockAgentMatch = { hasToolResult: false } | { toolResultContains: string }; + +/** What the agent gets back: another tool call, or the final text that ends the loop. */ +export type MockAgentResponse = { content: string } | { toolCall: MockToolCall }; + +export interface MockAgentTurn { + match: MockAgentMatch; + response: MockAgentResponse; +} + +function canConnect(port: number): Promise { + return new Promise((resolve) => { + const socket = connect({ host: '127.0.0.1', port }); + const settle = (reachable: boolean) => { + socket.destroy(); + resolve(reachable); + }; + + socket.once('connect', () => settle(true)); + socket.once('error', () => settle(false)); + }); +} + +/** + * Fails the run when `OPENAI_BASE_URL` does not point at this server. + * + * Silence here is not a failed test: `executeAgentBlock` reads the variable at call time and falls + * back to `openai(model)` against the real api.openai.com (@deepnote/runtime-core dist/index.js:102), + * so an unset or drifted value sends the suite's prompts — and whatever key is in SecretStorage — to + * the live API. + */ +function assertBaseUrlPointsAtMock(): void { + if (!process.env.OPENAI_BASE_URL?.includes(`:${MOCK_OPENAI_PORT}`)) { + throw new Error( + `OPENAI_BASE_URL must point at 127.0.0.1:${MOCK_OPENAI_PORT}; run via "npm run test:e2e". ` + + `Without it the agent would call the real OpenAI API. Current value: ` + + `${JSON.stringify(process.env.OPENAI_BASE_URL)}` + ); + } +} + +/** Writes `turns` as an aimock fixtures file in a fresh temp directory, and returns that directory. */ +function writeFixtures(turns: MockAgentTurn[]): string { + const fixtures = turns.map(({ match, response }) => ({ + match, + response: 'toolCall' in response ? { toolCalls: [response.toolCall] } : { content: response.content } + })); + + const directory = fs.mkdtempSync(path.join(os.tmpdir(), 'deepnote-e2e-aimock-')); + // The `fixtures` wrapper is required — aimock's loader rejects a bare array. + fs.writeFileSync(path.join(directory, 'fixtures.json'), JSON.stringify({ fixtures }, undefined, 4)); + + return directory; +} + +/** + * Starts aimock on the mock port, scripted with `turns` — each answering the request its `match` + * describes. Resolves once the port accepts connections, which the CLI only reaches after loading and + * validating the fixtures, so a served request can never race an unloaded fixture. + * + * Runs with `--strict`, so a request matching no leg is answered with an error rather than a default: + * a broken round-trip fails loudly instead of quietly taking a different path through the script. + */ +export async function startMockOpenAiServer(turns: MockAgentTurn[]): Promise { + assertBaseUrlPointsAtMock(); + + // Without this the readiness poll below cannot tell our server from someone else's: a leftover + // from a crashed run would satisfy it instantly, and the suite would then be asserting against + // that server's fixtures while ours had already died of EADDRINUSE. + if (await canConnect(MOCK_OPENAI_PORT)) { + throw new Error( + `Port ${MOCK_OPENAI_PORT} is already in use — most likely a mock server left behind by an ` + + `interrupted run. Kill it before running the suite.` + ); + } + + const fixturesDirectory = writeFixtures(turns); + + const child = spawn( + 'npx', + [ + // Use the cache `setup:e2e:mock` primed rather than re-checking the registry mid-run, but + // still fall back to fetching so a skipped setup step degrades to slow instead of broken. + '--prefer-offline', + '-y', + '-p', + `@copilotkit/aimock@${AIMOCK_VERSION}`, + AIMOCK_BIN, + '-f', + fixturesDirectory, + '-p', + String(MOCK_OPENAI_PORT), + // Answers an unmatched request with an error instead of letting the agent keep asking + // until runtime-core's 10-turn cap, which would bury the cause. + '--strict', + '--log-level', + 'warn' + ], + { + // npx runs the server two levels down (`npm exec` -> `sh -c` -> node), and a signal sent + // to npx alone leaves that grandchild holding the port. Its own process group makes the + // whole tree signalable; see `signalTree`. + detached: true, + stdio: ['ignore', 'inherit', 'inherit'] + } + ); + + let exitReason: string | undefined; + child.once('exit', (code, signal) => { + exitReason = `code ${code}, signal ${signal}`; + }); + // Node emits 'error' rather than 'exit' when the spawn itself fails (ENOENT for a missing npx, + // EACCES, …). An unhandled 'error' on a ChildProcess throws out of the event loop and takes the + // whole mocha process with it, losing every later suite and skipping ExTester's teardown; routing + // it through exitReason turns that into the readiness loop's ordinary failure. + child.once('error', (error) => { + exitReason = `spawn failed: ${error.message}`; + }); + + const signalTree = (signal: NodeJS.Signals) => { + try { + if (child.pid === undefined) { + return; + } + + process.kill(-child.pid, signal); + } catch { + // Already gone — nothing left to signal. + } + }; + + // A crashed runner would otherwise leave the port bound and fail every later run. + const killChild = () => signalTree('SIGKILL'); + process.once('exit', killChild); + + const hasShutDown = async () => + (child.exitCode !== null || child.signalCode !== null) && !(await canConnect(MOCK_OPENAI_PORT)); + + const waitForShutdown = async (): Promise => { + const deadline = Date.now() + STOP_TIMEOUT; + + while (Date.now() < deadline) { + if (await hasShutDown()) { + return true; + } + + await delay(POLL_INTERVAL); + } + + return false; + }; + + const stop = async () => { + fs.rmSync(fixturesDirectory, { force: true, recursive: true }); + + // Neither half of `hasShutDown` proves the node server is gone on its own: `server.close()` + // frees the listening socket while still draining open connections, and npx exits ahead of + // the server it spawned. So give SIGTERM a graceful window, then SIGKILL the group + // unconditionally — on an already-dead group that is a swallowed ESRCH, and it is the only + // step that guarantees nothing is left behind holding the port. + signalTree('SIGTERM'); + await waitForShutdown(); + signalTree('SIGKILL'); + + if (!(await waitForShutdown())) { + throw new Error( + `aimock did not shut down after SIGKILL: port ${MOCK_OPENAI_PORT} still accepts ` + + `connections, or the npx process has not exited.` + ); + } + + // Only once the shutdown is confirmed — until here the exit-time kill is the last safety net. + process.removeListener('exit', killChild); + }; + + const deadline = Date.now() + START_TIMEOUT; + while (Date.now() < deadline) { + if (exitReason) { + await stop(); + + throw new Error(`aimock exited before it started listening (${exitReason}); see its output above`); + } + + if (await canConnect(MOCK_OPENAI_PORT)) { + return { stop }; + } + + await delay(POLL_INTERVAL); + } + + await stop(); + + throw new Error( + `aimock did not listen on port ${MOCK_OPENAI_PORT} within ${START_TIMEOUT}ms. A stale server from an ` + + `earlier run may still hold the port.` + ); +} diff --git a/test/e2e/helpers/notebook.ts b/test/e2e/helpers/notebook.ts index 6fc561a7b3..b00bd6bf53 100644 --- a/test/e2e/helpers/notebook.ts +++ b/test/e2e/helpers/notebook.ts @@ -42,27 +42,57 @@ export async function clickRunAll(notebookFileName: string): Promise { } /** - * Reads the notebook cell output once. - * - * Output lives two iframes deep (iframe.webview.ready -> #active-frame). We only attempt to switch - * when an output webview iframe actually exists (`getViewToSwitchTo`), and we read output-specific - * elements inside the frame — so we never match the cell's source code that is visible in the editor - * of the main document. Returns '' when no output is present yet. + * Runs `read` inside the notebook webview (iframe.webview.ready -> #active-frame) and switches back + * afterwards. `read` only ever sees the webview, never the cell source in the main document — the + * guarantee callers rely on to avoid matching a cell's own text. Returns '' when the frame is absent, + * went stale, or has painted nothing yet, so callers can poll. */ -export async function readRenderedOutput(): Promise { +async function readInsideNotebookWebview(read: (webView: WebView) => Promise): Promise { + const driver = VSBrowser.instance.driver; const webView = new WebView(); - const outputFrame = await webView.getViewToSwitchTo().catch((error) => { - console.warn('[deepnote-e2e] locate notebook output webview:', error); + const frame = await webView.getViewToSwitchTo().catch((error) => { + console.warn('[deepnote-e2e] locate notebook webview:', error); return undefined; }); - if (!outputFrame) { + if (!frame) { return ''; } - let text = ''; try { await webView.switchToFrame(OUTPUT_FRAME_SWITCH_TIMEOUT); + + // switchToFrame re-resolves the view and returns silently when it has gone, leaving the + // driver on the workbench document — where a body read would scrape the editor, and the cell + // source with it. Confirm we actually descended before letting `read` run. + if (await driver.executeScript('return window.self === window.top')) { + return ''; + } + + return (await read(webView)).trim(); + } catch (error) { + console.warn('[deepnote-e2e] read inside notebook webview:', error); + + return ''; + } finally { + await webView.switchBack().catch((error) => { + console.warn('[deepnote-e2e] switch back from notebook webview:', error); + }); + } +} + +/** + * Reads everything the notebook webview currently paints — rendered markdown cells as well as cell + * outputs. Use it when a rendered markdown cell is part of the assertion, since `readRenderedOutput` + * deliberately narrows to output-only elements. + */ +export async function readNotebookWebviewText(): Promise { + return readInsideNotebookWebview(async (webView) => (await webView.findWebElement(By.css('body'))).getText()); +} + +/** Reads the notebook cell output once, falling back to the whole frame if the renderer used unexpected classes. */ +export async function readRenderedOutput(): Promise { + return readInsideNotebookWebview(async (webView) => { const elements = await webView.findWebElements(By.css(OUTPUT_SELECTOR)); const texts = await Promise.all( elements.map((element) => @@ -73,36 +103,11 @@ export async function readRenderedOutput(): Promise { }) ) ); - text = texts.join('\n').trim(); - - // Fallback: if the renderer used unexpected classes, read the frame body — safe here because - // we have confirmed we are inside the output iframe, not the editor. - if (!text) { - const body = await webView.findWebElement(By.css('body')).catch((error) => { - console.warn('[deepnote-e2e] read output frame body:', error); - - return undefined; - }); - text = body - ? ( - await body.getText().catch((error) => { - console.warn('[deepnote-e2e] read output frame body text:', error); - - return ''; - }) - ).trim() - : ''; - } - } catch (error) { - // Frame went stale or output not painted yet — treat as no output this tick. - console.warn('[deepnote-e2e] read rendered notebook output:', error); - } finally { - await webView.switchBack().catch((error) => { - console.warn('[deepnote-e2e] switch back from notebook output webview:', error); - }); - } + const text = texts.join('\n').trim(); - return text; + // Safe as a fallback because we have already confirmed we are inside the webview, not the editor. + return text || (await webView.findWebElement(By.css('body'))).getText(); + }); } /** diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts new file mode 100644 index 0000000000..ee1b28db7d --- /dev/null +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -0,0 +1,267 @@ +/** + * E2E (ExTester): one agent block driving a three-leg tool loop against a stand-in OpenAI API, so no + * network call is made. The agent asks for a code block, the extension inserts it as an ephemeral + * cell and runs it on the kernel, and the real stdout goes back as the tool result; the agent then + * asks for a markdown block and finally answers. + * + * The scripted legs 2 and 3 match on `toolResultContains`, so the agent can only advance if the + * extension genuinely executed the generated Python and returned its actual output. With aimock's + * `--strict`, a broken round-trip matches no leg and fails loudly. + * + * Executing generated code needs a real kernel: the first run provisions a venv and installs the + * Deepnote toolkit, which takes minutes. + */ + +import { expect } from 'chai'; +import { EditorView, InputBox, VSBrowser, WebView, Workbench } from 'vscode-extension-tester'; + +import { + FIRST_RUN_OUTPUT_TIMEOUT, + MockOpenAiServer, + OUTPUT_POLL_INTERVAL, + QUICK_PICK_TIMEOUT, + SUITE_TIMEOUT, + WORKBENCH_TIMEOUT, + clickRunAll, + confirmModalDialog, + copyFixtureToTempDir, + createEnvironment, + createScreenshotter, + dismissAllNotifications, + openFolderViaDialog, + openWorkspaceFile, + pointExtensionHostAtMockServer, + readNotebookWebviewText, + selectEnvironmentForNotebook, + startMockOpenAiServer, + waitForNotification +} from '../helpers'; + +// At module scope on purpose — VS Code is already running by the time `before` executes, and it +// inherits this at spawn time. See the function's contract. +pointExtensionHostAtMockServer(); + +const AGENT_FILE = 'agent-block.deepnote'; +const CODE_TOOL_NAME = 'add_code_block'; +const MARKDOWN_TOOL_NAME = 'add_markdown_block'; + +// The extension's tool result for a successful add_markdown_block (agentCellExecutionHandler.ts); +// leg 3 keys off it, so the wording is a coupling to that constant. +const MARKDOWN_BLOCK_ADDED_TEXT = 'Markdown block added.'; + +// A stable name: createEnvironment treats "already exists" as success, so a leftover environment from +// a previous or retried run is reused rather than colliding — and its provisioned venv with it. +const ENVIRONMENT_NAME = 'E2E Agent Env'; + +// Once the kernel is up the agent itself talks only to the local mock, so it is bounded by UI and +// extension-host latency. The kernel's own first run is bounded by FIRST_RUN_OUTPUT_TIMEOUT instead. +const AGENT_RUN_TIMEOUT = 60_000; + +// Printed by the Python the agent asks for. Only the executed ephemeral code cell can put it in the +// webview: the webview renders outputs and markdown previews, never cell source, and the agent's own +// transcript reports tool output by length rather than by content. +const PYTHON_OUTPUT_MARKER = 'agent-generated-python-ran'; +const GENERATED_PYTHON = `print("${PYTHON_OUTPUT_MARKER}")`; + +// Reaches the notebook only through the agent's tool call — the streamed transcript never echoes +// tool arguments — so seeing it rendered is what proves an ephemeral markdown cell was inserted. +const EPHEMERAL_MARKDOWN_TEXT = 'Ephemeral markdown written by the E2E agent run'; +const FINAL_AGENT_TEXT = 'Summary added as a markdown block.'; + +// The mock server ignores credentials, but the extension refuses to start an agent run without a +// stored key (and would otherwise block on an input box mid-execution). +const MOCK_API_KEY = 'sk-e2e-mock-key'; +// Exact palette label matters: `Workbench.executeCommand` silently runs the first palette entry on a +// mismatch. +const SET_API_KEY_COMMAND = 'Deepnote: Set OpenAI API Key'; +const CLEAR_API_KEY_COMMAND = 'Deepnote: Clear OpenAI API Key'; +const REVERT_FILE_COMMAND = 'File: Revert File'; +// VS Code's save prompt; the bundle stores it with a mnemonic marker ("Do&&n't Save") that is +// stripped before rendering, so the button's text is this. +const DISCARD_CHANGES_BUTTON = "Don't Save"; +const API_KEY_SAVED_NOTIFICATION = /OpenAI API key has been saved/; + +/** Polls the notebook webview until every marker is present, returning whatever it last read. */ +async function awaitWebviewMarkers(markers: string[], timeout: number): Promise { + const driver = VSBrowser.instance.driver; + const deadline = Date.now() + timeout; + let text = ''; + + while (Date.now() < deadline) { + text = await readNotebookWebviewText(); + if (markers.every((marker) => text.includes(marker))) { + return text; + } + + await driver.sleep(OUTPUT_POLL_INTERVAL); + } + + return text; +} + +/** Stores the throwaway key in SecretStorage so the agent run never opens the key prompt. */ +async function storeMockOpenAiApiKey(): Promise { + await new Workbench().executeCommand(SET_API_KEY_COMMAND); + + const input = await InputBox.create(QUICK_PICK_TIMEOUT); + await input.setText(MOCK_API_KEY); + await input.confirm(); + + // Confirm the key really landed. Had the palette missed the command, InputBox.create would have + // bound to the still-open palette and typed the key into it, and the suite would run keyless — + // surfacing a full AGENT_RUN_TIMEOUT later as a generic missing-marker failure. + await waitForNotification(API_KEY_SAVED_NOTIFICATION, QUICK_PICK_TIMEOUT, true); +} + +describe('Deepnote — running an agent block against a stand-in OpenAI API', function () { + this.timeout(SUITE_TIMEOUT); + + let cleanupTempDir: (() => void) | undefined; + let mockServer: MockOpenAiServer | undefined; + let screenshot: (label: string) => Promise; + + before(async function () { + screenshot = createScreenshotter(this); + + const copy = copyFixtureToTempDir(AGENT_FILE); + cleanupTempDir = copy.cleanup; + + await VSBrowser.instance.waitForWorkbench(WORKBENCH_TIMEOUT); + await openFolderViaDialog(copy.tempDir); + await VSBrowser.instance.waitForWorkbench(WORKBENCH_TIMEOUT); + + await openWorkspaceFile(AGENT_FILE); + await VSBrowser.instance.driver.wait( + async () => (await new EditorView().getOpenEditorTitles()).some((title) => title.includes(AGENT_FILE)), + WORKBENCH_TIMEOUT, + `${AGENT_FILE} did not open` + ); + + // Binds a real kernel for the code the agent generates, replacing the "Select Environment" + // placeholder controller the auto-selector picks on open. This is also the settle signal it waits + // on: selectEnvironmentForNotebook returns after the post-binding "switched successfully" + // toast, so Run All is not racing the auto-selection. + await createEnvironment(ENVIRONMENT_NAME); + await selectEnvironmentForNotebook(ENVIRONMENT_NAME, AGENT_FILE); + + // Toasts steal focus from the command palette. Safe to do after the environment flow, which + // has already driven extension commands and so guarantees `onNotebook:deepnote` activation. + await dismissAllNotifications(); + await storeMockOpenAiApiKey(); + await screenshot('kernel-connected'); + }); + + after(async function () { + // Process and filesystem cleanup goes FIRST. The UI steps below can each stall for minutes — + // this is the one suite that ends with a dirty notebook, so `closeAllEditors` retries against + // a save modal and the `.catch()`-swallowed revert cannot stop it — and a wedged UI step + // would otherwise burn SUITE_TIMEOUT with the server and temp dirs never released. + await mockServer?.stop().catch((error) => { + console.warn('[agent-block] stop the mock OpenAI server during cleanup:', error); + }); + try { + cleanupTempDir?.(); + } catch (error) { + console.warn('[agent-block] remove temp workspace dir during cleanup:', error); + } + + await new WebView().switchBack().catch((error) => { + console.warn('[agent-block] switch back from webview during cleanup:', error); + }); + // The inserted ephemeral cell leaves the notebook dirty, and the resulting modal save prompt + // outlives this suite and blocks the next one in the shared VS Code instance. + await new Workbench().executeCommand(REVERT_FILE_COMMAND).catch((error) => { + console.warn('[agent-block] revert notebook during cleanup:', error); + }); + await new EditorView().closeAllEditors().catch((error) => { + console.warn('[agent-block] close all editors during cleanup:', error); + }); + + // Backstop for a revert that did not land: an unanswered save modal blocks the next suite. + // Gated on an editor surviving the close, because confirmModalDialog polls for the full + // WORKBENCH_TIMEOUT when no dialog is up — dead time on every green run otherwise. + const openEditors = await new EditorView().getOpenEditorTitles().catch(() => [] as string[]); + if (openEditors.length > 0) { + await confirmModalDialog(DISCARD_CHANGES_BUTTON).catch((error) => { + console.warn('[agent-block] discard unsaved changes during cleanup:', error); + }); + } + // SecretStorage outlives this suite in the shared VS Code instance, so leave no key behind. + await new Workbench().executeCommand(CLEAR_API_KEY_COMMAND).catch((error) => { + console.warn('[agent-block] clear the stored OpenAI API key during cleanup:', error); + }); + }); + + // Known limitation of the Mocha retry (`.mocharc.js` sets `retries: 1`): if this times out with an + // execution still in flight, the retry's clickRunAll may find Interrupt where Run All was and fail + // for an unrelated reason, losing the original signal. Only reachable on an already-failing test, + // so it costs debuggability rather than correctness. + it('executes the code block the agent generates, then inserts its markdown block', async function () { + // Leg 2 and leg 3 are reachable only via what the extension sends back, so the script itself + // asserts the round-trip: leg 2 needs the kernel's real stdout, leg 3 the markdown tool's reply. + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: GENERATED_PYTHON }), + id: 'call_e2e_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: PYTHON_OUTPUT_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: EPHEMERAL_MARKDOWN_TEXT }), + id: 'call_e2e_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: FINAL_AGENT_TEXT } + } + ]); + + // Clears the "OpenAI API key has been saved." and "switched successfully" toasts, which would + // otherwise intercept the toolbar click. + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); + + // Every marker is asserted below, so poll for all of them — a missing one then fails on its + // own assertion rather than on whichever runs first. Split in two waits because the stages + // have very different budgets: the generated cell is the first thing to touch the kernel, and + // that first execution carries the connect cost, while the rest is local. Waiting on the + // Python marker first also reports a kernel failure as a kernel failure rather than as a + // missing agent marker. + await awaitWebviewMarkers([PYTHON_OUTPUT_MARKER], FIRST_RUN_OUTPUT_TIMEOUT); + + const agentMarkers = [ + `[Agent] Tool called: ${CODE_TOOL_NAME}`, + `[Agent] Tool called: ${MARKDOWN_TOOL_NAME}`, + EPHEMERAL_MARKDOWN_TEXT, + FINAL_AGENT_TEXT + ]; + const webviewText = await awaitWebviewMarkers(agentMarkers, AGENT_RUN_TIMEOUT); + + await screenshot('agent-run'); + + expect(webviewText, 'the agent cell did not stream its add_code_block call into the cell output').to.contain( + `[Agent] Tool called: ${CODE_TOOL_NAME}` + ); + expect(webviewText, 'the generated code cell did not run on the kernel').to.contain(PYTHON_OUTPUT_MARKER); + expect( + webviewText, + 'the agent cell did not stream its add_markdown_block call into the cell output' + ).to.contain(`[Agent] Tool called: ${MARKDOWN_TOOL_NAME}`); + expect(webviewText, 'the tool call did not insert an ephemeral markdown cell').to.contain( + EPHEMERAL_MARKDOWN_TEXT + ); + expect(webviewText, "the agent's final message was not streamed into the cell output").to.contain( + FINAL_AGENT_TEXT + ); + }); +}); From 4c897f58b7f4d7885319aae87a87ece966a6d553 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 4 Aug 2026 15:54:30 +0000 Subject: [PATCH 29/37] test(agent-block): scope the mock server to the attempt, not the suite MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mocha retries the test but not `before`/`after`, so a server started once outlived a failed attempt and still held the port when the retry began — where the pre-flight check rejected it as a leftover, failing the retry for a different reason than the original and losing the real signal. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- test/e2e/suite/agentBlock.e2e.test.ts | 82 +++++++++++++++------------ 1 file changed, 47 insertions(+), 35 deletions(-) diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts index ee1b28db7d..fd2c30a45d 100644 --- a/test/e2e/suite/agentBlock.e2e.test.ts +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -151,14 +151,55 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu await screenshot('kernel-connected'); }); - after(async function () { - // Process and filesystem cleanup goes FIRST. The UI steps below can each stall for minutes — - // this is the one suite that ends with a dirty notebook, so `closeAllEditors` retries against - // a save modal and the `.catch()`-swallowed revert cannot stop it — and a wedged UI step - // would otherwise burn SUITE_TIMEOUT with the server and temp dirs never released. + // Per attempt, not per suite: `.mocharc.js` sets `retries: 1` and `before`/`after` do not run + // between attempts, so a server started once would still hold the port when the retry starts — + // and `startMockOpenAiServer`'s pre-flight check would reject it as a leftover, failing the retry + // for a different reason than the original and losing the real signal. + // + // Leg 2 and leg 3 are reachable only via what the extension sends back, so the script itself + // asserts the round-trip: leg 2 needs the kernel's real stdout, leg 3 the markdown tool's reply. + beforeEach(async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: GENERATED_PYTHON }), + id: 'call_e2e_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: PYTHON_OUTPUT_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: EPHEMERAL_MARKDOWN_TEXT }), + id: 'call_e2e_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: FINAL_AGENT_TEXT } + } + ]); + }); + + afterEach(async function () { await mockServer?.stop().catch((error) => { - console.warn('[agent-block] stop the mock OpenAI server during cleanup:', error); + console.warn('[agent-block] stop the mock OpenAI server:', error); }); + // Cleared so a failed start cannot leave the next attempt stopping a dead handle. + mockServer = undefined; + }); + + after(async function () { + // Filesystem cleanup goes FIRST. The UI steps below can each stall for minutes — this is the + // one suite that ends with a dirty notebook, so `closeAllEditors` retries against a save modal + // and the `.catch()`-swallowed revert cannot stop it — and a wedged UI step would otherwise + // burn SUITE_TIMEOUT with the temp dir never released. try { cleanupTempDir?.(); } catch (error) { @@ -197,35 +238,6 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu // for an unrelated reason, losing the original signal. Only reachable on an already-failing test, // so it costs debuggability rather than correctness. it('executes the code block the agent generates, then inserts its markdown block', async function () { - // Leg 2 and leg 3 are reachable only via what the extension sends back, so the script itself - // asserts the round-trip: leg 2 needs the kernel's real stdout, leg 3 the markdown tool's reply. - mockServer = await startMockOpenAiServer([ - { - match: { hasToolResult: false }, - response: { - toolCall: { - arguments: JSON.stringify({ code: GENERATED_PYTHON }), - id: 'call_e2e_code', - name: CODE_TOOL_NAME - } - } - }, - { - match: { toolResultContains: PYTHON_OUTPUT_MARKER }, - response: { - toolCall: { - arguments: JSON.stringify({ content: EPHEMERAL_MARKDOWN_TEXT }), - id: 'call_e2e_markdown', - name: MARKDOWN_TOOL_NAME - } - } - }, - { - match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, - response: { content: FINAL_AGENT_TEXT } - } - ]); - // Clears the "OpenAI API key has been saved." and "switched successfully" toasts, which would // otherwise intercept the toolbar click. await dismissAllNotifications(); From 8339146570cfba1160e82cb899034a7f5231bdd9 Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 4 Aug 2026 16:31:30 +0000 Subject: [PATCH 30/37] test(agent-block): script the mock in the test, release it around each attempt MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The scripted legs are what a test is about, so they belong with it rather than in a shared hook — a second test would script different ones. The release runs on both sides of the test, not just after: Mocha retries the test but not before/after, so a server surviving a failed attempt still holds the port when the retry starts, where the pre-flight check rejects it as a leftover and the retry fails for a reason unrelated to the original. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- test/e2e/suite/agentBlock.e2e.test.ts | 87 +++++++++++++++------------ 1 file changed, 48 insertions(+), 39 deletions(-) diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts index fd2c30a45d..a01ee9a637 100644 --- a/test/e2e/suite/agentBlock.e2e.test.ts +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -151,49 +151,27 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu await screenshot('kernel-connected'); }); - // Per attempt, not per suite: `.mocharc.js` sets `retries: 1` and `before`/`after` do not run - // between attempts, so a server started once would still hold the port when the retry starts — - // and `startMockOpenAiServer`'s pre-flight check would reject it as a leftover, failing the retry - // for a different reason than the original and losing the real signal. - // - // Leg 2 and leg 3 are reachable only via what the extension sends back, so the script itself - // asserts the round-trip: leg 2 needs the kernel's real stdout, leg 3 the markdown tool's reply. - beforeEach(async function () { - mockServer = await startMockOpenAiServer([ - { - match: { hasToolResult: false }, - response: { - toolCall: { - arguments: JSON.stringify({ code: GENERATED_PYTHON }), - id: 'call_e2e_code', - name: CODE_TOOL_NAME - } - } - }, - { - match: { toolResultContains: PYTHON_OUTPUT_MARKER }, - response: { - toolCall: { - arguments: JSON.stringify({ content: EPHEMERAL_MARKDOWN_TEXT }), - id: 'call_e2e_markdown', - name: MARKDOWN_TOOL_NAME - } - } - }, - { - match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, - response: { content: FINAL_AGENT_TEXT } - } - ]); - }); - - afterEach(async function () { + /** + * Releases the server the running test started, if any. + * + * Runs on both sides of the test rather than only after it. `.mocharc.js` sets `retries: 1` and + * `before`/`after` do not run between attempts, so a server surviving a failed attempt would still + * hold the port when the retry starts — and `startMockOpenAiServer`'s pre-flight check would then + * reject it as a leftover, failing the retry for a different reason than the original and losing + * the real signal. `afterEach` normally prevents that; `beforeEach` covers the case where it was + * itself interrupted. + */ + async function releaseMockServer(): Promise { await mockServer?.stop().catch((error) => { console.warn('[agent-block] stop the mock OpenAI server:', error); }); - // Cleared so a failed start cannot leave the next attempt stopping a dead handle. + // Cleared so a later release cannot stop an already-dead handle. mockServer = undefined; - }); + } + + beforeEach(releaseMockServer); + + afterEach(releaseMockServer); after(async function () { // Filesystem cleanup goes FIRST. The UI steps below can each stall for minutes — this is the @@ -238,6 +216,37 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu // for an unrelated reason, losing the original signal. Only reachable on an already-failing test, // so it costs debuggability rather than correctness. it('executes the code block the agent generates, then inserts its markdown block', async function () { + // Scripted here because the conversation is what a given test is about — a different test + // scripts different legs. Leg 2 and leg 3 are reachable only via what the extension sends + // back, so the script itself asserts the round-trip: leg 2 needs the kernel's real stdout, + // leg 3 the markdown tool's reply. + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: GENERATED_PYTHON }), + id: 'call_e2e_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: PYTHON_OUTPUT_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: EPHEMERAL_MARKDOWN_TEXT }), + id: 'call_e2e_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: FINAL_AGENT_TEXT } + } + ]); + // Clears the "OpenAI API key has been saved." and "switched successfully" toasts, which would // otherwise intercept the toolbar click. await dismissAllNotifications(); From eb2a868a8ef7c54eb644d185f93aa21b56bf3a3b Mon Sep 17 00:00:00 2001 From: tomas Date: Tue, 4 Aug 2026 16:31:38 +0000 Subject: [PATCH 31/37] fix(e2e): make the mocha root hooks actually run MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `require` is resolved by the mocha CLI's handleRequires, not by the Mocha constructor. ExTester hands this config straight to `new Mocha(config)`, which reads `rootHooks` and ignores `require` — so rootHooks.js was never loaded and the between-test toast dismissal it defines has never run for any suite. Resolving the module here and passing `rootHooks` works under both the programmatic API and the CLI. It also means a missing build fails at config load rather than silently dropping the hooks. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- test/e2e/.mocharc.js | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/test/e2e/.mocharc.js b/test/e2e/.mocharc.js index 14040de246..70a9f99b81 100644 --- a/test/e2e/.mocharc.js +++ b/test/e2e/.mocharc.js @@ -3,12 +3,22 @@ // tests are the real guard rails; this is a generous suite-level safety net. const path = require('path'); +// Loaded here rather than declared via mocha's `require` option, which only the mocha CLI acts on: +// ExTester hands this config straight to `new Mocha(config)` (vscode-extension-tester +// suite/runner.js), and the constructor reads `rootHooks` — already-resolved hook objects — while +// ignoring `require` entirely. Declared the other way the file is never loaded and the hooks below +// silently never run. +// +// Requires compiled output, so compile-e2e must run first; a missing build now fails here rather +// than passing with the hooks quietly absent. +const { mochaHooks } = require(path.resolve(__dirname, '..', '..', 'out', 'e2e', 'rootHooks.js')); + module.exports = { timeout: 1500000, // 25 min — env creation + first kernel start (venv + toolkit) can be slow retries: 1, // absorb transient UI flakiness with a single retry reporter: 'spec', color: true, - // Dismiss notification toasts between tests (rootHooks) so they don't accumulate across the one - // shared VS Code instance. Points at compiled output, so compile-e2e must run first. - require: [path.resolve(__dirname, '..', '..', 'out', 'e2e', 'rootHooks.js')] + // Dismiss notification toasts between tests so they don't accumulate across the one shared + // VS Code instance. + rootHooks: mochaHooks }; From 84c3cedfe77f5c412e5db08d5a8cfedb3eaa586e Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 5 Aug 2026 09:38:13 +0000 Subject: [PATCH 32/37] fix(agent-block): clear the previous run before the batch executes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Run All queues the ephemeral cells a previous agent run generated, and that queue is a snapshot the agent cannot mutate. Keeping them off the kernel relied on their index going negative once the agent deleted them — which never happens when the agent aborts before its cleanup, most easily by dismissing the OpenAI key prompt. The previous run's generated code was then handed to the kernel. Clearing now happens once per batch, before anything runs, so it no longer depends on the agent getting far enough to do it. It is scoped to the agents in that batch: running one agent block must not throw away a sibling's output, and the cell an agent executes while generating it arrives without its agent, so it has to survive. executeAgentCell verifies the precondition rather than assuming it — nothing enforces it across call sites, and running against a dirty notebook fails silently, feeding the agent its own previous output and appending a second copy below the stale cells. The placeholder controller no longer executes anything. Running an agent block with no environment configured raised the environment picker mid-run, from inside the agent's own code tool, and then disposed the placeholder that owned the running execution. It now prompts and tells the user to run again, which also leaves one execution loop in the codebase instead of two that had already drifted apart. Block id reads are consolidated into getBlockId, which was spelled out inline in four places with three different fallback chains. Includes three edits that were already in the working tree: dropping @deepnote/runtime-core from the web externals, defaulting integrations to an empty list, and no longer coercing agent block content to ''. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- build/esbuild/build.ts | 3 +- .../controllers/vscodeNotebookController.ts | 14 +- .../deepnote/agentCellExecutionHandler.ts | 103 ++++++-- .../agentCellExecutionHandler.unit.test.ts | 241 ++++++++++++------ .../converters/agentBlockConverter.ts | 2 +- src/notebooks/deepnote/dataConversionUtils.ts | 26 ++ .../deepnote/dataConversionUtils.unit.test.ts | 57 ++++- .../deepnote/deepnoteDataConverter.ts | 4 +- .../deepnote/deepnoteFileChangeWatcher.ts | 13 +- .../deepnoteKernelAutoSelector.node.ts | 67 +---- ...epnoteKernelAutoSelector.node.unit.test.ts | 78 +++--- src/platform/deepnote/pocket.ts | 7 +- src/platform/deepnote/pocket.unit.test.ts | 16 ++ 13 files changed, 402 insertions(+), 229 deletions(-) diff --git a/build/esbuild/build.ts b/build/esbuild/build.ts index 77b7fa66ab..d9bf27dd73 100644 --- a/build/esbuild/build.ts +++ b/build/esbuild/build.ts @@ -72,8 +72,7 @@ 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 - '@deepnote/runtime-core' // Uses tcp-port-used → net, only needed in desktop for agent block execution + 'mathjax-electron' // Uses Node.js path module, MathJax rendering handled differently in browser ]; const desktopExternals = [...commonExternals, ...deskTopNodeModulesToExternalize]; const bundleConfig = getBundleConfiguration(); diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index b472167ded..4c3c770a5e 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -91,7 +91,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 } from '../deepnote/agentCellExecutionHandler'; +import { executeAgentCell, removeEphemeralCellsForAgentBlocks } from '../deepnote/agentCellExecutionHandler'; import { isAgentCell } from '../deepnote/dataConversionUtils'; /** @@ -626,17 +626,19 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont if (!this.cellQueue.has(doc)) { return; } - const allCells = this.cellQueue.get(doc) || []; + const queuedCells = 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 cellsToExecute = await removeEphemeralCellsForAgentBlocks(doc, queuedCells); + // 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[] = []; - for (const cell of allCells) { + for (const cell of cellsToExecute) { if (!isAgentCell(cell)) { pendingKernelCells.push(cell); continue; @@ -657,9 +659,9 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont // 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. + // `pendingKernelCells` holds NotebookCell references captured earlier in the batch; the document + // may have changed meanwhile (user delete, overlapping run that calls removeEphemeralCellsForAgentBlocks, + // etc.). Stale handles report index -1; createNotebookCellExecution throws and would abort the batch. const kernelCells = cells.filter((cell) => cell.index >= 0); if (kernelCells.length === 0) { diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index c79742919c..0feb58d52f 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -30,7 +30,13 @@ import { ServiceContainer } from '../../platform/ioc/container'; import { logger } from '../../platform/logging'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; import { IDeepnoteNotebookManager } from '../types'; -import { generateBlockId, generateSortingKey, isEphemeralCell } from './dataConversionUtils'; +import { + generateBlockId, + generateSortingKey, + getBlockId, + getEphemeralCellAgentSourceBlockId, + isAgentCell +} from './dataConversionUtils'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; import { getOrPromptOpenAiApiKey } from './deepnoteSecretStore'; @@ -49,7 +55,7 @@ function getProjectAgentContext(notebook: NotebookDocument): Pick(IDeepnoteNotebookManager); @@ -158,6 +164,14 @@ export interface ExecuteAgentCellOptions { executeAgentBlockFn?: typeof executeAgentBlock; } +/** + * Runs an agent block, streaming its progress into the cell's output and inserting the cells it + * generates below itself. + * + * Requires the cell's previous run to have been cleared first — call + * `removeEphemeralCellsForAgentBatch` on the batch. Never rejects: failures, including an uncleared + * previous run, are reported on the cell as stderr output and end the execution unsuccessfully. + */ export async function executeAgentCell( cell: NotebookCell, controller: NotebookController, @@ -198,16 +212,26 @@ 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(); + // Verify rather than assume the caller cleared them: nothing enforces the precondition across + // the three call sites, and running dirty fails silently — the agent would be handed its own + // previous output as context, and insertEphemeralCell appends below the stale cells rather + // than replacing them, so every run would leave another copy behind. + const staleCellCount = cell.notebook + .getCells() + .filter((c) => getEphemeralCellAgentSourceBlockId(c) === agentBlock.id).length; + + if (staleCellCount > 0) { + throw new Error( + `Agent block ${agentBlock.id} still has ${staleCellCount} generated cell(s) from its previous run` + ); + } - await removeEphemeralCellsForAgent(cell.notebook, agentBlock.id); + const openAiToken = await getOrPromptOpenAiApiKey(); 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. + // serializeNotebookContextFromBlocks does no ephemeral filtering, so this is safe only + // because of the precondition checked above. const notebookContext = serializeNotebookContext({ cells: cell.notebook.getCells().filter((c) => c.index !== cell.index), notebookName: (cell.notebook.metadata?.deepnoteNotebookName as string | undefined) ?? '' @@ -317,8 +341,7 @@ function getInsertIndexAfterAgentCell( let index = agentCellIndex + 1; while (index < notebook.cellCount) { - const cell = notebook.cellAt(index); - if (isEphemeralCell(cell) && cell.metadata?.agent_source_block_id === agentBlockId) { + if (getEphemeralCellAgentSourceBlockId(notebook.cellAt(index)) === agentBlockId) { index++; } else { break; @@ -366,9 +389,7 @@ async function insertEphemeralCell( 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); + const insertedCell = notebook.getCells().find((c) => getBlockId(c) === block.id); if (!insertedCell) { throw new Error(`Inserted ephemeral ${blockType} cell ${block.id} not found in notebook`); @@ -454,30 +475,64 @@ export async function executeEphemeralCell( } } -async function removeEphemeralCellsForAgent(notebook: NotebookDocument, agentBlockId: string): Promise { +/** + * Deletes the scratch cells the agent cells in `cells` generated on their previous run, and returns + * the batch without them. Call this before executing a batch of cells; `executeAgentCell` requires it. + * + * Ephemeral cells are agent-owned: the agent regenerates them on every run and the serializer never + * persists them. Left in the batch they would run the previous run's generated code against the + * kernel — so they are dropped up front, whether or not the agent that owns them gets far enough to + * replace them. + * + * Scoped to the agents present in the batch, because two callers legitimately run an ephemeral cell + * on its own: the user selecting one, and the agent executing the code cell it just generated via + * `notebook.cell.execute`. Neither carries its agent, so neither is touched. + * + * A rejected edit is logged rather than thrown: the agent cells re-check the notebook themselves and + * report it on the cell, and a failed edit is no reason to hold back the batch's ordinary code cells. + */ +export async function removeEphemeralCellsForAgentBlocks( + notebook: NotebookDocument, + cells: NotebookCell[] +): Promise { + const agentBlockIds = new Set( + cells + .filter(isAgentCell) + .map(getBlockId) + .filter((id): id is string => typeof id === 'string') + ); + + if (agentBlockIds.size === 0) { + return cells; + } + + const isOwnedScratch = (cell: NotebookCell) => { + const owner = getEphemeralCellAgentSourceBlockId(cell); + + return owner !== undefined && agentBlockIds.has(owner); + }; + + const remainingCells = cells.filter((cell) => !isOwnedScratch(cell)); const deletions: NotebookEdit[] = []; for (let i = notebook.cellCount - 1; i >= 0; i--) { - const cell = notebook.cellAt(i); - - if (isEphemeralCell(cell) && cell.metadata?.agent_source_block_id === agentBlockId) { + if (isOwnedScratch(notebook.cellAt(i))) { deletions.push(NotebookEdit.deleteCells(new NotebookRange(i, i + 1))); } } if (deletions.length === 0) { - return; + return remainingCells; } const edit = new WorkspaceEdit(); edit.set(notebook.uri, deletions); - // Fatal rather than a warning: the notebook context the agent receives is read off the live - // document, and Run All keeps these cells out of its kernel batch only by their index going - // negative once they are deleted. - if (!(await workspace.applyEdit(edit))) { - throw new Error(`Failed to remove ephemeral cells for agent block ${agentBlockId}`); + if (await workspace.applyEdit(edit)) { + logger.info(`Removed ${deletions.length} ephemeral cell(s) for ${agentBlockIds.size} agent block(s)`); + } else { + logger.error(`Failed to remove ephemeral cells for agent blocks ${[...agentBlockIds].join(', ')}`); } - logger.info(`Removed ${deletions.length} ephemeral cell(s) for agent block ${agentBlockId}`); + return remainingCells; } diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index 8d2a1290c0..109ffc11c1 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -12,6 +12,7 @@ import { NotebookCellOutput, NotebookCellOutputItem, NotebookController, + NotebookDocument, SecretStorage, SecretStorageChangeEvent, Uri, @@ -27,7 +28,12 @@ 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 { describeExecutionOutputs, executeAgentCell, executeEphemeralCell } from './agentCellExecutionHandler'; +import { + describeExecutionOutputs, + executeAgentCell, + executeEphemeralCell, + removeEphemeralCellsForAgentBlocks +} from './agentCellExecutionHandler'; import { IDeepnoteNotebookManager } from '../types'; import { createDeepnoteFile, createDeepnoteProject, createMockCell, createMockNotebook } from './deepnoteTestHelpers'; @@ -56,6 +62,51 @@ function stubSecretStorage(secretStorage: Map): ServiceContainer return serviceContainer; } +/** + * Makes `workspace.applyEdit` apply the notebook edits it is given to `cells`, so the code under test + * observes its own inserts and deletes. + * + * 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. + * + * Returns the number of edits applied so far — the shared `workspace` mock is never reset between + * tests, so its own call counts are useless here. + */ +function applyNotebookEditsTo(cells: NotebookCell[], notebook: NotebookDocument) { + type RecordedEdit = { range: { start: number; end: number }; newCells: NotebookCellData[] }; + let recordedEdits: RecordedEdit[] = []; + let appliedEdits = 0; + + sinon.stub(WorkspaceEdit.prototype, 'set').callsFake((_uri, edits) => { + recordedEdits = edits as unknown as RecordedEdit[]; + }); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { + appliedEdits++; + + 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: NotebookDocument }).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 { appliedEdits: () => appliedEdits }; +} + suite('AgentCellExecutionHandler', () => { const secretStorage = new Map(); let disposables: IDisposable[] = []; @@ -164,9 +215,6 @@ 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 }); @@ -179,33 +227,7 @@ suite('AgentCellExecutionHandler', () => { (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); - }); + applyNotebookEditsTo(cells, notebook); return { agentCell, cells, notebook }; } @@ -377,12 +399,10 @@ suite('AgentCellExecutionHandler', () => { 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)); - + // Clearing the previous run belongs to the caller. Running against a dirty notebook fails + // silently — the agent gets its own old output as context and appends a second copy below it + // — so the precondition is checked rather than assumed. + test('refuses to run rather than clearing the previous run itself', async () => { const previousResult = createMockCell({ text: 'print("previous run")', metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }, @@ -392,8 +412,13 @@ suite('AgentCellExecutionHandler', () => { await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + expect(executeAgentBlockStub.called).to.be.false; expect(mockExecution.end.firstCall.args[0]).to.be.false; expect(cells).to.include(previousResult); + + const [outputs] = mockExecution.appendOutput.firstCall.args as [NotebookCellOutput[]]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('previous run'); }); test('inserts a markdown cell after the agent cell with ephemeral metadata', async () => { @@ -450,46 +475,6 @@ suite('AgentCellExecutionHandler', () => { 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); - }); - - // The notebook context is read off the live document, and Run All keeps stale ephemeral cells - // out of the kernel batch only by their index going negative on deletion. - test('fails the run without calling the agent when the cleanup edit is rejected', async () => { - const previousResult = createMockCell({ - text: 'print("previous run")', - metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }, - index: 1 - }); - const { agentCell } = createAgentCellInMutableNotebook([previousResult]); - - when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(false)); - - await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - - expect(executeAgentBlockStub.called).to.be.false; - expect(mockExecution.end.firstCall.args[0]).to.be.false; - }); - test('passes project MCP servers and integrations to the agent', async () => { const integrations = [{ id: 'warehouse', name: 'Warehouse', type: 'postgres' }]; const mcpServers = [{ name: 'files', command: 'mcp-files', args: [] }]; @@ -516,6 +501,108 @@ suite('AgentCellExecutionHandler', () => { }); }); + suite('removeEphemeralCellsForAgentBatch', () => { + teardown(() => { + sinon.restore(); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(true)); + }); + + function createAgentCell(agentBlockId: string) { + return createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' }, id: agentBlockId }, + text: 'Test prompt' + }); + } + + function createEphemeralCell(agentBlockId: string, text: string) { + return createMockCell({ + metadata: { is_ephemeral: true, agent_source_block_id: agentBlockId }, + text + }); + } + + /** Wires the cells into a notebook whose list the applied edits actually mutate. */ + function createMutableNotebook(cells: NotebookCell[]) { + const notebook = createMockNotebook({ cells }); + + cells.forEach((cell, index) => { + (cell as { notebook: NotebookDocument }).notebook = notebook; + (cell as { index: number }).index = index; + }); + + return { notebook, ...applyNotebookEditsTo(cells, notebook) }; + } + + test('drops the previous run from the batch and deletes it from the notebook', async () => { + const agentCell = createAgentCell('agent-block-1'); + const previousResult = createEphemeralCell('agent-block-1', 'print("previous run")'); + const cells = [agentCell, previousResult]; + const { notebook } = createMutableNotebook(cells); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [...cells]); + + expect(batch).to.deep.equal([agentCell]); + expect(cells).to.deep.equal([agentCell]); + }); + + test('keeps another agent and ordinary cells', async () => { + const agentCell = createAgentCell('agent-block-1'); + const ownResult = createEphemeralCell('agent-block-1', 'own'); + const otherAgentResult = createEphemeralCell('agent-block-2', 'other agent'); + const userCell = createMockCell({ text: 'user code', metadata: {} }); + const cells = [agentCell, ownResult, otherAgentResult, userCell]; + const { notebook } = createMutableNotebook(cells); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [...cells]); + + expect(batch).to.deep.equal([agentCell, otherAgentResult, userCell]); + expect(cells).to.deep.equal([agentCell, otherAgentResult, userCell]); + }); + + // An agent runs the code cell it just generated through `notebook.cell.execute`, which arrives + // here as a batch of that cell alone. Dropping it would hang the agent until its timeout. + test('leaves an ephemeral cell whose agent is not in the batch', async () => { + const agentCell = createAgentCell('agent-block-1'); + const generatedCell = createEphemeralCell('agent-block-1', 'print("just generated")'); + const cells = [agentCell, generatedCell]; + const { notebook, appliedEdits } = createMutableNotebook(cells); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [generatedCell]); + + expect(batch).to.deep.equal([generatedCell]); + expect(cells).to.deep.equal([agentCell, generatedCell]); + expect(appliedEdits()).to.equal(0); + }); + + test('applies no edit when the batch has no agent cell', async () => { + const userCell = createMockCell({ text: 'user code', metadata: {} }); + const cells = [userCell]; + const { notebook, appliedEdits } = createMutableNotebook(cells); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [...cells]); + + expect(batch).to.deep.equal([userCell]); + expect(appliedEdits()).to.equal(0); + }); + + // The agent cells re-check the notebook and report it on the cell, so a rejected edit must not + // hold back the batch's ordinary code cells. + test('still drops the previous run from the batch when the edit is rejected', async () => { + const agentCell = createAgentCell('agent-block-1'); + const previousResult = createEphemeralCell('agent-block-1', 'print("previous run")'); + const userCell = createMockCell({ text: 'user code', metadata: {} }); + const cells = [agentCell, previousResult, userCell]; + const { notebook } = createMutableNotebook(cells); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(false)); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [...cells]); + + expect(batch).to.deep.equal([agentCell, userCell]); + expect(cells).to.deep.equal([agentCell, previousResult, userCell]); + }); + }); + suite('executeEphemeralCell', () => { teardown(() => { reset(mockedVSCodeNamespaces.commands); diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.ts b/src/notebooks/deepnote/converters/agentBlockConverter.ts index 6f9ebbd31c..3f5c6a4db8 100644 --- a/src/notebooks/deepnote/converters/agentBlockConverter.ts +++ b/src/notebooks/deepnote/converters/agentBlockConverter.ts @@ -15,7 +15,7 @@ import type { BlockConverter } from './blockConverter'; */ export class AgentBlockConverter implements BlockConverter { applyChangesToBlock(block: DeepnoteBlock, cell: NotebookCellData): void { - block.content = cell.value || ''; + block.content = cell.value; } canConvert(blockType: string): boolean { diff --git a/src/notebooks/deepnote/dataConversionUtils.ts b/src/notebooks/deepnote/dataConversionUtils.ts index fc9d227ccd..69fc7a8d86 100644 --- a/src/notebooks/deepnote/dataConversionUtils.ts +++ b/src/notebooks/deepnote/dataConversionUtils.ts @@ -45,6 +45,32 @@ export function isEphemeralCell(cell: NotebookCell | NotebookCellData): boolean return cell.metadata?.is_ephemeral === true; } +/** + * Returns the id of the block a cell is backed by, or undefined for a cell that has never been + * serialized. + * + * `__deepnoteBlockId` wins because VS Code may rewrite `id`, which is why the converter mirrors it. + * `deepnoteBlockId` is a third name the fallback-cell path writes; it is only ever set alongside the + * other two, so it resolves nothing new in practice and exists to tolerate metadata that has lost + * them. Losing an id here is worse than reading a redundant one: callers mint a fresh one, which + * reassigns the block on save. + */ +export function getBlockId(cell: NotebookCell | NotebookCellData): string | undefined { + return ( + (cell.metadata?.__deepnoteBlockId as string | undefined) || + (cell.metadata?.id as string | undefined) || + (cell.metadata?.deepnoteBlockId as string | undefined) + ); +} + +/** + * Returns the id of the agent block that generated this ephemeral cell, or undefined if the cell + * isn't agent-generated scratch. + */ +export function getEphemeralCellAgentSourceBlockId(cell: NotebookCell): string | undefined { + return isEphemeralCell(cell) ? (cell.metadata?.agent_source_block_id as string | undefined) : undefined; +} + /** * Generate sorting key based on index (format: a0, a1, ..., a99, b0, b1, ...) */ diff --git a/src/notebooks/deepnote/dataConversionUtils.unit.test.ts b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts index ab2f94efc5..0b9af65731 100644 --- a/src/notebooks/deepnote/dataConversionUtils.unit.test.ts +++ b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts @@ -1,6 +1,6 @@ import { expect } from 'chai'; -import { isAgentCell } from './dataConversionUtils'; +import { getBlockId, getEphemeralCellAgentSourceBlockId, isAgentCell } from './dataConversionUtils'; import { createMockCell } from './deepnoteTestHelpers'; suite('DataConversionUtils', () => { @@ -35,4 +35,59 @@ suite('DataConversionUtils', () => { expect(isAgentCell(cell)).to.be.false; }); }); + + suite('getBlockId', () => { + test('prefers the backup id VS Code cannot rewrite', () => { + const cell = createMockCell({ metadata: { __deepnoteBlockId: 'backup-id', id: 'rewritten-id' } }); + + expect(getBlockId(cell)).to.equal('backup-id'); + }); + + test('falls back to id when the backup is absent', () => { + const cell = createMockCell({ metadata: { id: 'block-id' } }); + + expect(getBlockId(cell)).to.equal('block-id'); + }); + + // The fallback-cell path writes this third name. Reading it beats minting a fresh id, which + // would reassign the block on save. + test('falls back to the legacy deepnoteBlockId when both are absent', () => { + const cell = createMockCell({ metadata: { deepnoteBlockId: 'legacy-id' } }); + + expect(getBlockId(cell)).to.equal('legacy-id'); + }); + + test('ranks the legacy name below both current ones', () => { + const cell = createMockCell({ metadata: { id: 'block-id', deepnoteBlockId: 'legacy-id' } }); + + expect(getBlockId(cell)).to.equal('block-id'); + }); + + test('returns undefined for a cell that was never serialized', () => { + const cell = createMockCell({ metadata: {} }); + + expect(getBlockId(cell)).to.be.undefined; + }); + }); + + suite('getEphemeralCellOwner', () => { + test('returns the agent block that generated the cell', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' } }); + + expect(getEphemeralCellAgentSourceBlockId(cell)).to.equal('agent-block-1'); + }); + + // An ordinary cell that happens to carry the metadata is not the agent's to delete. + test('returns undefined when the cell is not marked ephemeral', () => { + const cell = createMockCell({ metadata: { agent_source_block_id: 'agent-block-1' } }); + + expect(getEphemeralCellAgentSourceBlockId(cell)).to.be.undefined; + }); + + test('returns undefined for an ordinary cell', () => { + const cell = createMockCell({ metadata: {} }); + + expect(getEphemeralCellAgentSourceBlockId(cell)).to.be.undefined; + }); + }); }); diff --git a/src/notebooks/deepnote/deepnoteDataConverter.ts b/src/notebooks/deepnote/deepnoteDataConverter.ts index 51007cae9d..f8fae71b49 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.ts @@ -1,7 +1,7 @@ import { isExecutableBlock, type DeepnoteBlock } from '@deepnote/blocks'; import { NotebookCellData, NotebookCellKind, NotebookCellOutput, NotebookCellOutputItem } from 'vscode'; -import { generateBlockId, generateSortingKey } from './dataConversionUtils'; +import { generateBlockId, generateSortingKey, getBlockId } from './dataConversionUtils'; import type { DeepnoteOutput } from '../../platform/deepnote/deepnoteTypes'; import { ConverterRegistry } from './converters/converterRegistry'; import { BlockConverter } from './converters/blockConverter'; @@ -439,7 +439,7 @@ export class DeepnoteDataConverter { private createFallbackBlock(cell: NotebookCellData, index: number): DeepnoteBlock { const meta = cell.metadata as Record | undefined; - const preservedId = (meta?.__deepnoteBlockId ?? meta?.id ?? meta?.deepnoteBlockId) as string | undefined; + const preservedId = getBlockId(cell); const preservedSortingKey = (meta?.sortingKey ?? meta?.deepnoteSortingKey) as string | undefined; const preservedBlockGroup = meta?.blockGroup as string | undefined; diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index c5abf5a442..0a82635870 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -18,6 +18,7 @@ import { IExtensionSyncActivationService } from '../../platform/activation/types import { IDisposableRegistry } from '../../platform/common/types'; import { logger } from '../../platform/logging'; import { IDeepnoteNotebookManager } from '../types'; +import { getBlockId } from './dataConversionUtils'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; import { DeepnoteNotebookSerializer } from './deepnoteSerializer'; @@ -279,14 +280,14 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic const liveCells = notebook.getCells(); const liveOutputsByBlockId = new Map(); for (const liveCell of liveCells) { - const blockId = this.getBlockIdFromMetadata(liveCell.metadata); + const blockId = getBlockId(liveCell); if (blockId && liveCell.outputs.length > 0) { liveOutputsByBlockId.set(blockId, liveCell.outputs); } } for (const cell of newCells) { - const blockId = this.getBlockIdFromMetadata(cell.metadata); + const blockId = getBlockId(cell); if (blockId && (!cell.outputs || cell.outputs.length === 0)) { const liveOutputs = liveOutputsByBlockId.get(blockId); if (liveOutputs) { @@ -300,7 +301,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic edits.push(NotebookEdit.replaceCells(new NotebookRange(0, notebook.cellCount), newCells)); for (let i = 0; i < newCells.length; i++) { - const blockId = this.getBlockIdFromMetadata(newCells[i].metadata); + const blockId = getBlockId(newCells[i]); if (blockId) { edits.push( NotebookEdit.updateCellMetadata(i, { @@ -376,7 +377,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic for (let i = 0; i < liveCells.length; i++) { try { const cell = liveCells[i]; - let blockId = this.getBlockIdFromMetadata(cell.metadata); + let blockId = getBlockId(cell); let blockIdFromFallback = false; // Fallback to original project blocks when metadata was lost @@ -502,10 +503,6 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic logger.info(`[FileChangeWatcher] Updated notebook outputs from external snapshot: ${notebook.uri.path}`); } - private getBlockIdFromMetadata(metadata: Record | undefined): string | undefined { - return (metadata?.__deepnoteBlockId ?? metadata?.id) as string | undefined; - } - private handleFileChange(uri: Uri): void { // Deterministic self-write check — no timers involved if (this.consumeSelfWrite(uri)) { diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index aed1e61bba..d8bedbea79 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -57,8 +57,6 @@ import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { IControllerRegistration, IVSCodeNotebookController } from '../controllers/types'; import { IDeepnoteNotebookManager } from '../types'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; -import { executeAgentCell } from './agentCellExecutionHandler'; -import { isAgentCell } from './dataConversionUtils'; import { computeRequirementsHash } from './deepnoteProjectUtils'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; @@ -1093,7 +1091,9 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, controller.supportsExecutionOrder = true; controller.supportedLanguages = ['python', 'sql', 'markdown', 'plaintext']; - // Execution handler that shows environment picker when user tries to run without an environment + // Turns a Run gesture into the environment picker and nothing else. Executing here means + // executing without a kernel: configuring the environment disposes this controller mid-run + // (see ensureKernelSelectedWithConfiguration), orphaning any execution created from it. controller.executeHandler = async (cells, doc) => { logger.info( `Placeholder controller execute handler called for ${getDisplayPath(doc.uri)} with ${ @@ -1101,25 +1101,10 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } cells` ); - // 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. + // Setting up an environment runs a workspace-provided Python interpreter and installs into + // it, 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)}`); - - return; - } - - 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); - } - } + logger.info(`Workspace is not trusted, skipping environment setup for ${getDisplayPath(doc.uri)}`); return; } @@ -1142,45 +1127,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return; } - // Environment is now configured, execute the cells through the kernel - const docNotebookKey = getNotebookKey(doc.uri); - const realController = this.notebookControllers.get(docNotebookKey); - - if (!realController) { - logger.error(`No controller found after environment configuration for ${docNotebookKey}`); - - return; - } - - 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, { - metadata: realController.connection, - controller: realController.controller, - resourceUri: doc.uri - }); - - // Execute cells through the kernel - const kernelExecution = this.kernelProvider.getKernelExecution(kernel); - - // 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 { - if (isAgentCell(cell)) { - // Configuring the environment disposed this placeholder and handed the - // notebook to the real controller, which now owns its executions. - await executeAgentCell(cell, realController.controller); - } else { - await kernelExecution.executeCell(cell); - } - } catch (cellError) { - logger.error(`Error executing cell ${cell.index}`, cellError); - } - } - - logger.info(`Finished executing ${cells.length} cells`); + void window.showInformationMessage(l10n.t('Environment ready. Run the cells again to execute them.')); } catch (error) { if (isCancellationError(error)) { logger.info(`Environment setup cancelled for ${getDisplayPath(doc.uri)}`); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts index 2f6c4cae23..90592b09a7 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -1036,80 +1036,70 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { }); suite('Placeholder controller execution', () => { - function createExecutionStub() { - return { - start: sandbox.stub(), - end: sandbox.stub(), - clearOutput: sandbox.stub().resolves(), - replaceOutput: sandbox.stub().resolves(), - appendOutput: sandbox.stub().resolves(), - appendOutputItems: sandbox.stub().resolves() - }; - } - - // Configuring the environment disposes and deselects the placeholder, and VS Code throws for - // executions created on either a disposed or an unassociated controller. - test('runs agent cells on the real controller after configuring the environment', async () => { - const placeholderExecution = createExecutionStub(); + function createPlaceholder() { const placeholder = { supportsExecutionOrder: false, supportedLanguages: [] as string[], updateNotebookAffinity: sandbox.stub(), dispose: sandbox.stub(), - createNotebookCellExecution: sandbox.stub().returns(placeholderExecution) + createNotebookCellExecution: sandbox.stub() } as unknown as NotebookController; when( mockedVSCodeNamespaces.notebooks!.createNotebookController(anything(), anything(), anything()) ).thenReturn(placeholder); - when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(true); const onDidCloseNotebookDocument = new EventEmitter(); when(mockedVSCodeNamespaces.workspace.onDidCloseNotebookDocument).thenReturn( onDidCloseNotebookDocument.event ); - const realExecution = createExecutionStub(); - const realNotebookController = { - createNotebookCellExecution: sandbox.stub().returns(realExecution) - } as unknown as NotebookController; - const realController = mock(); - when(realController.controller).thenReturn(realNotebookController); - const internals = selector as unknown as { createPlaceholderController(notebook: NotebookDocument): NotebookController; - notebookControllers: Map; }; internals.createPlaceholderController(mockNotebook); - internals.notebookControllers.set(getNotebookKey(mockNotebook.uri), instance(realController)); - sandbox.stub(selector, 'ensureEnvironmentConfiguredBeforeExecution').resolves(true); + return placeholder; + } - const kernelExecution = { executeCell: sandbox.stub().resolves() }; - when(mockKernelProvider.getOrCreate(anything(), anything())).thenReturn(instance(mock())); - when(mockKernelProvider.getKernelExecution(anything())).thenReturn( - kernelExecution as unknown as ReturnType - ); + const agentCell = { + index: 0, + metadata: { __deepnotePocket: { type: 'agent' } } + } as unknown as NotebookCell; + const codeCell = { index: 1, metadata: {} } as unknown as NotebookCell; - const agentCell = { - index: 0, - metadata: { __deepnotePocket: { type: 'agent' } } - } as unknown as NotebookCell; - const codeCell = { index: 1, metadata: {} } as unknown as NotebookCell; + // This controller has no kernel, and configuring one disposes it mid-run — so it prompts and + // stops, rather than executing anything itself or handing the batch on. + test('configures the environment and executes nothing', async () => { + when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(true); + const placeholder = createPlaceholder(); + const ensureEnvironment = sandbox + .stub(selector, 'ensureEnvironmentConfiguredBeforeExecution') + .resolves(true); await placeholder.executeHandler!([agentCell, codeCell], mockNotebook, placeholder); - assert.isTrue( - (realNotebookController.createNotebookCellExecution as sinon.SinonStub).calledOnceWithExactly( - agentCell - ), - 'agent cell execution should be created on the real controller' - ); + assert.isTrue(ensureEnvironment.calledOnce, 'should prompt for an environment'); assert.isTrue( (placeholder.createNotebookCellExecution as sinon.SinonStub).notCalled, - 'no execution should be created on the disposed placeholder' + 'placeholder must not create executions' ); + verify(mockKernelProvider.getOrCreate(anything(), anything())).never(); + }); + + // Agent blocks spawn MCP servers declared by the workspace file, and setting up an environment + // runs a workspace-provided interpreter. + test('does nothing at all in an untrusted workspace', async () => { + when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(false); + const placeholder = createPlaceholder(); + const ensureEnvironment = sandbox + .stub(selector, 'ensureEnvironmentConfiguredBeforeExecution') + .resolves(true); + + await placeholder.executeHandler!([agentCell, codeCell], mockNotebook, placeholder); + + assert.isTrue(ensureEnvironment.notCalled, 'should not prompt in an untrusted workspace'); }); }); }); diff --git a/src/platform/deepnote/pocket.ts b/src/platform/deepnote/pocket.ts index 1bf7b6286c..fd2d4bbeb7 100644 --- a/src/platform/deepnote/pocket.ts +++ b/src/platform/deepnote/pocket.ts @@ -2,7 +2,7 @@ import type { DeepnoteBlock, ExecutableBlock } from '@deepnote/blocks'; import { isExecutableBlockType } from '@deepnote/blocks'; import { NotebookCellKind, type NotebookCellData } from 'vscode'; -import { generateBlockId, generateSortingKey } from '../../notebooks/deepnote/dataConversionUtils'; +import { generateBlockId, generateSortingKey, getBlockId } from '../../notebooks/deepnote/dataConversionUtils'; import { logger } from '../logging'; import { generateUuid } from '../common/uuid'; @@ -74,9 +74,8 @@ export function createBlockFromPocket(cell: NotebookCellData, index: number): De const pocket = extractPocketFromCellMetadata(cell); const metadata = cell.metadata ? { ...cell.metadata } : undefined; - // Get id from top-level metadata before cleaning it up - // Check both 'id' and backup '__deepnoteBlockId' in case VS Code modifies 'id' - const cellId = (metadata?.__deepnoteBlockId as string | undefined) || (metadata?.id as string | undefined); + // Read the id before the copy below is stripped of it + const cellId = getBlockId(cell); logger.debug( `[Pocket] createBlockFromPocket index=${index}: cell.metadata.id=${metadata?.id}, __deepnoteBlockId=${metadata?.__deepnoteBlockId}, using cellId=${cellId}, metadata keys=${ diff --git a/src/platform/deepnote/pocket.unit.test.ts b/src/platform/deepnote/pocket.unit.test.ts index b8cff6033f..ae9a699f32 100644 --- a/src/platform/deepnote/pocket.unit.test.ts +++ b/src/platform/deepnote/pocket.unit.test.ts @@ -135,6 +135,22 @@ suite('Pocket', () => { assert.strictEqual((block as any).outputs, undefined); }); + // VS Code may rewrite `id`, which is the whole reason the converter mirrors it into + // `__deepnoteBlockId`. Losing this preference silently reassigns block ids on every save. + test('takes the id from the backup rather than a rewritten id', () => { + const cell = new NotebookCellData(NotebookCellKind.Code, 'print("hello")', 'python'); + + cell.metadata = { + __deepnotePocket: { type: 'code', sortingKey: 'a0' }, + __deepnoteBlockId: 'block-123', + id: 'rewritten-by-vscode' + }; + + const block = createBlockFromPocket(cell, 0); + + assert.strictEqual(block.id, 'block-123'); + }); + test('creates block with generated ID and sortingKey when no pocket exists', () => { const cell = new NotebookCellData(NotebookCellKind.Code, 'print("hello")', 'python'); From 604d90c8722e414699984a96da86625fbeb47c7c Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 5 Aug 2026 10:57:25 +0000 Subject: [PATCH 33/37] fix(build): keep runtime-core out of the web bundle MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bundling it for the browser pulls in `net` via tcp-port-used and `child_process` via the MCP stdio transport, which spawns servers — neither resolves for a browser target, and esbuild fails the web build outright. Agent blocks execute on desktop only, so the web bundle has no use for the package. Restores the external dropped in 84c3cedfe, which broke CI, CD and E2E: all three package the extension, and only `compile-tsc` was run before pushing. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- build/esbuild/build.ts | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/build/esbuild/build.ts b/build/esbuild/build.ts index d9bf27dd73..f55108ae4c 100644 --- a/build/esbuild/build.ts +++ b/build/esbuild/build.ts @@ -72,7 +72,11 @@ 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 + // Reaches Node built-ins the browser has no answer for — `net` via tcp-port-used, and + // `child_process` via the MCP stdio transport, which spawns servers. Agent blocks execute on + // desktop only. + '@deepnote/runtime-core' ]; const desktopExternals = [...commonExternals, ...deskTopNodeModulesToExternalize]; const bundleConfig = getBundleConfiguration(); From b7f9ac794954b4a1b3aabfa862a4cae75ffd51c8 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 5 Aug 2026 14:18:24 +0000 Subject: [PATCH 34/37] refactor(agent-block): tighten PR code per engineering standards Trim redundant comments, DRY test helpers, and improve e2e assertions without mirror expects; align secret store tests with Chai assert. Co-authored-by: Cursor --- .github/workflows/e2e.yml | 5 +- build/esbuild/build.ts | 5 +- .../controllers/vscodeNotebookController.ts | 10 +- .../deepnote/agentCellExecutionHandler.ts | 50 +++---- .../agentCellExecutionHandler.unit.test.ts | 33 ++--- .../deepnote/agentCellStatusBarProvider.ts | 21 +-- .../converters/agentBlockConverter.ts | 11 +- .../deepnote/dataConversionUtils.unit.test.ts | 2 +- .../deepnote/deepnoteDataConverter.ts | 2 +- .../deepnoteKernelAutoSelector.node.ts | 6 +- ...epnoteKernelAutoSelector.node.unit.test.ts | 4 - src/notebooks/deepnote/deepnoteSecretStore.ts | 2 - .../deepnote/deepnoteSecretStore.unit.test.ts | 52 +++---- src/notebooks/deepnote/deepnoteTestHelpers.ts | 4 - .../ephemeralCellDecorationProvider.ts | 45 +++--- .../ephemeralCellStatusBarProvider.ts | 6 +- src/platform/deepnote/pocket.unit.test.ts | 2 - src/renderers/client/markdown.ts | 5 +- src/test/mocks/deepnoteRuntimeCore.ts | 2 +- test/e2e/.mocharc.js | 11 +- test/e2e/helpers/mockOpenAiServer.ts | 83 +---------- test/e2e/helpers/notebook.ts | 17 +-- test/e2e/suite/agentBlock.e2e.test.ts | 137 +++--------------- 23 files changed, 132 insertions(+), 383 deletions(-) diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index 3c5b674fc8..81f46448d7 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -72,10 +72,7 @@ jobs: run: npm run setup:e2e:deps - name: Pre-download the mock LLM server - # The agent-block suite starts this mid-run via npx. setup-node's cache is keyed on the - # lockfile, which aimock is deliberately absent from, so ~/.npm/_npx is never restored and the - # fetch would otherwise happen inside the test — where a registry blip surfaces as an opaque - # start-up timeout. Failing here instead points straight at the cause. + # aimock is npx-only (not in lockfile), so setup-node never restores ~/.npm/_npx. run: npm run setup:e2e:mock - name: Cache pip wheel downloads diff --git a/build/esbuild/build.ts b/build/esbuild/build.ts index f55108ae4c..2c3955037a 100644 --- a/build/esbuild/build.ts +++ b/build/esbuild/build.ts @@ -73,10 +73,7 @@ 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 - // Reaches Node built-ins the browser has no answer for — `net` via tcp-port-used, and - // `child_process` via the MCP stdio transport, which spawns servers. Agent blocks execute on - // desktop only. - '@deepnote/runtime-core' + '@deepnote/runtime-core' // Node built-ins (net, child_process); agent blocks run on desktop only ]; const desktopExternals = [...commonExternals, ...deskTopNodeModulesToExternalize]; const bundleConfig = getBundleConfiguration(); diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index 4c3c770a5e..7f61d0ab59 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -627,15 +627,11 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont return; } const queuedCells = 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. + // Clear before await so agent-driven re-entrant runs start with an empty queue. this.cellQueue.delete(doc); const cellsToExecute = await removeEphemeralCellsForAgentBlocks(doc, queuedCells); - // 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[] = []; for (const cell of cellsToExecute) { @@ -659,9 +655,7 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont // Creating these execution objects marks the cell as queued for execution (vscode will update cell UI). type CellExec = { cell: NotebookCell; exec: NotebookCellExecution }; - // `pendingKernelCells` holds NotebookCell references captured earlier in the batch; the document - // may have changed meanwhile (user delete, overlapping run that calls removeEphemeralCellsForAgentBlocks, - // etc.). Stale handles report index -1; createNotebookCellExecution throws and would abort the batch. + // Stale cell handles report index -1; createNotebookCellExecution would abort the batch. const kernelCells = cells.filter((cell) => cell.index >= 0); if (kernelCells.length === 0) { diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index 0feb58d52f..81f54d77b0 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -2,6 +2,7 @@ import { CancellationError, CancellationToken, NotebookCell, + NotebookCellData, NotebookCellOutput, NotebookCellOutputItem, NotebookController, @@ -86,6 +87,20 @@ function getProjectAgentContext(notebook: NotebookDocument): Pick((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 - ); + const block = converter.convertCellToBlock(notebookCellDataFromCell(cell), cell.index); acc.push(block); } catch (error) { logger.error(`Error converting cell to block: ${error}`); @@ -169,7 +175,7 @@ export interface ExecuteAgentCellOptions { * generates below itself. * * Requires the cell's previous run to have been cleared first — call - * `removeEphemeralCellsForAgentBatch` on the batch. Never rejects: failures, including an uncleared + * `removeEphemeralCellsForAgentBlocks` on the batch. Never rejects: failures, including an uncleared * previous run, are reported on the cell as stderr output and end the execution unsuccessfully. */ export async function executeAgentCell( @@ -195,20 +201,10 @@ export async function executeAgentCell( await execution.replaceOutput([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 deepnoteBlock = dataConverter.convertCellToBlock(notebookCellDataFromCell(cell), 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'); } @@ -247,9 +243,7 @@ export async function executeAgentCell( 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}`; + return `Failed to add markdown block: ${toError(error).message}`; } }, addAndExecuteCodeBlock: async ({ code }: { code: string }) => { @@ -267,9 +261,7 @@ export async function executeAgentCell( 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}`; + return `Execution error: ${toError(error).message}`; } }, onAgentEvent: async (event: AgentStreamEvent) => { @@ -398,7 +390,7 @@ async function insertEphemeralCell( return insertedCell; } -const EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS = 5 * 60 * 1000; +export const EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS = 5 * 60 * 1000; export interface EphemeralCellExecutionResult { success: boolean; diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index 109ffc11c1..3c3dbd8d9c 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -24,12 +24,13 @@ 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 { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; +import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; import { describeExecutionOutputs, + EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS, executeAgentCell, executeEphemeralCell, removeEphemeralCellsForAgentBlocks @@ -232,6 +233,12 @@ suite('AgentCellExecutionHandler', () => { return { agentCell, cells, notebook }; } + function getStdoutChunkText(callIndex: number): string { + const item = mockExecution.appendOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; + + return Buffer.from(item.data).toString('utf-8'); + } + test('creates execution and starts it', async () => { const cell = createAgentCell('Analyze data'); @@ -297,14 +304,8 @@ suite('AgentCellExecutionHandler', () => { await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - const getChunkText = (callIndex: number): string => { - const item = mockExecution.appendOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; - - return Buffer.from(item.data).toString('utf-8'); - }; - - expect(getChunkText(0)).to.equal('[Agent] Text:\nfirst'); - expect(getChunkText(1)).to.equal(' second'); + expect(getStdoutChunkText(0)).to.equal('[Agent] Text:\nfirst'); + expect(getStdoutChunkText(1)).to.equal(' second'); }); test('separates different event types with blank lines', async () => { @@ -319,13 +320,7 @@ suite('AgentCellExecutionHandler', () => { await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - const getChunkText = (callIndex: number): string => { - const item = mockExecution.appendOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; - - return Buffer.from(item.data).toString('utf-8'); - }; - - const chunk2 = getChunkText(1); + const chunk2 = getStdoutChunkText(1); expect(chunk2).to.include('\n\n'); expect(chunk2).to.include('[Agent] Tool called: search'); }); @@ -501,7 +496,7 @@ suite('AgentCellExecutionHandler', () => { }); }); - suite('removeEphemeralCellsForAgentBatch', () => { + suite('removeEphemeralCellsForAgentBlocks', () => { teardown(() => { sinon.restore(); when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(true)); @@ -680,7 +675,7 @@ suite('AgentCellExecutionHandler', () => { ); const resultPromise = executeEphemeralCell(cell); - await clock.tickAsync(5 * 60 * 1000); + await clock.tickAsync(EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS); const result = await resultPromise; diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index 6ccd26fe6a..7e36a1aca0 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -22,15 +22,14 @@ import { clearOpenAiApiKey, promptForOpenAiApiKey } from './deepnoteSecretStore' /** 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. */ +/** Must be stored explicitly; a missing key becomes `undefined` in runtime-core and is passed to openai() as the model name. */ 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 - * and the AI model picker. - */ +const AGENT_INDICATOR_PRIORITY = 100; +const MODEL_PICKER_PRIORITY = 90; + @injectable() export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService { private readonly disposables: Disposable[] = []; @@ -78,9 +77,7 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv } public dispose(): void { - for (const disposable of this.disposables) { - disposable.dispose(); - } + this.disposables.forEach((disposable) => disposable.dispose()); } public provideCellStatusBarItems( @@ -105,7 +102,7 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv return { text: `$(hubot) ${l10n.t('Agent Block')}`, alignment: 1, - priority: 100, + priority: AGENT_INDICATOR_PRIORITY, tooltip: l10n.t('Deepnote Agent Block\nAI-powered block that autonomously generates code and analysis') }; } @@ -114,7 +111,7 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv return { text: `$(symbol-enum) ${l10n.t('Model: {0}', model)}`, alignment: 1, - priority: 90, + priority: MODEL_PICKER_PRIORITY, tooltip: l10n.t('AI Model: {0}\nClick to change', model), command: { title: l10n.t('Switch Model'), @@ -163,16 +160,12 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv return; } - // 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 { 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]; diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.ts b/src/notebooks/deepnote/converters/agentBlockConverter.ts index 3f5c6a4db8..ebf6cdff2c 100644 --- a/src/notebooks/deepnote/converters/agentBlockConverter.ts +++ b/src/notebooks/deepnote/converters/agentBlockConverter.ts @@ -3,16 +3,7 @@ import { NotebookCellData, NotebookCellKind } from 'vscode'; import type { BlockConverter } from './blockConverter'; -/** - * Converter for agent blocks. - * - * 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. - */ +/** Agent prompts render as plaintext code cells; metadata passes through in DeepnoteDataConverter. */ export class AgentBlockConverter implements BlockConverter { applyChangesToBlock(block: DeepnoteBlock, cell: NotebookCellData): void { block.content = cell.value; diff --git a/src/notebooks/deepnote/dataConversionUtils.unit.test.ts b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts index 0b9af65731..2b39574c45 100644 --- a/src/notebooks/deepnote/dataConversionUtils.unit.test.ts +++ b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts @@ -70,7 +70,7 @@ suite('DataConversionUtils', () => { }); }); - suite('getEphemeralCellOwner', () => { + suite('getEphemeralCellAgentSourceBlockId', () => { test('returns the agent block that generated the cell', () => { const cell = createMockCell({ metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' } }); diff --git a/src/notebooks/deepnote/deepnoteDataConverter.ts b/src/notebooks/deepnote/deepnoteDataConverter.ts index f8fae71b49..09b06cc2be 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.ts @@ -3,6 +3,7 @@ import { NotebookCellData, NotebookCellKind, NotebookCellOutput, NotebookCellOut import { generateBlockId, generateSortingKey, getBlockId } from './dataConversionUtils'; import type { DeepnoteOutput } from '../../platform/deepnote/deepnoteTypes'; +import { AgentBlockConverter } from './converters/agentBlockConverter'; import { ConverterRegistry } from './converters/converterRegistry'; import { BlockConverter } from './converters/blockConverter'; import { CodeBlockConverter } from './converters/codeBlockConverter'; @@ -11,7 +12,6 @@ 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 diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index d8bedbea79..eecb22ec2c 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -1091,9 +1091,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, controller.supportsExecutionOrder = true; controller.supportedLanguages = ['python', 'sql', 'markdown', 'plaintext']; - // Turns a Run gesture into the environment picker and nothing else. Executing here means - // executing without a kernel: configuring the environment disposes this controller mid-run - // (see ensureKernelSelectedWithConfiguration), orphaning any execution created from it. + // Run here only prompts for an environment; kernel execution uses the real controller afterward. controller.executeHandler = async (cells, doc) => { logger.info( `Placeholder controller execute handler called for ${getDisplayPath(doc.uri)} with ${ @@ -1101,8 +1099,6 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } cells` ); - // Setting up an environment runs a workspace-provided Python interpreter and installs into - // it, so gate this path the same way VSCodeNotebookController gates its own execute handler. if (!workspace.isTrusted) { logger.info(`Workspace is not trusted, skipping environment setup for ${getDisplayPath(doc.uri)}`); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts index 90592b09a7..159124e7a0 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -1069,8 +1069,6 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { } as unknown as NotebookCell; const codeCell = { index: 1, metadata: {} } as unknown as NotebookCell; - // This controller has no kernel, and configuring one disposes it mid-run — so it prompts and - // stops, rather than executing anything itself or handing the batch on. test('configures the environment and executes nothing', async () => { when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(true); const placeholder = createPlaceholder(); @@ -1088,8 +1086,6 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { verify(mockKernelProvider.getOrCreate(anything(), anything())).never(); }); - // Agent blocks spawn MCP servers declared by the workspace file, and setting up an environment - // runs a workspace-provided interpreter. test('does nothing at all in an untrusted workspace', async () => { when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(false); const placeholder = createPlaceholder(); diff --git a/src/notebooks/deepnote/deepnoteSecretStore.ts b/src/notebooks/deepnote/deepnoteSecretStore.ts index fadf11bd42..9e940557f6 100644 --- a/src/notebooks/deepnote/deepnoteSecretStore.ts +++ b/src/notebooks/deepnote/deepnoteSecretStore.ts @@ -87,8 +87,6 @@ export async function getOrPromptSecret( return value; } -// OpenAI API key - specific wrappers - const OPENAI_API_KEY = 'openAiApiKey'; const OPENAI_PROMPT_OPTIONS: SecretPromptOptions = { diff --git a/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts index e99bafc638..3ba41edd7f 100644 --- a/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts @@ -1,8 +1,9 @@ -import { expect } from 'chai'; +import { assert } from 'chai'; import * as sinon from 'sinon'; import { anything, instance, mock, when } from 'ts-mockito'; import { EventEmitter, ExtensionMode, SecretStorage, SecretStorageChangeEvent } from 'vscode'; +import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; import { IExtensionContext } from '../../platform/common/types'; import { ServiceContainer } from '../../platform/ioc/container'; import { @@ -17,7 +18,6 @@ import { setOpenAiApiKey, setSecret } from './deepnoteSecretStore'; -import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; suite('deepnoteSecretStore', () => { const secretStorage = new Map(); @@ -61,13 +61,13 @@ suite('deepnoteSecretStore', () => { const value = await getSecret('customKey'); - expect(value).to.equal('custom-value'); + assert.strictEqual(value, 'custom-value'); }); test('returns undefined when not set', async () => { const value = await getSecret('customKey'); - expect(value).to.be.undefined; + assert.isUndefined(value); }); test('returns undefined when value is empty string', async () => { @@ -75,7 +75,7 @@ suite('deepnoteSecretStore', () => { const value = await getSecret('customKey'); - expect(value).to.be.undefined; + assert.isUndefined(value); }); }); @@ -83,7 +83,7 @@ suite('deepnoteSecretStore', () => { test('stores value in secrets', async () => { await setSecret('customKey', 'custom-value'); - expect(secretStorage.get('customKey')).to.equal('custom-value'); + assert.strictEqual(secretStorage.get('customKey'), 'custom-value'); }); }); @@ -93,7 +93,7 @@ suite('deepnoteSecretStore', () => { await clearSecret('customKey'); - expect(secretStorage.has('customKey')).to.be.false; + assert.isFalse(secretStorage.has('customKey')); }); }); @@ -107,8 +107,8 @@ suite('deepnoteSecretStore', () => { password: false }); - expect(value).to.equal('user-input'); - expect(secretStorage.get('customKey')).to.equal('user-input'); + assert.strictEqual(value, 'user-input'); + assert.strictEqual(secretStorage.get('customKey'), 'user-input'); }); test('returns undefined when user cancels', async () => { @@ -116,7 +116,7 @@ suite('deepnoteSecretStore', () => { const value = await promptForSecret('customKey', { prompt: 'Enter value' }); - expect(value).to.be.undefined; + assert.isUndefined(value); }); }); @@ -126,7 +126,7 @@ suite('deepnoteSecretStore', () => { const value = await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); - expect(value).to.equal('stored-value'); + assert.strictEqual(value, 'stored-value'); }); test('throws when value missing and user cancels prompt', async () => { @@ -134,9 +134,9 @@ suite('deepnoteSecretStore', () => { try { await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); - expect.fail('Should have thrown'); + assert.fail('Should have thrown'); } catch (e) { - expect((e as Error).message).to.equal('Value is required'); + assert.strictEqual((e as Error).message, 'Value is required'); } }); }); @@ -147,13 +147,13 @@ suite('deepnoteSecretStore', () => { const key = await getOpenAiApiKey(); - expect(key).to.equal('test-key'); + assert.strictEqual(key, 'test-key'); }); test('returns undefined when not set', async () => { const key = await getOpenAiApiKey(); - expect(key).to.be.undefined; + assert.isUndefined(key); }); test('returns undefined when key is empty string', async () => { @@ -161,7 +161,7 @@ suite('deepnoteSecretStore', () => { const key = await getOpenAiApiKey(); - expect(key).to.be.undefined; + assert.isUndefined(key); }); }); @@ -169,7 +169,7 @@ suite('deepnoteSecretStore', () => { test('stores key in secrets', async () => { await setOpenAiApiKey('my-api-key'); - expect(secretStorage.get('openAiApiKey')).to.equal('my-api-key'); + assert.strictEqual(secretStorage.get('openAiApiKey'), 'my-api-key'); }); }); @@ -179,7 +179,7 @@ suite('deepnoteSecretStore', () => { await clearOpenAiApiKey(); - expect(secretStorage.has('openAiApiKey')).to.be.false; + assert.isFalse(secretStorage.has('openAiApiKey')); }); }); @@ -189,8 +189,8 @@ suite('deepnoteSecretStore', () => { const key = await promptForOpenAiApiKey(); - expect(key).to.equal('sk-abc123'); - expect(secretStorage.get('openAiApiKey')).to.equal('sk-abc123'); + assert.strictEqual(key, 'sk-abc123'); + assert.strictEqual(secretStorage.get('openAiApiKey'), 'sk-abc123'); }); test('returns undefined when user cancels', async () => { @@ -198,7 +198,7 @@ suite('deepnoteSecretStore', () => { const key = await promptForOpenAiApiKey(); - expect(key).to.be.undefined; + assert.isUndefined(key); }); test('returns undefined when user enters empty string', async () => { @@ -206,7 +206,7 @@ suite('deepnoteSecretStore', () => { const key = await promptForOpenAiApiKey(); - expect(key).to.be.undefined; + assert.isUndefined(key); }); }); @@ -216,7 +216,7 @@ suite('deepnoteSecretStore', () => { const key = await getOrPromptOpenAiApiKey(); - expect(key).to.equal('stored-key'); + assert.strictEqual(key, 'stored-key'); }); test('prompts and returns key when missing', async () => { @@ -224,7 +224,7 @@ suite('deepnoteSecretStore', () => { const key = await getOrPromptOpenAiApiKey(); - expect(key).to.equal('prompted-key'); + assert.strictEqual(key, 'prompted-key'); }); test('throws when key missing and user cancels prompt', async () => { @@ -232,9 +232,9 @@ suite('deepnoteSecretStore', () => { try { await getOrPromptOpenAiApiKey(); - expect.fail('Should have thrown'); + assert.fail('Should have thrown'); } catch (e) { - expect((e as Error).message).to.include('OpenAI API key is not set'); + assert.include((e as Error).message, 'OpenAI API key is not set'); } }); }); diff --git a/src/notebooks/deepnote/deepnoteTestHelpers.ts b/src/notebooks/deepnote/deepnoteTestHelpers.ts index e00762286e..b93e0bae59 100644 --- a/src/notebooks/deepnote/deepnoteTestHelpers.ts +++ b/src/notebooks/deepnote/deepnoteTestHelpers.ts @@ -47,10 +47,6 @@ 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[]; } diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts index 19ca8b76fc..2e1108b52e 100644 --- a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts @@ -12,21 +12,12 @@ import { } from 'vscode'; import { injectable } from 'inversify'; -import { isEphemeralCell } from './dataConversionUtils'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { isEphemeralCell } from './dataConversionUtils'; 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`. - */ +/** Code cell editor decorations for `is_ephemeral` blocks; markup uses `src/renderers/client/markdown.ts`. */ @injectable() export class EphemeralCellDecorationProvider implements IExtensionSyncActivationService { private readonly disposables: Disposable[] = []; @@ -104,23 +95,27 @@ export class EphemeralCellDecorationProvider implements IExtensionSyncActivation private updateDecorations(): void { for (const editor of window.visibleTextEditors) { - if (editor.document.uri.scheme !== NOTEBOOK_CELL_SCHEME) { - continue; - } + try { + if (editor.document.uri.scheme !== NOTEBOOK_CELL_SCHEME) { + continue; + } - const cell = this.findCellForEditor(editor); - if (!cell || !isEphemeralCell(cell)) { - editor.setDecorations(this.ephemeralDecorationType, []); - continue; - } + const cell = this.findCellForEditor(editor); + if (!cell || !isEphemeralCell(cell)) { + 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); - } + 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); + editor.setDecorations(this.ephemeralDecorationType, lineRanges); + } catch { + continue; + } } } } diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts index 60c0a67fba..12a7f08d0b 100644 --- a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts @@ -11,15 +11,11 @@ import { } from 'vscode'; import { injectable } from 'inversify'; -import { isEphemeralCell } from './dataConversionUtils'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { isEphemeralCell } from './dataConversionUtils'; 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 diff --git a/src/platform/deepnote/pocket.unit.test.ts b/src/platform/deepnote/pocket.unit.test.ts index ae9a699f32..73b317561a 100644 --- a/src/platform/deepnote/pocket.unit.test.ts +++ b/src/platform/deepnote/pocket.unit.test.ts @@ -135,8 +135,6 @@ suite('Pocket', () => { assert.strictEqual((block as any).outputs, undefined); }); - // VS Code may rewrite `id`, which is the whole reason the converter mirrors it into - // `__deepnoteBlockId`. Losing this preference silently reassigns block ids on every save. test('takes the id from the backup rather than a rewritten id', () => { const cell = new NotebookCellData(NotebookCellKind.Code, 'print("hello")', 'python'); diff --git a/src/renderers/client/markdown.ts b/src/renderers/client/markdown.ts index f18a4392df..8537598b5e 100644 --- a/src/renderers/client/markdown.ts +++ b/src/renderers/client/markdown.ts @@ -1,7 +1,6 @@ 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. +// Minimal markdown-it surface (no package types; transitive dependency only). interface MarkdownItToken { content: string; } @@ -87,8 +86,6 @@ export const activate: ActivationFunction = async (ctx) => { document.head.appendChild(template); const markdownRenderer = await ctx.getRenderer('vscode.markdown-it-renderer'); - // 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') { diff --git a/src/test/mocks/deepnoteRuntimeCore.ts b/src/test/mocks/deepnoteRuntimeCore.ts index 9470a8a8ce..4dccd7953f 100644 --- a/src/test/mocks/deepnoteRuntimeCore.ts +++ b/src/test/mocks/deepnoteRuntimeCore.ts @@ -1,5 +1,5 @@ -import type { AgentBlockContext, ServerInfo, ServerOptions } from '@deepnote/runtime-core'; import type { AgentBlock } from '@deepnote/blocks'; +import type { AgentBlockContext, ServerInfo, ServerOptions } from '@deepnote/runtime-core'; import type { ChildProcess } from 'child_process'; /** diff --git a/test/e2e/.mocharc.js b/test/e2e/.mocharc.js index 70a9f99b81..50d9e789fd 100644 --- a/test/e2e/.mocharc.js +++ b/test/e2e/.mocharc.js @@ -3,14 +3,7 @@ // tests are the real guard rails; this is a generous suite-level safety net. const path = require('path'); -// Loaded here rather than declared via mocha's `require` option, which only the mocha CLI acts on: -// ExTester hands this config straight to `new Mocha(config)` (vscode-extension-tester -// suite/runner.js), and the constructor reads `rootHooks` — already-resolved hook objects — while -// ignoring `require` entirely. Declared the other way the file is never loaded and the hooks below -// silently never run. -// -// Requires compiled output, so compile-e2e must run first; a missing build now fails here rather -// than passing with the hooks quietly absent. +// ExTester uses `new Mocha(config)` and ignores `require`; load rootHooks here as `rootHooks`. const { mochaHooks } = require(path.resolve(__dirname, '..', '..', 'out', 'e2e', 'rootHooks.js')); module.exports = { @@ -18,7 +11,5 @@ module.exports = { retries: 1, // absorb transient UI flakiness with a single retry reporter: 'spec', color: true, - // Dismiss notification toasts between tests so they don't accumulate across the one shared - // VS Code instance. rootHooks: mochaHooks }; diff --git a/test/e2e/helpers/mockOpenAiServer.ts b/test/e2e/helpers/mockOpenAiServer.ts index 8d7833ffaf..b998e99af5 100644 --- a/test/e2e/helpers/mockOpenAiServer.ts +++ b/test/e2e/helpers/mockOpenAiServer.ts @@ -5,48 +5,26 @@ import * as os from 'os'; import * as path from 'path'; import { setTimeout as delay } from 'timers/promises'; -// Fetched by npx rather than installed: aimock declares `jest` and `vitest` as peers, and resolving -// those against this repo's tree forces overrides that would outlive the test. npx resolves in its -// own cache, so the dependency graph here is untouched. -// -// Pinned exactly — a range would let the mock the suite asserts against change underneath it. +// aimock is npx-only so its peer deps (jest/vitest) never enter this repo's lockfile. const AIMOCK_VERSION = '1.37.4'; -// `llmock` is the bin that takes `-f`/`-p`; the package's `aimock` bin takes `--config` instead. const AIMOCK_BIN = 'llmock'; -// Deliberately below `ip_local_port_range` (32768-60999 here and on GitHub runners): inside it an -// unrelated outbound connection can hold the number as its source port, which the pre-flight check -// below would not see (a client socket does not accept) and `listen` would then fail with EADDRINUSE. +// Below typical ephemeral port range so a stray outbound source port cannot fake the pre-flight check. const MOCK_OPENAI_PORT = 18_937; /** - * Points the extension host at the mock server instead of the real OpenAI API. - * - * MUST be called at a spec file's module scope, never from `before`. ExTester launches VS Code from a - * root `beforeAll` (`vscode-extension-tester/out/suite/runner.js`) and the extension host inherits its - * environment at spawn time, so a hook runs too late — while Mocha loads spec files before it runs any - * hook, which is what makes module scope early enough. - * - * `rootHooks.ts` would be the tidier home, but ExTester builds Mocha through `new Mocha(config)`, and - * the programmatic API ignores the `require` option that file is wired up with. + * Set `OPENAI_BASE_URL` for the extension host. Call at spec module scope — ExTester spawns VS Code + * before Mocha `before` hooks, and the host inherits env at spawn time. */ export function pointExtensionHostAtMockServer(): void { process.env.OPENAI_BASE_URL = `http://127.0.0.1:${MOCK_OPENAI_PORT}/v1`; } -// `npm run setup:e2e:mock` primes `~/.npm/_npx` with this exact spec, so a warm start resolves from -// cache without a registry round-trip. The ceiling still covers a cold fetch: the setup step is not -// enforced, and if the two specs ever drift the run silently falls back to downloading here. const START_TIMEOUT = 90_000; const POLL_INTERVAL = 200; - -// How long to wait for the tree to go down after each of SIGTERM and SIGKILL. Short because the -// graceful signal is not what we rely on: `server.close()` releases the listening socket before the -// process is gone, so a freed port arrives long before the shutdown it appears to signal. const STOP_TIMEOUT = 2_000; export interface MockOpenAiServer { - /** Stops the server and removes its fixtures. Idempotent; safe to call more than once. */ stop: () => Promise; } @@ -56,19 +34,9 @@ export interface MockToolCall { name: string; } -/** - * Which request a scripted leg answers. Both alternatives are predicates over the request's own - * messages, with no server-side counter — unlike aimock's `sequenceIndex`, which would run past the - * end of the script on a Mocha retry (`.mocharc.js` sets `retries: 1` and `before` does not re-run - * between attempts) and fail the retry for a different reason than the original. - * - * `toolResultContains` additionally requires the last message to be a tool result, so it is what - * proves a round-trip: the leg is only reachable if the extension really ran the previous tool and - * fed its real output back. - */ +/** Predicate per scripted leg (not sequence index — safe across Mocha retries). */ export type MockAgentMatch = { hasToolResult: false } | { toolResultContains: string }; -/** What the agent gets back: another tool call, or the final text that ends the loop. */ export type MockAgentResponse = { content: string } | { toolCall: MockToolCall }; export interface MockAgentTurn { @@ -89,14 +57,6 @@ function canConnect(port: number): Promise { }); } -/** - * Fails the run when `OPENAI_BASE_URL` does not point at this server. - * - * Silence here is not a failed test: `executeAgentBlock` reads the variable at call time and falls - * back to `openai(model)` against the real api.openai.com (@deepnote/runtime-core dist/index.js:102), - * so an unset or drifted value sends the suite's prompts — and whatever key is in SecretStorage — to - * the live API. - */ function assertBaseUrlPointsAtMock(): void { if (!process.env.OPENAI_BASE_URL?.includes(`:${MOCK_OPENAI_PORT}`)) { throw new Error( @@ -107,7 +67,6 @@ function assertBaseUrlPointsAtMock(): void { } } -/** Writes `turns` as an aimock fixtures file in a fresh temp directory, and returns that directory. */ function writeFixtures(turns: MockAgentTurn[]): string { const fixtures = turns.map(({ match, response }) => ({ match, @@ -115,26 +74,14 @@ function writeFixtures(turns: MockAgentTurn[]): string { })); const directory = fs.mkdtempSync(path.join(os.tmpdir(), 'deepnote-e2e-aimock-')); - // The `fixtures` wrapper is required — aimock's loader rejects a bare array. fs.writeFileSync(path.join(directory, 'fixtures.json'), JSON.stringify({ fixtures }, undefined, 4)); return directory; } -/** - * Starts aimock on the mock port, scripted with `turns` — each answering the request its `match` - * describes. Resolves once the port accepts connections, which the CLI only reaches after loading and - * validating the fixtures, so a served request can never race an unloaded fixture. - * - * Runs with `--strict`, so a request matching no leg is answered with an error rather than a default: - * a broken round-trip fails loudly instead of quietly taking a different path through the script. - */ export async function startMockOpenAiServer(turns: MockAgentTurn[]): Promise { assertBaseUrlPointsAtMock(); - // Without this the readiness poll below cannot tell our server from someone else's: a leftover - // from a crashed run would satisfy it instantly, and the suite would then be asserting against - // that server's fixtures while ours had already died of EADDRINUSE. if (await canConnect(MOCK_OPENAI_PORT)) { throw new Error( `Port ${MOCK_OPENAI_PORT} is already in use — most likely a mock server left behind by an ` + @@ -147,8 +94,6 @@ export async function startMockOpenAiServer(turns: MockAgentTurn[]): Promise `sh -c` -> node), and a signal sent - // to npx alone leaves that grandchild holding the port. Its own process group makes the - // whole tree signalable; see `signalTree`. detached: true, stdio: ['ignore', 'inherit', 'inherit'] } @@ -177,10 +117,6 @@ export async function startMockOpenAiServer(turns: MockAgentTurn[]): Promise { exitReason = `code ${code}, signal ${signal}`; }); - // Node emits 'error' rather than 'exit' when the spawn itself fails (ENOENT for a missing npx, - // EACCES, …). An unhandled 'error' on a ChildProcess throws out of the event loop and takes the - // whole mocha process with it, losing every later suite and skipping ExTester's teardown; routing - // it through exitReason turns that into the readiness loop's ordinary failure. child.once('error', (error) => { exitReason = `spawn failed: ${error.message}`; }); @@ -193,11 +129,10 @@ export async function startMockOpenAiServer(turns: MockAgentTurn[]): Promise signalTree('SIGKILL'); process.once('exit', killChild); @@ -221,11 +156,6 @@ export async function startMockOpenAiServer(turns: MockAgentTurn[]): Promise { fs.rmSync(fixturesDirectory, { force: true, recursive: true }); - // Neither half of `hasShutDown` proves the node server is gone on its own: `server.close()` - // frees the listening socket while still draining open connections, and npx exits ahead of - // the server it spawned. So give SIGTERM a graceful window, then SIGKILL the group - // unconditionally — on an already-dead group that is a swallowed ESRCH, and it is the only - // step that guarantees nothing is left behind holding the port. signalTree('SIGTERM'); await waitForShutdown(); signalTree('SIGKILL'); @@ -237,7 +167,6 @@ export async function startMockOpenAiServer(turns: MockAgentTurn[]): Promise { ); } -/** - * Runs `read` inside the notebook webview (iframe.webview.ready -> #active-frame) and switches back - * afterwards. `read` only ever sees the webview, never the cell source in the main document — the - * guarantee callers rely on to avoid matching a cell's own text. Returns '' when the frame is absent, - * went stale, or has painted nothing yet, so callers can poll. - */ +/** Runs `read` inside the notebook output webview; returns '' when the frame is missing or not ready. */ async function readInsideNotebookWebview(read: (webView: WebView) => Promise): Promise { const driver = VSBrowser.instance.driver; const webView = new WebView(); @@ -62,9 +57,6 @@ async function readInsideNotebookWebview(read: (webView: WebView) => Promise('return window.self === window.top')) { return ''; } @@ -81,11 +73,7 @@ async function readInsideNotebookWebview(read: (webView: WebView) => Promise { return readInsideNotebookWebview(async (webView) => (await webView.findWebElement(By.css('body'))).getText()); } @@ -105,7 +93,6 @@ export async function readRenderedOutput(): Promise { ); const text = texts.join('\n').trim(); - // Safe as a fallback because we have already confirmed we are inside the webview, not the editor. return text || (await webView.findWebElement(By.css('body'))).getText(); }); } diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts index a01ee9a637..f34f483fed 100644 --- a/test/e2e/suite/agentBlock.e2e.test.ts +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -1,18 +1,9 @@ /** - * E2E (ExTester): one agent block driving a three-leg tool loop against a stand-in OpenAI API, so no - * network call is made. The agent asks for a code block, the extension inserts it as an ephemeral - * cell and runs it on the kernel, and the real stdout goes back as the tool result; the agent then - * asks for a markdown block and finally answers. - * - * The scripted legs 2 and 3 match on `toolResultContains`, so the agent can only advance if the - * extension genuinely executed the generated Python and returned its actual output. With aimock's - * `--strict`, a broken round-trip matches no leg and fails loudly. - * - * Executing generated code needs a real kernel: the first run provisions a venv and installs the - * Deepnote toolkit, which takes minutes. + * Agent block E2E: three-leg tool loop against a local aimock server (no live OpenAI calls). + * Legs 2–3 match on tool results, so the mock only advances after real kernel stdout and + * markdown tool replies. First kernel run can take minutes (venv + toolkit). */ -import { expect } from 'chai'; import { EditorView, InputBox, VSBrowser, WebView, Workbench } from 'vscode-extension-tester'; import { @@ -37,69 +28,48 @@ import { waitForNotification } from '../helpers'; -// At module scope on purpose — VS Code is already running by the time `before` executes, and it -// inherits this at spawn time. See the function's contract. pointExtensionHostAtMockServer(); const AGENT_FILE = 'agent-block.deepnote'; const CODE_TOOL_NAME = 'add_code_block'; const MARKDOWN_TOOL_NAME = 'add_markdown_block'; - -// The extension's tool result for a successful add_markdown_block (agentCellExecutionHandler.ts); -// leg 3 keys off it, so the wording is a coupling to that constant. +// Coupled to agentCellExecutionHandler tool result for add_markdown_block (leg 3 match). const MARKDOWN_BLOCK_ADDED_TEXT = 'Markdown block added.'; - -// A stable name: createEnvironment treats "already exists" as success, so a leftover environment from -// a previous or retried run is reused rather than colliding — and its provisioned venv with it. const ENVIRONMENT_NAME = 'E2E Agent Env'; - -// Once the kernel is up the agent itself talks only to the local mock, so it is bounded by UI and -// extension-host latency. The kernel's own first run is bounded by FIRST_RUN_OUTPUT_TIMEOUT instead. const AGENT_RUN_TIMEOUT = 60_000; - -// Printed by the Python the agent asks for. Only the executed ephemeral code cell can put it in the -// webview: the webview renders outputs and markdown previews, never cell source, and the agent's own -// transcript reports tool output by length rather than by content. const PYTHON_OUTPUT_MARKER = 'agent-generated-python-ran'; const GENERATED_PYTHON = `print("${PYTHON_OUTPUT_MARKER}")`; - -// Reaches the notebook only through the agent's tool call — the streamed transcript never echoes -// tool arguments — so seeing it rendered is what proves an ephemeral markdown cell was inserted. const EPHEMERAL_MARKDOWN_TEXT = 'Ephemeral markdown written by the E2E agent run'; const FINAL_AGENT_TEXT = 'Summary added as a markdown block.'; - -// The mock server ignores credentials, but the extension refuses to start an agent run without a -// stored key (and would otherwise block on an input box mid-execution). const MOCK_API_KEY = 'sk-e2e-mock-key'; -// Exact palette label matters: `Workbench.executeCommand` silently runs the first palette entry on a -// mismatch. const SET_API_KEY_COMMAND = 'Deepnote: Set OpenAI API Key'; const CLEAR_API_KEY_COMMAND = 'Deepnote: Clear OpenAI API Key'; const REVERT_FILE_COMMAND = 'File: Revert File'; -// VS Code's save prompt; the bundle stores it with a mnemonic marker ("Do&&n't Save") that is -// stripped before rendering, so the button's text is this. const DISCARD_CHANGES_BUTTON = "Don't Save"; const API_KEY_SAVED_NOTIFICATION = /OpenAI API key has been saved/; -/** Polls the notebook webview until every marker is present, returning whatever it last read. */ -async function awaitWebviewMarkers(markers: string[], timeout: number): Promise { +async function awaitWebviewMarkers(markers: string[], timeout: number, context: string): Promise { const driver = VSBrowser.instance.driver; const deadline = Date.now() + timeout; let text = ''; while (Date.now() < deadline) { text = await readNotebookWebviewText(); - if (markers.every((marker) => text.includes(marker))) { + const missing = markers.filter((marker) => !text.includes(marker)); + if (missing.length === 0) { return text; } await driver.sleep(OUTPUT_POLL_INTERVAL); } - return text; + const missing = markers.filter((marker) => !text.includes(marker)); + throw new Error( + `Timed out after ${timeout}ms waiting for notebook webview (${context}). Missing: ${JSON.stringify(missing)}. ` + + `Last text: ${JSON.stringify(text)}` + ); } -/** Stores the throwaway key in SecretStorage so the agent run never opens the key prompt. */ async function storeMockOpenAiApiKey(): Promise { await new Workbench().executeCommand(SET_API_KEY_COMMAND); @@ -107,9 +77,6 @@ async function storeMockOpenAiApiKey(): Promise { await input.setText(MOCK_API_KEY); await input.confirm(); - // Confirm the key really landed. Had the palette missed the command, InputBox.create would have - // bound to the still-open palette and typed the key into it, and the suite would run keyless — - // surfacing a full AGENT_RUN_TIMEOUT later as a generic missing-marker failure. await waitForNotification(API_KEY_SAVED_NOTIFICATION, QUICK_PICK_TIMEOUT, true); } @@ -137,47 +104,25 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu `${AGENT_FILE} did not open` ); - // Binds a real kernel for the code the agent generates, replacing the "Select Environment" - // placeholder controller the auto-selector picks on open. This is also the settle signal it waits - // on: selectEnvironmentForNotebook returns after the post-binding "switched successfully" - // toast, so Run All is not racing the auto-selection. await createEnvironment(ENVIRONMENT_NAME); await selectEnvironmentForNotebook(ENVIRONMENT_NAME, AGENT_FILE); - // Toasts steal focus from the command palette. Safe to do after the environment flow, which - // has already driven extension commands and so guarantees `onNotebook:deepnote` activation. await dismissAllNotifications(); await storeMockOpenAiApiKey(); await screenshot('kernel-connected'); }); - /** - * Releases the server the running test started, if any. - * - * Runs on both sides of the test rather than only after it. `.mocharc.js` sets `retries: 1` and - * `before`/`after` do not run between attempts, so a server surviving a failed attempt would still - * hold the port when the retry starts — and `startMockOpenAiServer`'s pre-flight check would then - * reject it as a leftover, failing the retry for a different reason than the original and losing - * the real signal. `afterEach` normally prevents that; `beforeEach` covers the case where it was - * itself interrupted. - */ async function releaseMockServer(): Promise { await mockServer?.stop().catch((error) => { console.warn('[agent-block] stop the mock OpenAI server:', error); }); - // Cleared so a later release cannot stop an already-dead handle. mockServer = undefined; } beforeEach(releaseMockServer); - afterEach(releaseMockServer); after(async function () { - // Filesystem cleanup goes FIRST. The UI steps below can each stall for minutes — this is the - // one suite that ends with a dirty notebook, so `closeAllEditors` retries against a save modal - // and the `.catch()`-swallowed revert cannot stop it — and a wedged UI step would otherwise - // burn SUITE_TIMEOUT with the temp dir never released. try { cleanupTempDir?.(); } catch (error) { @@ -187,8 +132,6 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu await new WebView().switchBack().catch((error) => { console.warn('[agent-block] switch back from webview during cleanup:', error); }); - // The inserted ephemeral cell leaves the notebook dirty, and the resulting modal save prompt - // outlives this suite and blocks the next one in the shared VS Code instance. await new Workbench().executeCommand(REVERT_FILE_COMMAND).catch((error) => { console.warn('[agent-block] revert notebook during cleanup:', error); }); @@ -196,30 +139,18 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu console.warn('[agent-block] close all editors during cleanup:', error); }); - // Backstop for a revert that did not land: an unanswered save modal blocks the next suite. - // Gated on an editor surviving the close, because confirmModalDialog polls for the full - // WORKBENCH_TIMEOUT when no dialog is up — dead time on every green run otherwise. const openEditors = await new EditorView().getOpenEditorTitles().catch(() => [] as string[]); if (openEditors.length > 0) { await confirmModalDialog(DISCARD_CHANGES_BUTTON).catch((error) => { console.warn('[agent-block] discard unsaved changes during cleanup:', error); }); } - // SecretStorage outlives this suite in the shared VS Code instance, so leave no key behind. await new Workbench().executeCommand(CLEAR_API_KEY_COMMAND).catch((error) => { console.warn('[agent-block] clear the stored OpenAI API key during cleanup:', error); }); }); - // Known limitation of the Mocha retry (`.mocharc.js` sets `retries: 1`): if this times out with an - // execution still in flight, the retry's clickRunAll may find Interrupt where Run All was and fail - // for an unrelated reason, losing the original signal. Only reachable on an already-failing test, - // so it costs debuggability rather than correctness. it('executes the code block the agent generates, then inserts its markdown block', async function () { - // Scripted here because the conversation is what a given test is about — a different test - // scripts different legs. Leg 2 and leg 3 are reachable only via what the extension sends - // back, so the script itself asserts the round-trip: leg 2 needs the kernel's real stdout, - // leg 3 the markdown tool's reply. mockServer = await startMockOpenAiServer([ { match: { hasToolResult: false }, @@ -247,42 +178,22 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu } ]); - // Clears the "OpenAI API key has been saved." and "switched successfully" toasts, which would - // otherwise intercept the toolbar click. await dismissAllNotifications(); await clickRunAll(AGENT_FILE); - // Every marker is asserted below, so poll for all of them — a missing one then fails on its - // own assertion rather than on whichever runs first. Split in two waits because the stages - // have very different budgets: the generated cell is the first thing to touch the kernel, and - // that first execution carries the connect cost, while the rest is local. Waiting on the - // Python marker first also reports a kernel failure as a kernel failure rather than as a - // missing agent marker. - await awaitWebviewMarkers([PYTHON_OUTPUT_MARKER], FIRST_RUN_OUTPUT_TIMEOUT); - - const agentMarkers = [ - `[Agent] Tool called: ${CODE_TOOL_NAME}`, - `[Agent] Tool called: ${MARKDOWN_TOOL_NAME}`, - EPHEMERAL_MARKDOWN_TEXT, - FINAL_AGENT_TEXT - ]; - const webviewText = await awaitWebviewMarkers(agentMarkers, AGENT_RUN_TIMEOUT); + await awaitWebviewMarkers([PYTHON_OUTPUT_MARKER], FIRST_RUN_OUTPUT_TIMEOUT, 'generated code cell stdout'); + + await awaitWebviewMarkers( + [ + `[Agent] Tool called: ${CODE_TOOL_NAME}`, + `[Agent] Tool called: ${MARKDOWN_TOOL_NAME}`, + EPHEMERAL_MARKDOWN_TEXT, + FINAL_AGENT_TEXT + ], + AGENT_RUN_TIMEOUT, + 'agent tool loop and ephemeral markdown' + ); await screenshot('agent-run'); - - expect(webviewText, 'the agent cell did not stream its add_code_block call into the cell output').to.contain( - `[Agent] Tool called: ${CODE_TOOL_NAME}` - ); - expect(webviewText, 'the generated code cell did not run on the kernel').to.contain(PYTHON_OUTPUT_MARKER); - expect( - webviewText, - 'the agent cell did not stream its add_markdown_block call into the cell output' - ).to.contain(`[Agent] Tool called: ${MARKDOWN_TOOL_NAME}`); - expect(webviewText, 'the tool call did not insert an ephemeral markdown cell').to.contain( - EPHEMERAL_MARKDOWN_TEXT - ); - expect(webviewText, "the agent's final message was not streamed into the cell output").to.contain( - FINAL_AGENT_TEXT - ); }); }); From 53398c49deaf3bd3a255576ffc7158bbe653c46b Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 5 Aug 2026 14:20:33 +0000 Subject: [PATCH 35/37] refactor(agent-block): split OpenAI key commands from status bar Move palette commands into AgentOpenAiApiKeyCommandHandler, teach the agent execution test harness to apply delete notebook edits, and drop the misplaced createMockNotebook suite. Co-authored-by: Cursor --- .../agentCellExecutionHandler.unit.test.ts | 32 ++++++++----------- .../deepnote/agentCellStatusBarProvider.ts | 17 ---------- .../agentOpenAiApiKeyCommandHandler.ts | 30 +++++++++++++++++ src/notebooks/serviceRegistry.node.ts | 5 +++ src/notebooks/serviceRegistry.web.ts | 5 +++ 5 files changed, 53 insertions(+), 36 deletions(-) create mode 100644 src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index 3c3dbd8d9c..f8e70c30df 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -15,7 +15,6 @@ import { NotebookDocument, SecretStorage, SecretStorageChangeEvent, - Uri, WorkspaceEdit } from 'vscode'; @@ -74,7 +73,7 @@ function stubSecretStorage(secretStorage: Map): ServiceContainer * tests, so its own call counts are useless here. */ function applyNotebookEditsTo(cells: NotebookCell[], notebook: NotebookDocument) { - type RecordedEdit = { range: { start: number; end: number }; newCells: NotebookCellData[] }; + type RecordedEdit = { range: { start: number; end: number }; newCells?: NotebookCellData[] }; let recordedEdits: RecordedEdit[] = []; let appliedEdits = 0; @@ -87,7 +86,17 @@ function applyNotebookEditsTo(cells: NotebookCell[], notebook: NotebookDocument) for (const notebookEdit of recordedEdits) { const { start, end } = notebookEdit.range; - const inserted = notebookEdit.newCells.map((cellData) => { + const deleteCount = end - start; + const newCellData = notebookEdit.newCells; + + if (!newCellData || newCellData.length === 0) { + if (deleteCount > 0) { + cells.splice(start, deleteCount); + } + continue; + } + + const inserted = newCellData.map((cellData) => { const created = createMockCell({ text: cellData.value, metadata: cellData.metadata @@ -97,7 +106,7 @@ function applyNotebookEditsTo(cells: NotebookCell[], notebook: NotebookDocument) return created; }); - cells.splice(start, end - start, ...inserted); + cells.splice(start, deleteCount, ...inserted); } cells.forEach((cell, index) => ((cell as { index: number }).index = index)); recordedEdits = []; @@ -687,18 +696,3 @@ suite('AgentCellExecutionHandler', () => { }); }); }); - -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 7e36a1aca0..6cbb2797f5 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -17,7 +17,6 @@ import { injectable } from 'inversify'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { isAgentCell } from './dataConversionUtils'; -import { clearOpenAiApiKey, promptForOpenAiApiKey } from './deepnoteSecretStore'; /** The key `agentBlockSchema` defines and `executeAgentBlock` reads. */ const AGENT_MODEL_METADATA_KEY = 'deepnote_agent_model'; @@ -57,22 +56,6 @@ 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/agentOpenAiApiKeyCommandHandler.ts b/src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts new file mode 100644 index 0000000000..09ada758df --- /dev/null +++ b/src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts @@ -0,0 +1,30 @@ +import { inject, injectable } from 'inversify'; +import { commands, l10n, window } from 'vscode'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { IExtensionContext } from '../../platform/common/types'; +import { clearOpenAiApiKey, promptForOpenAiApiKey } from './deepnoteSecretStore'; + +@injectable() +export class AgentOpenAiApiKeyCommandHandler implements IExtensionSyncActivationService { + constructor(@inject(IExtensionContext) private readonly extensionContext: IExtensionContext) {} + + public activate(): void { + this.extensionContext.subscriptions.push( + commands.registerCommand('deepnote.setOpenAiApiKey', () => this.setApiKey()), + commands.registerCommand('deepnote.clearOpenAiApiKey', () => this.clearApiKey()) + ); + } + + private async setApiKey(): Promise { + const key = await promptForOpenAiApiKey(); + if (key) { + void window.showInformationMessage(l10n.t('OpenAI API key has been saved.')); + } + } + + private async clearApiKey(): Promise { + await clearOpenAiApiKey(); + void window.showInformationMessage(l10n.t('OpenAI API key has been cleared.')); + } +} diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index ceadef4d78..a70e84ac1c 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -95,6 +95,7 @@ import { DeepnoteNotebookEnvironmentMapper } from '../kernels/deepnote/environme import { DeepnoteNotebookCommandListener } from './deepnote/deepnoteNotebookCommandListener'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; import { AgentCellStatusBarProvider } from './deepnote/agentCellStatusBarProvider'; +import { AgentOpenAiApiKeyCommandHandler } from './deepnote/agentOpenAiApiKeyCommandHandler'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; import { EphemeralCellDecorationProvider } from './deepnote/ephemeralCellDecorationProvider'; import { EphemeralCellStatusBarProvider } from './deepnote/ephemeralCellStatusBarProvider'; @@ -264,6 +265,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentOpenAiApiKeyCommandHandler + ); serviceManager.addSingleton( IExtensionSyncActivationService, AgentCellStatusBarProvider diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index 1ca01e8512..8e7b5e665a 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/src/notebooks/serviceRegistry.web.ts @@ -51,6 +51,7 @@ import { } from './deepnote/integrations/types'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; import { AgentCellStatusBarProvider } from './deepnote/agentCellStatusBarProvider'; +import { AgentOpenAiApiKeyCommandHandler } from './deepnote/agentOpenAiApiKeyCommandHandler'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; import { EphemeralCellDecorationProvider } from './deepnote/ephemeralCellDecorationProvider'; import { EphemeralCellStatusBarProvider } from './deepnote/ephemeralCellStatusBarProvider'; @@ -130,6 +131,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentOpenAiApiKeyCommandHandler + ); serviceManager.addSingleton( IExtensionSyncActivationService, AgentCellStatusBarProvider From d6338dc5a0c161971266c83bd4a6a416d9a88330 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 5 Aug 2026 18:39:59 +0000 Subject: [PATCH 36/37] fix(agent-block): keep ephemeral cells when the watcher reads back our own save MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit serializeNotebook filters ephemeral cells out, so an ordinary Ctrl+S or Auto Save writes a file with fewer cells than the live document. contentActuallyChanged compared raw cell counts, read that difference as an external edit, and executeMainFileSync replaced the whole document from disk — deleting the agent's generated cells about half a second after the user saved, or mid-run under Auto Save. Compare against the same view the serializer persists. Preferred over marking serializer writes as self-writes: onDidSaveNotebookDocument fires after the write and races the fs event, whereas the content comparison is deterministic and is already the documented guard for saves. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- .../deepnote/deepnoteFileChangeWatcher.ts | 6 ++-- .../deepnoteFileChangeWatcher.unit.test.ts | 34 +++++++++++++++++++ 2 files changed, 38 insertions(+), 2 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index 0a82635870..159b4b4b7c 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -18,7 +18,7 @@ import { IExtensionSyncActivationService } from '../../platform/activation/types import { IDisposableRegistry } from '../../platform/common/types'; import { logger } from '../../platform/logging'; import { IDeepnoteNotebookManager } from '../types'; -import { getBlockId } from './dataConversionUtils'; +import { getBlockId, isEphemeralCell } from './dataConversionUtils'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; import { DeepnoteNotebookSerializer } from './deepnoteSerializer'; @@ -164,7 +164,9 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic * has fewer/no outputs), it's an auto-save of stripped content — skip reload. */ private contentActuallyChanged(notebook: NotebookDocument, newCells: NotebookCellData[]): boolean { - const liveCells = notebook.getCells(); + // Compare against what the serializer persists: ephemeral cells are never written, so + // counting them reads our own save back as an external edit that deletes them. + const liveCells = notebook.getCells().filter((cell) => !isEphemeralCell(cell)); if (liveCells.length !== newCells.length) { return true; } diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index cd9eb4d818..2a95a82051 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -178,6 +178,40 @@ project: assert.strictEqual(applyEditCount, 0, 'applyEdit should not be called when cells match'); }); + test('should skip reload when the live notebook only adds ephemeral cells', async () => { + const uri = Uri.file('/workspace/test.deepnote'); + // The serializer never persists ephemeral cells, so a plain save produces a file with + // fewer cells than the live document. That difference must not read as an external edit. + const notebook = createMockNotebook({ + uri, + cells: [ + { + metadata: { id: 'block-1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")', languageId: 'python' } + }, + { + metadata: { id: 'eph-1', is_ephemeral: true, agent_source_block_id: 'agent-1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("agent generated")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(validYaml); + + onDidChangeFile.fire(uri); + + await waitFor(() => readFileCalls > 0); + await new Promise((resolve) => setTimeout(resolve, autoSaveGraceMs)); + + assert.strictEqual(applyEditCount, 0, 'ephemeral-only difference should not trigger a reload'); + assert.strictEqual(saveCount, 0, 'ephemeral-only difference should not trigger a save'); + }); + test('should reload on external change', async () => { const uri = Uri.file('/workspace/test.deepnote'); const notebook = createMockNotebook({ uri, cellCount: 0 }); From c9f7c60bd923b362f9c4e45ffa12d96aa3b9100d Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 5 Aug 2026 20:43:49 +0000 Subject: [PATCH 37/37] test(agent-block): pin the streamed transcript rendering in E2E MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The agent transcript is streamed as one appended stdout item per event rather than re-sent whole, which only pays off if the renderer joins those items back into a single block. That was reasoned from the API and left as a manual check: the unit tests assert against a stubbed appendOutputItems, so they cannot see what the renderer does with the items. The agent E2E run now keeps the rendered webview text and asserts two substrings that each span several appended items — the blank line between two tool sections, and the final answer, lengthened so aimock's 20-character chunking splits it across four text_delta events. Replaying the same items joined per-item instead of concatenated fails both, so a fragmenting renderer is caught. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- .../deepnote/agentCellExecutionHandler.ts | 3 +- test/e2e/suite/agentBlock.e2e.test.ts | 29 +++++++++++++++++-- 2 files changed, 28 insertions(+), 4 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index 81f54d77b0..a91c58c348 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -196,7 +196,8 @@ export async function executeAgentCell( // 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. + // Appended stdout items render as one continuous block, the same way kernel output streams; + // agentBlock.e2e.test.ts pins that, since only a real renderer can show it. const output = new NotebookCellOutput([NotebookCellOutputItem.stdout(`[Agent] Planning next steps...`)]); await execution.replaceOutput([output]); diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts index f34f483fed..6f60263260 100644 --- a/test/e2e/suite/agentBlock.e2e.test.ts +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -40,7 +40,9 @@ const AGENT_RUN_TIMEOUT = 60_000; const PYTHON_OUTPUT_MARKER = 'agent-generated-python-ran'; const GENERATED_PYTHON = `print("${PYTHON_OUTPUT_MARKER}")`; const EPHEMERAL_MARKDOWN_TEXT = 'Ephemeral markdown written by the E2E agent run'; -const FINAL_AGENT_TEXT = 'Summary added as a markdown block.'; +// aimock streams content in 20-character chunks, so an answer this long reaches the agent cell as +// several text_delta events — one appended output item each. +const FINAL_AGENT_TEXT = 'Summary added as a markdown block, streamed across several deltas.'; const MOCK_API_KEY = 'sk-e2e-mock-key'; const SET_API_KEY_COMMAND = 'Deepnote: Set OpenAI API Key'; const CLEAR_API_KEY_COMMAND = 'Deepnote: Clear OpenAI API Key'; @@ -70,6 +72,17 @@ async function awaitWebviewMarkers(markers: string[], timeout: number, context: ); } +function assertRenderedContiguously(transcript: string, expected: string): void { + if (transcript.includes(expected)) { + return; + } + + throw new Error( + `Agent transcript does not contain ${JSON.stringify(expected)} as one unbroken run — appended stdout ` + + `items are not rendering as a single block. Full transcript: ${JSON.stringify(transcript)}` + ); +} + async function storeMockOpenAiApiKey(): Promise { await new Workbench().executeCommand(SET_API_KEY_COMMAND); @@ -150,7 +163,7 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu }); }); - it('executes the code block the agent generates, then inserts its markdown block', async function () { + it('executes the generated code block, inserts its markdown block, and streams one transcript', async function () { mockServer = await startMockOpenAiServer([ { match: { hasToolResult: false }, @@ -183,7 +196,7 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu await awaitWebviewMarkers([PYTHON_OUTPUT_MARKER], FIRST_RUN_OUTPUT_TIMEOUT, 'generated code cell stdout'); - await awaitWebviewMarkers( + const transcript = await awaitWebviewMarkers( [ `[Agent] Tool called: ${CODE_TOOL_NAME}`, `[Agent] Tool called: ${MARKDOWN_TOOL_NAME}`, @@ -195,5 +208,15 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu ); await screenshot('agent-run'); + + // agentCellExecutionHandler appends one stdout item per agent event instead of re-sending the + // transcript, which only pays off if the renderer joins those items back into one block — + // something only a real renderer can show: a section boundary keeps its blank line, and an + // answer split across several deltas arrives unbroken. + assertRenderedContiguously( + transcript, + `[Agent] Tool called: ${MARKDOWN_TOOL_NAME}\n\n[Agent] Tool output: ${MARKDOWN_TOOL_NAME}` + ); + assertRenderedContiguously(transcript, `[Agent] Text:\n${FINAL_AGENT_TEXT}`); }); });