From ffe09588ac299d7ad88f852d707ee91eb5041181 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 9 Mar 2026 08:30:26 +0000 Subject: [PATCH 01/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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/80] 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}`); }); }); From b22c79c34c87bf0e488941ab2043b1e52f673cbd Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 6 Aug 2026 06:09:13 +0000 Subject: [PATCH 38/80] test(agent-block): cover the re-run that clears the previous run's cells MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Running an agent block twice is the only state in which the batch-level clearing and the precondition executeAgentCell verifies do anything, and nothing outside the unit tests exercised it. The first spec leaves the block owning a generated code cell and a markdown cell, so a second Run All over that notebook is what the fix in 84c3cedfe has to survive in the real product. The new spec re-runs with a different script and counts occurrences rather than asserting presence — on a Mocha retry the markers repeat, so only the count separates one copy from two — and checks the precondition never surfaces as stderr on the agent cell. Seen to fail: with removeEphemeralCellsForAgentBlocks dropped from executeQueuedCells, the spec times out waiting for the second run's markers and the notebook reads "still has 2 generated cell(s) from its previous run" next to the first run's stdout, still rendered. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- test/e2e/suite/agentBlock.e2e.test.ts | 77 +++++++++++++++++++++++++++ 1 file changed, 77 insertions(+) diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts index 6f60263260..7bc61545bc 100644 --- a/test/e2e/suite/agentBlock.e2e.test.ts +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -2,6 +2,7 @@ * 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). + * The specs run in order: the second re-runs the same block over the first one's generated cells. */ import { EditorView, InputBox, VSBrowser, WebView, Workbench } from 'vscode-extension-tester'; @@ -43,6 +44,13 @@ const EPHEMERAL_MARKDOWN_TEXT = 'Ephemeral markdown written by the E2E agent run // 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.'; +// The re-run's own markers, chosen so neither run's text is a substring of the other's. +const RERUN_PYTHON_OUTPUT_MARKER = 'rerun-python-ran'; +const RERUN_GENERATED_PYTHON = `print("${RERUN_PYTHON_OUTPUT_MARKER}")`; +const RERUN_MARKDOWN_TEXT = 'Second-run markdown from the E2E agent'; +const RERUN_FINAL_AGENT_TEXT = 'Re-run summary added as a markdown block.'; +// Coupled to the precondition executeAgentCell reports when a previous run was left in place. +const STALE_CELLS_ERROR_TEXT = 'from its previous run'; 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'; @@ -83,6 +91,19 @@ function assertRenderedContiguously(transcript: string, expected: string): void ); } +function assertOccurrences(rendered: string, needle: string, expected: number): void { + const actual = rendered.split(needle).length - 1; + + if (actual === expected) { + return; + } + + throw new Error( + `Expected ${expected} occurrence(s) of ${JSON.stringify(needle)} in the notebook, found ${actual}. ` + + `Full text: ${JSON.stringify(rendered)}` + ); +} + async function storeMockOpenAiApiKey(): Promise { await new Workbench().executeCommand(SET_API_KEY_COMMAND); @@ -219,4 +240,60 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu ); assertRenderedContiguously(transcript, `[Agent] Text:\n${FINAL_AGENT_TEXT}`); }); + + // Runs against the notebook the previous test left behind, so the agent block starts this run + // owning a full set of generated cells — the only state in which the batch-level clearing and the + // precondition executeAgentCell verifies can be observed at all. + it('clears the cells its previous run generated instead of stacking a second copy', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: RERUN_GENERATED_PYTHON }), + id: 'call_e2e_rerun_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: RERUN_PYTHON_OUTPUT_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: RERUN_MARKDOWN_TEXT }), + id: 'call_e2e_rerun_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: RERUN_FINAL_AGENT_TEXT } + } + ]); + + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); + + const rendered = await awaitWebviewMarkers( + [RERUN_PYTHON_OUTPUT_MARKER, RERUN_MARKDOWN_TEXT, RERUN_FINAL_AGENT_TEXT], + AGENT_RUN_TIMEOUT, + 'second agent run' + ); + + await screenshot('agent-rerun'); + + // The first run's cells are gone rather than pushed down by the second run's, which is also + // what keeps them out of the notebook context the agent is handed. Counting rather than + // asserting presence: on a Mocha retry the markers are the same, so only the count separates + // one copy from two. + assertOccurrences(rendered, PYTHON_OUTPUT_MARKER, 0); + assertOccurrences(rendered, EPHEMERAL_MARKDOWN_TEXT, 0); + assertOccurrences(rendered, RERUN_PYTHON_OUTPUT_MARKER, 1); + assertOccurrences(rendered, RERUN_MARKDOWN_TEXT, 1); + + // Clearing happens per batch, before anything runs, so the agent cell never reports the + // precondition — it would surface here as stderr on the agent cell. + assertOccurrences(rendered, STALE_CELLS_ERROR_TEXT, 0); + }); }); From 1101487347ac18aba54a7fec943067c8b1d4b42e Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 6 Aug 2026 06:22:49 +0000 Subject: [PATCH 39/80] refactor(agent-block): require cancellation token for ephemeral cell runs Call sites always pass a token; tighten the API and align unit tests and status bar helper typing. Co-authored-by: Cursor --- src/notebooks/deepnote/agentCellExecutionHandler.ts | 8 +++----- .../deepnote/agentCellExecutionHandler.unit.test.ts | 13 ++++++++++--- .../deepnote/ephemeralCellStatusBarProvider.ts | 4 ++-- 3 files changed, 15 insertions(+), 10 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index a91c58c348..d0e03b316a 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -403,11 +403,11 @@ export interface EphemeralCellExecutionResult { export async function executeEphemeralCell( cell: NotebookCell, - token?: CancellationToken + token: CancellationToken ): Promise { // Bail before dispatching: rejecting the deferred alone would abandon the wait but still hand the // generated code to the kernel. - if (token?.isCancellationRequested) { + if (token.isCancellationRequested) { throw new CancellationError(); } @@ -422,9 +422,7 @@ export async function executeEphemeralCell( }) ); - if (token) { - disposables.push(token.onCancellationRequested(() => completionDeferred.reject(new CancellationError()))); - } + disposables.push(token.onCancellationRequested(() => completionDeferred.reject(new CancellationError()))); const timeout = setTimeout(() => { completionDeferred.reject(new Error('Ephemeral cell execution timed out')); diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index f8e70c30df..185eba47bb 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -608,7 +608,14 @@ suite('AgentCellExecutionHandler', () => { }); suite('executeEphemeralCell', () => { + let tokenSource: CancellationTokenSource; + + setup(() => { + tokenSource = new CancellationTokenSource(); + }); + teardown(() => { + tokenSource.dispose(); reset(mockedVSCodeNamespaces.commands); }); @@ -625,7 +632,7 @@ suite('AgentCellExecutionHandler', () => { notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Idle); }); - await executeEphemeralCell(cell); + await executeEphemeralCell(cell, tokenSource.token); const [commandName, commandArg] = capture( mockedVSCodeNamespaces.commands.executeCommand as (cmd: string, arg: unknown) => Thenable @@ -666,7 +673,7 @@ suite('AgentCellExecutionHandler', () => { new Error('kernel is dead') ); - const result = await executeEphemeralCell(cell); + const result = await executeEphemeralCell(cell, tokenSource.token); expect(result.success).to.be.false; expect(result.error).to.equal('kernel is dead'); @@ -683,7 +690,7 @@ suite('AgentCellExecutionHandler', () => { () => new Promise(() => undefined) ); - const resultPromise = executeEphemeralCell(cell); + const resultPromise = executeEphemeralCell(cell, tokenSource.token); await clock.tickAsync(EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS); const result = await resultPromise; diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts index 12a7f08d0b..d3da9be862 100644 --- a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts @@ -57,10 +57,10 @@ export class EphemeralCellStatusBarProvider const agentSourceBlockId = cell.metadata?.agent_source_block_id as string | undefined; - return this.createEphemeralIndicatorItem(agentSourceBlockId); + return this.createEphemeralIndicatorItem(agentSourceBlockId ?? null); } - private createEphemeralIndicatorItem(agentSourceBlockId?: string): NotebookCellStatusBarItem { + private createEphemeralIndicatorItem(agentSourceBlockId: string | null): NotebookCellStatusBarItem { const tooltipLines = [l10n.t('Auto-generated ephemeral block')]; if (agentSourceBlockId) { tooltipLines.push(l10n.t('Source agent block: {0}', agentSourceBlockId)); From 22e1fe54e9306a3bd1a95c61414d1c0533c66b06 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 6 Aug 2026 08:30:07 +0000 Subject: [PATCH 40/80] docs(agent-block): cut the comments to the facts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drops the narration and the counterfactuals, and the stale claim that executeAgentCell has three call sites — it has one. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- .../deepnote/agentCellExecutionHandler.ts | 13 +++------- test/e2e/suite/agentBlock.e2e.test.ts | 24 ++++--------------- 2 files changed, 8 insertions(+), 29 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index d0e03b316a..e7870e2524 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -192,12 +192,8 @@ export async function executeAgentCell( const prompt = cell.document.getText(); - // 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. - // 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. + // Per-event appends keep this O(n) bytes across the extension-host boundary, on the run's wall + // clock: runtime-core awaits `onAgentEvent`. The renderer concatenates appended stdout items. const output = new NotebookCellOutput([NotebookCellOutputItem.stdout(`[Agent] Planning next steps...`)]); await execution.replaceOutput([output]); @@ -209,10 +205,7 @@ export async function executeAgentCell( throw new Error('Cell is not an agent cell'); } - // 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. + // The caller clears the previous run per batch. const staleCellCount = cell.notebook .getCells() .filter((c) => getEphemeralCellAgentSourceBlockId(c) === agentBlock.id).length; diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts index 7bc61545bc..a9aa076ba3 100644 --- a/test/e2e/suite/agentBlock.e2e.test.ts +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -2,7 +2,6 @@ * 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). - * The specs run in order: the second re-runs the same block over the first one's generated cells. */ import { EditorView, InputBox, VSBrowser, WebView, Workbench } from 'vscode-extension-tester'; @@ -41,15 +40,14 @@ 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'; -// 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. +// aimock chunks content at 20 characters, so this length arrives as several text_delta events. const FINAL_AGENT_TEXT = 'Summary added as a markdown block, streamed across several deltas.'; -// The re-run's own markers, chosen so neither run's text is a substring of the other's. +// Not substrings of the first run's markers; the counts below depend on that. const RERUN_PYTHON_OUTPUT_MARKER = 'rerun-python-ran'; const RERUN_GENERATED_PYTHON = `print("${RERUN_PYTHON_OUTPUT_MARKER}")`; const RERUN_MARKDOWN_TEXT = 'Second-run markdown from the E2E agent'; const RERUN_FINAL_AGENT_TEXT = 'Re-run summary added as a markdown block.'; -// Coupled to the precondition executeAgentCell reports when a previous run was left in place. +// Coupled to executeAgentCell's uncleared-previous-run error. const STALE_CELLS_ERROR_TEXT = 'from its previous run'; const MOCK_API_KEY = 'sk-e2e-mock-key'; const SET_API_KEY_COMMAND = 'Deepnote: Set OpenAI API Key'; @@ -230,10 +228,6 @@ 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}` @@ -241,9 +235,7 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu assertRenderedContiguously(transcript, `[Agent] Text:\n${FINAL_AGENT_TEXT}`); }); - // Runs against the notebook the previous test left behind, so the agent block starts this run - // owning a full set of generated cells — the only state in which the batch-level clearing and the - // precondition executeAgentCell verifies can be observed at all. + // Depends on the previous spec: the block starts this run owning the cells it generated there. it('clears the cells its previous run generated instead of stacking a second copy', async function () { mockServer = await startMockOpenAiServer([ { @@ -283,17 +275,11 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu await screenshot('agent-rerun'); - // The first run's cells are gone rather than pushed down by the second run's, which is also - // what keeps them out of the notebook context the agent is handed. Counting rather than - // asserting presence: on a Mocha retry the markers are the same, so only the count separates - // one copy from two. + // Counts, not presence: a Mocha retry repeats the markers. assertOccurrences(rendered, PYTHON_OUTPUT_MARKER, 0); assertOccurrences(rendered, EPHEMERAL_MARKDOWN_TEXT, 0); assertOccurrences(rendered, RERUN_PYTHON_OUTPUT_MARKER, 1); assertOccurrences(rendered, RERUN_MARKDOWN_TEXT, 1); - - // Clearing happens per batch, before anything runs, so the agent cell never reports the - // precondition — it would surface here as stderr on the agent cell. assertOccurrences(rendered, STALE_CELLS_ERROR_TEXT, 0); }); }); From c165b2b7f27b8ed407681b43ef4939292c879a2d Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 6 Aug 2026 14:04:29 +0000 Subject: [PATCH 41/80] docs(agent-block): trim comments to minimal why-only notes Align agent-block PR comments with senior-audience style: shorter public docstrings and fewer inline narrations without changing behavior. Co-authored-by: Cursor --- .../controllers/vscodeNotebookController.ts | 4 +- .../deepnote/agentCellExecutionHandler.ts | 77 ++++--------------- .../agentCellExecutionHandler.unit.test.ts | 52 +++---------- .../deepnote/agentCellStatusBarProvider.ts | 4 +- .../converters/agentBlockConverter.ts | 2 +- src/notebooks/deepnote/dataConversionUtils.ts | 26 ++----- .../deepnote/dataConversionUtils.unit.test.ts | 5 +- .../deepnote/deepnoteFileChangeWatcher.ts | 3 +- .../deepnoteFileChangeWatcher.unit.test.ts | 3 +- .../deepnoteKernelAutoSelector.node.ts | 2 +- src/notebooks/deepnote/deepnoteSerializer.ts | 1 - .../ephemeralCellDecorationProvider.ts | 2 +- src/renderers/client/markdown.ts | 2 +- src/test/mocks/deepnoteRuntimeCore.ts | 9 +-- test/e2e/.mocharc.js | 2 +- test/e2e/helpers/mockOpenAiServer.ts | 11 +-- test/e2e/helpers/notebook.ts | 6 +- test/e2e/suite/agentBlock.e2e.test.ts | 18 ++--- 18 files changed, 61 insertions(+), 168 deletions(-) diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index 7f61d0ab59..502a3b1009 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -627,7 +627,7 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont return; } const queuedCells = this.cellQueue.get(doc) || []; - // Clear before await so agent-driven re-entrant runs start with an empty queue. + // Clear before await — agent runs can re-enter with an empty queue. this.cellQueue.delete(doc); const cellsToExecute = await removeEphemeralCellsForAgentBlocks(doc, queuedCells); @@ -655,7 +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 }; - // Stale cell handles report index -1; createNotebookCellExecution would abort the batch. + // Stale handles use index -1 and abort the whole batch in createNotebookCellExecution. 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 e7870e2524..58d1d562a8 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -41,16 +41,7 @@ import { import { DeepnoteDataConverter } from './deepnoteDataConverter'; import { getOrPromptOpenAiApiKey } from './deepnoteSecretStore'; -/** - * 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 MCP servers is arbitrary local command execution declared by a workspace file, so every - * caller must already be behind a `workspace.isTrusted` check. - */ +/** Project MCP servers and integrations from the `.deepnote` file (CLI ExecutionEngine parity). Callers must gate on `workspace.isTrusted` — MCP spawn is arbitrary command execution. */ function getProjectAgentContext(notebook: NotebookDocument): Pick { const projectId = notebook.metadata?.deepnoteProjectId as string | undefined; const notebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; @@ -81,9 +72,7 @@ function getProjectAgentContext(notebook: NotebookDocument): Pick typeof entry === 'string') ? value.join('') : value; } -/** - * `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()` string representation doesn't reach the agent with a - * comma glued to the start of every line. - */ +/** Join nbformat line arrays before `extractOutputsText` — `String(array)` inserts commas between lines. */ function normalizeOutputsForTextExtraction(outputs: unknown[]): unknown[] { return outputs.map((output) => { const candidate = output as { output_type?: unknown; text?: unknown; data?: unknown } | null; @@ -171,12 +153,8 @@ export interface ExecuteAgentCellOptions { } /** - * 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 - * `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. + * Runs an agent block into the cell output and inserts generated cells below. + * Call `removeEphemeralCellsForAgentBlocks` on the batch first. Never rejects — errors become stderr on the cell. */ export async function executeAgentCell( cell: NotebookCell, @@ -192,8 +170,7 @@ export async function executeAgentCell( const prompt = cell.document.getText(); - // Per-event appends keep this O(n) bytes across the extension-host boundary, on the run's wall - // clock: runtime-core awaits `onAgentEvent`. The renderer concatenates appended stdout items. + // runtime-core awaits each event; append deltas only (O(n) over the EH boundary). const output = new NotebookCellOutput([NotebookCellOutputItem.stdout(`[Agent] Planning next steps...`)]); await execution.replaceOutput([output]); @@ -205,7 +182,6 @@ export async function executeAgentCell( throw new Error('Cell is not an agent cell'); } - // The caller clears the previous run per batch. const staleCellCount = cell.notebook .getCells() .filter((c) => getEphemeralCellAgentSourceBlockId(c) === agentBlock.id).length; @@ -220,8 +196,7 @@ export async function executeAgentCell( let lastAgentEventType: AgentStreamEvent['type'] | undefined; - // serializeNotebookContextFromBlocks does no ephemeral filtering, so this is safe only - // because of the precondition checked above. + // Caller must clear scratch cells before the batch; context serialization does not filter them. const notebookContext = serializeNotebookContext({ cells: cell.notebook.getCells().filter((c) => c.index !== cell.index), notebookName: (cell.notebook.metadata?.deepnoteNotebookName as string | undefined) ?? '' @@ -300,8 +275,7 @@ 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 does not print stacks unless isJupyterError — log stack explicitly. logger.error('Agent cell execution failed', error); if (error instanceof Error) { if (error.cause) { @@ -337,13 +311,7 @@ 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. - */ +/** Inserts an ephemeral cell after the agent; returns the created cell resolved by block id (`cellAt` clamps bad indices). */ async function insertEphemeralCell( notebook: NotebookDocument, agentCellIndex: number, @@ -398,8 +366,7 @@ export async function executeEphemeralCell( cell: NotebookCell, token: CancellationToken ): Promise { - // Bail before dispatching: rejecting the deferred alone would abandon the wait but still hand the - // generated code to the kernel. + // Cancel before dispatch — a rejected wait alone still runs the cell in the kernel. if (token.isCancellationRequested) { throw new CancellationError(); } @@ -424,9 +391,7 @@ export async function executeEphemeralCell( try { const cellIndex = cell.index; - // 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. + // Race dispatch with Idle wait — command can hang while the cell is still running. await Promise.all([ commands.executeCommand('notebook.cell.execute', { ranges: [{ start: cellIndex, end: cellIndex + 1 }], @@ -445,8 +410,7 @@ export async function executeEphemeralCell( 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. + // Surface timeout/cancel reason — "(no output)" makes the agent retry while the cell still runs. return { success: false, outputs: [], @@ -460,20 +424,9 @@ export async function executeEphemeralCell( } /** - * 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. + * Removes prior-run scratch cells owned by agents in `cells` and returns the batch without them. + * Required before `executeAgentCell` — otherwise stale generated code would run. Only agents in + * `cells` are scoped so standalone ephemeral runs stay untouched. Edit failures are logged, not thrown. */ export async function removeEphemeralCellsForAgentBlocks( notebook: NotebookDocument, diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index 185eba47bb..475381f57d 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -37,10 +37,7 @@ import { 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. - */ +// ExtensionMode.Test skips secrets; Production + in-memory store exercises real paths. function stubSecretStorage(secretStorage: Map): ServiceContainer { const context = mock(); const secrets = mock(); @@ -62,16 +59,8 @@ 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. - */ +// Mocked WorkspaceEdit.set drops edits — capture on prototype and apply to `cells`. +// `appliedEdits` is local because the shared workspace mock is never reset between tests. function applyNotebookEditsTo(cells: NotebookCell[], notebook: NotebookDocument) { type RecordedEdit = { range: { start: number; end: number }; newCells?: NotebookCellData[] }; let recordedEdits: RecordedEdit[] = []; @@ -132,9 +121,7 @@ suite('AgentCellExecutionHandler', () => { expect(describeExecutionOutputs([output])).to.equal('hello\nworld\n'); }); - // 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. + // String(line[]) joins with commas — breaks DataFrame text/plain for the agent. test('joins nbformat line arrays in execute_result text/plain', () => { const output = { output_type: 'execute_result', @@ -210,8 +197,7 @@ 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. + // Restore applyEdit default; don't reset() the shared workspace mock. when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(true)); }); @@ -222,10 +208,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. - */ function createAgentCellInMutableNotebook(cells: NotebookCell[] = [], agentBlockId = 'agent-block-1') { const notebook = createMockNotebook({ cells }); const agentCell = createMockCell({ @@ -299,8 +281,7 @@ suite('AgentCellExecutionHandler', () => { expect(item.mime).to.equal('application/vnd.code.notebook.stdout'); }); - // 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. + // Incremental deltas only — full transcript per event is O(n²) over the EH boundary. 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' }); @@ -403,9 +384,7 @@ suite('AgentCellExecutionHandler', () => { expect(text).to.include('OpenAI API key is not set'); }); - // 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. + // Caller clears prior ephemeral output; dirty notebook would duplicate agent context. test('refuses to run rather than clearing the previous run itself', async () => { const previousResult = createMockCell({ text: 'print("previous run")', @@ -457,8 +436,7 @@ suite('AgentCellExecutionHandler', () => { 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. + // cellAt clamps — failed insert must not run an existing cell at that index. test('fails the tool call without executing anything when the insert edit is rejected', async () => { const { agentCell } = createAgentCellInMutableNotebook(); let toolResult: string | undefined; @@ -525,7 +503,6 @@ suite('AgentCellExecutionHandler', () => { }); } - /** Wires the cells into a notebook whose list the applied edits actually mutate. */ function createMutableNotebook(cells: NotebookCell[]) { const notebook = createMockNotebook({ cells }); @@ -563,8 +540,7 @@ suite('AgentCellExecutionHandler', () => { 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. + // Generated cell may execute alone via notebook.cell.execute — must stay in batch. 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")'); @@ -589,8 +565,7 @@ suite('AgentCellExecutionHandler', () => { 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. + // Rejected delete edit must not block running ordinary cells in the batch. 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")'); @@ -625,7 +600,6 @@ suite('AgentCellExecutionHandler', () => { 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 () => { @@ -645,8 +619,7 @@ suite('AgentCellExecutionHandler', () => { }); }); - // Rejecting the deferred alone abandons only the wait — the generated code would still reach - // the kernel after the user cancelled. + // Pre-cancelled token must skip executeCommand, not only the idle wait. test('throws without dispatching to the kernel when the token is pre-cancelled', async () => { const cell = createMockCell({ index: 0 }); const tokenSource = new CancellationTokenSource(); @@ -679,8 +652,7 @@ suite('AgentCellExecutionHandler', () => { 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. + // Timeout must fire even when executeCommand never resolves. test('times out while the dispatch is still pending', async () => { const cell = createMockCell({ index: 0 }); const clock = sinon.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }); diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index 6cbb2797f5..96d6aeac55 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -18,10 +18,10 @@ import { injectable } from 'inversify'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { isAgentCell } from './dataConversionUtils'; -/** The key `agentBlockSchema` defines and `executeAgentBlock` reads. */ +/** Same key as `agentBlockSchema` / `executeAgentBlock`. */ const AGENT_MODEL_METADATA_KEY = 'deepnote_agent_model'; -/** Must be stored explicitly; a missing key becomes `undefined` in runtime-core and is passed to openai() as the model name. */ +/** Persisted default — absent key becomes `undefined` and breaks openai() model selection. */ const AGENT_MODEL_AUTO = 'auto'; const AGENT_MODEL_OPTIONS = [AGENT_MODEL_AUTO, 'gpt-4o', 'gpt-5']; diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.ts b/src/notebooks/deepnote/converters/agentBlockConverter.ts index ebf6cdff2c..416e01a91d 100644 --- a/src/notebooks/deepnote/converters/agentBlockConverter.ts +++ b/src/notebooks/deepnote/converters/agentBlockConverter.ts @@ -3,7 +3,7 @@ import { NotebookCellData, NotebookCellKind } from 'vscode'; import type { BlockConverter } from './blockConverter'; -/** Agent prompts render as plaintext code cells; metadata passes through in DeepnoteDataConverter. */ +/** Agent blocks as plaintext code cells; pocket metadata from DeepnoteDataConverter. */ export class AgentBlockConverter implements BlockConverter { applyChangesToBlock(block: DeepnoteBlock, cell: NotebookCellData): void { block.content = cell.value; diff --git a/src/notebooks/deepnote/dataConversionUtils.ts b/src/notebooks/deepnote/dataConversionUtils.ts index 69fc7a8d86..ea9948eaa8 100644 --- a/src/notebooks/deepnote/dataConversionUtils.ts +++ b/src/notebooks/deepnote/dataConversionUtils.ts @@ -26,34 +26,21 @@ 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. - */ +/** Agent block cell. Lives here so importers avoid `@deepnote/runtime-core`. */ 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). - */ +/** Agent-generated scratch cell (`metadata.is_ephemeral`). */ 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. + * Serialized block id, or undefined. Prefer `__deepnoteBlockId` — VS Code may rewrite `id`. + * Missing id makes callers mint a new one and reassign the block on save. */ export function getBlockId(cell: NotebookCell | NotebookCellData): string | undefined { return ( @@ -63,10 +50,7 @@ export function getBlockId(cell: NotebookCell | NotebookCellData): string | unde ); } -/** - * Returns the id of the agent block that generated this ephemeral cell, or undefined if the cell - * isn't agent-generated scratch. - */ +/** Owning agent block id when `isEphemeralCell`; otherwise undefined. */ export function getEphemeralCellAgentSourceBlockId(cell: NotebookCell): string | undefined { return isEphemeralCell(cell) ? (cell.metadata?.agent_source_block_id as string | undefined) : undefined; } diff --git a/src/notebooks/deepnote/dataConversionUtils.unit.test.ts b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts index 2b39574c45..1b68c61595 100644 --- a/src/notebooks/deepnote/dataConversionUtils.unit.test.ts +++ b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts @@ -49,8 +49,7 @@ suite('DataConversionUtils', () => { 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. + // Fallback-cell metadata; minting a new id would reassign the block on save. test('falls back to the legacy deepnoteBlockId when both are absent', () => { const cell = createMockCell({ metadata: { deepnoteBlockId: 'legacy-id' } }); @@ -77,7 +76,7 @@ suite('DataConversionUtils', () => { expect(getEphemeralCellAgentSourceBlockId(cell)).to.equal('agent-block-1'); }); - // An ordinary cell that happens to carry the metadata is not the agent's to delete. + // agent_source_block_id alone does not mark a cell for agent cleanup. test('returns undefined when the cell is not marked ephemeral', () => { const cell = createMockCell({ metadata: { agent_source_block_id: 'agent-block-1' } }); diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index 159b4b4b7c..102b4dfaf9 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -164,8 +164,7 @@ 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 { - // 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. + // Ephemeral cells aren't persisted; counting them looks like an external delete. 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 2a95a82051..e45ca19899 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -180,8 +180,7 @@ project: 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. + // Ephemeral cells are not serialized — extra live cells must not look like an external edit. const notebook = createMockNotebook({ uri, cells: [ diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index eecb22ec2c..6cdaa35d31 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -1091,7 +1091,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, controller.supportsExecutionOrder = true; controller.supportedLanguages = ['python', 'sql', 'markdown', 'plaintext']; - // Run here only prompts for an environment; kernel execution uses the real controller afterward. + // Environment picker only; execution goes through the real controller on retry. controller.executeHandler = async (cells, doc) => { logger.info( `Placeholder controller execute handler called for ${getDisplayPath(doc.uri)} with ${ diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 2c6df36a15..0365216810 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -231,7 +231,6 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { throw new Error(`Notebook with ID ${notebookId} not found in project`); } - // Exclude ephemeral cells (agent-generated) from persistence const nonEphemeralCells = data.cells.filter((cell) => !isEphemeralCell(cell)); logger.debug( diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts index 2e1108b52e..30878e73a7 100644 --- a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts @@ -17,7 +17,7 @@ import { isEphemeralCell } from './dataConversionUtils'; const NOTEBOOK_CELL_SCHEME = 'vscode-notebook-cell'; -/** Code cell editor decorations for `is_ephemeral` blocks; markup uses `src/renderers/client/markdown.ts`. */ +/** Ephemeral styling in code-cell editors; markdown cells use `src/renderers/client/markdown.ts`. */ @injectable() export class EphemeralCellDecorationProvider implements IExtensionSyncActivationService { private readonly disposables: Disposable[] = []; diff --git a/src/renderers/client/markdown.ts b/src/renderers/client/markdown.ts index 8537598b5e..4374444265 100644 --- a/src/renderers/client/markdown.ts +++ b/src/renderers/client/markdown.ts @@ -1,6 +1,6 @@ import type { ActivationFunction } from 'vscode-notebook-renderer'; -// Minimal markdown-it surface (no package types; transitive dependency only). +// Local markdown-it shape — transitive dep, no types. interface MarkdownItToken { content: string; } diff --git a/src/test/mocks/deepnoteRuntimeCore.ts b/src/test/mocks/deepnoteRuntimeCore.ts index 4dccd7953f..e759d3bf2b 100644 --- a/src/test/mocks/deepnoteRuntimeCore.ts +++ b/src/test/mocks/deepnoteRuntimeCore.ts @@ -3,13 +3,8 @@ import type { AgentBlockContext, ServerInfo, ServerOptions } from '@deepnote/run import type { ChildProcess } from 'child_process'; /** - * Mock of @deepnote/runtime-core for unit tests: the real startServer/stopServer spawn and - * 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. - * The exports are typed against the real package so the mock cannot drift from its API. + * Mock @deepnote/runtime-core: no Python spawns or live agent API; records calls, returns stubs. + * build/mocha-esm-loader.js aliases the package here for one shared instance; exports match real types. */ type RuntimeCore = typeof import('@deepnote/runtime-core'); diff --git a/test/e2e/.mocharc.js b/test/e2e/.mocharc.js index 50d9e789fd..caa8289f15 100644 --- a/test/e2e/.mocharc.js +++ b/test/e2e/.mocharc.js @@ -3,7 +3,7 @@ // tests are the real guard rails; this is a generous suite-level safety net. const path = require('path'); -// ExTester uses `new Mocha(config)` and ignores `require`; load rootHooks here as `rootHooks`. +// ExTester ignores Mocha `require`; wire rootHooks from compiled output. const { mochaHooks } = require(path.resolve(__dirname, '..', '..', 'out', 'e2e', 'rootHooks.js')); module.exports = { diff --git a/test/e2e/helpers/mockOpenAiServer.ts b/test/e2e/helpers/mockOpenAiServer.ts index b998e99af5..bb718dc284 100644 --- a/test/e2e/helpers/mockOpenAiServer.ts +++ b/test/e2e/helpers/mockOpenAiServer.ts @@ -5,17 +5,14 @@ import * as os from 'os'; import * as path from 'path'; import { setTimeout as delay } from 'timers/promises'; -// aimock is npx-only so its peer deps (jest/vitest) never enter this repo's lockfile. +// npx aimock — keep jest/vitest peers out of the lockfile. const AIMOCK_VERSION = '1.37.4'; const AIMOCK_BIN = 'llmock'; -// Below typical ephemeral port range so a stray outbound source port cannot fake the pre-flight check. +// Fixed port below ephemeral range (connect pre-flight). const MOCK_OPENAI_PORT = 18_937; -/** - * 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. - */ +/** Set OPENAI_BASE_URL at module scope — ExTester spawns the host before `before` hooks. */ export function pointExtensionHostAtMockServer(): void { process.env.OPENAI_BASE_URL = `http://127.0.0.1:${MOCK_OPENAI_PORT}/v1`; } @@ -34,7 +31,7 @@ export interface MockToolCall { name: string; } -/** Predicate per scripted leg (not sequence index — safe across Mocha retries). */ +/** Per-leg match predicate (not call order); Mocha-retry-safe. */ export type MockAgentMatch = { hasToolResult: false } | { toolResultContains: string }; export type MockAgentResponse = { content: string } | { toolCall: MockToolCall }; diff --git a/test/e2e/helpers/notebook.ts b/test/e2e/helpers/notebook.ts index 83c912ec0b..8f6af68872 100644 --- a/test/e2e/helpers/notebook.ts +++ b/test/e2e/helpers/notebook.ts @@ -41,7 +41,7 @@ export async function clickRunAll(notebookFileName: string): Promise { ); } -/** Runs `read` inside the notebook output webview; returns '' when the frame is missing or not ready. */ +/** Run `read` in the notebook output webview; '' if the frame is missing. */ async function readInsideNotebookWebview(read: (webView: WebView) => Promise): Promise { const driver = VSBrowser.instance.driver; const webView = new WebView(); @@ -73,12 +73,12 @@ async function readInsideNotebookWebview(read: (webView: WebView) => 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. */ +/** Cell output once; falls back to frame body if output selectors miss. */ export async function readRenderedOutput(): Promise { return readInsideNotebookWebview(async (webView) => { const elements = await webView.findWebElements(By.css(OUTPUT_SELECTOR)); diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts index a9aa076ba3..4e3551c91f 100644 --- a/test/e2e/suite/agentBlock.e2e.test.ts +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -1,8 +1,4 @@ -/** - * 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). - */ +/** Agent block E2E vs local aimock; legs 2–3 advance on real tool results (no live OpenAI). */ import { EditorView, InputBox, VSBrowser, WebView, Workbench } from 'vscode-extension-tester'; @@ -33,21 +29,21 @@ pointExtensionHostAtMockServer(); const AGENT_FILE = 'agent-block.deepnote'; const CODE_TOOL_NAME = 'add_code_block'; const MARKDOWN_TOOL_NAME = 'add_markdown_block'; -// Coupled to agentCellExecutionHandler tool result for add_markdown_block (leg 3 match). +// Leg 3 match: agentCellExecutionHandler add_markdown_block tool result. const MARKDOWN_BLOCK_ADDED_TEXT = 'Markdown block added.'; const ENVIRONMENT_NAME = 'E2E Agent Env'; 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'; -// aimock chunks content at 20 characters, so this length arrives as several text_delta events. +// aimock emits 20-char chunks (multiple text_delta). const FINAL_AGENT_TEXT = 'Summary added as a markdown block, streamed across several deltas.'; -// Not substrings of the first run's markers; the counts below depend on that. +// Disjoint from first-run markers (assertOccurrences below). const RERUN_PYTHON_OUTPUT_MARKER = 'rerun-python-ran'; const RERUN_GENERATED_PYTHON = `print("${RERUN_PYTHON_OUTPUT_MARKER}")`; const RERUN_MARKDOWN_TEXT = 'Second-run markdown from the E2E agent'; const RERUN_FINAL_AGENT_TEXT = 'Re-run summary added as a markdown block.'; -// Coupled to executeAgentCell's uncleared-previous-run error. +// executeAgentCell stale-run error substring. const STALE_CELLS_ERROR_TEXT = 'from its previous run'; const MOCK_API_KEY = 'sk-e2e-mock-key'; const SET_API_KEY_COMMAND = 'Deepnote: Set OpenAI API Key'; @@ -235,7 +231,7 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu assertRenderedContiguously(transcript, `[Agent] Text:\n${FINAL_AGENT_TEXT}`); }); - // Depends on the previous spec: the block starts this run owning the cells it generated there. + // Serial with prior it — block still owns first-run cells. it('clears the cells its previous run generated instead of stacking a second copy', async function () { mockServer = await startMockOpenAiServer([ { @@ -275,7 +271,7 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu await screenshot('agent-rerun'); - // Counts, not presence: a Mocha retry repeats the markers. + // assertOccurrences — retries would duplicate markers. assertOccurrences(rendered, PYTHON_OUTPUT_MARKER, 0); assertOccurrences(rendered, EPHEMERAL_MARKDOWN_TEXT, 0); assertOccurrences(rendered, RERUN_PYTHON_OUTPUT_MARKER, 1); From 377a8fe15023d20783ced3cbb98dbb202bee5005 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 6 Aug 2026 15:40:09 +0000 Subject: [PATCH 42/80] test(agent-block): merge overlapping unit tests to cut PR noise Consolidate duplicate setup and assertions across agent-block test files while keeping distinct failure paths covered. Co-authored-by: Cursor --- .../agentCellExecutionHandler.unit.test.ts | 112 ++++---------- .../agentCellStatusBarProvider.unit.test.ts | 62 +------- .../agentBlockConverter.unit.test.ts | 98 +++--------- .../deepnote/dataConversionUtils.unit.test.ts | 42 ++---- .../deepnote/deepnoteSecretStore.unit.test.ts | 139 +++++------------- src/notebooks/deepnote/deepnoteTestHelpers.ts | 24 +-- ...phemeralCellStatusBarProvider.unit.test.ts | 87 +++-------- src/platform/deepnote/pocket.unit.test.ts | 17 +-- 8 files changed, 139 insertions(+), 442 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index 475381f57d..b0598893ea 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -19,7 +19,7 @@ import { } from 'vscode'; import type { AgentBlock } from '@deepnote/blocks'; -import type { AgentBlockContext, AgentBlockResult } from '@deepnote/runtime-core'; +import type { AgentBlockContext } from '@deepnote/runtime-core'; import type { IDisposable } from '../../platform/common/types'; import { IExtensionContext } from '../../platform/common/types'; @@ -122,25 +122,21 @@ suite('AgentCellExecutionHandler', () => { }); // String(line[]) joins with commas — breaks DataFrame text/plain for the agent. - test('joins nbformat line arrays in execute_result text/plain', () => { - const output = { + test('joins nbformat line arrays in execute_result and display_data text/plain', () => { + const executeResult = { output_type: 'execute_result', data: { 'text/plain': [' a b\n', '0 1 4\n', '1 2 5'] }, metadata: {}, execution_count: 1 }; - - expect(describeExecutionOutputs([output])).to.equal(' a b\n0 1 4\n1 2 5'); - }); - - test('joins nbformat line arrays in display_data text/plain', () => { - const output = { + const displayData = { output_type: 'display_data', data: { 'text/plain': ['line one\n', 'line two'] }, metadata: {} }; - expect(describeExecutionOutputs([output])).to.equal('line one\nline two'); + expect(describeExecutionOutputs([executeResult])).to.equal(' a b\n0 1 4\n1 2 5'); + expect(describeExecutionOutputs([displayData])).to.equal('line one\nline two'); }); test('leaves single-line text/plain untouched', () => { @@ -191,7 +187,7 @@ suite('AgentCellExecutionHandler', () => { createNotebookCellExecution: sinon.stub().returns(mockExecution) } as unknown as NotebookController; - executeAgentBlockStub = sinon.stub().resolves({ finalOutput: 'done' } as AgentBlockResult); + executeAgentBlockStub = sinon.stub().resolves({ finalOutput: 'done' }); }); teardown(() => { @@ -199,6 +195,7 @@ suite('AgentCellExecutionHandler', () => { reset(mockedVSCodeNamespaces.commands); // Restore applyEdit default; don't reset() the shared workspace mock. when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(true)); + secretStorage.clear(); }); function createAgentCell(text: string = 'Test prompt') { @@ -230,29 +227,15 @@ suite('AgentCellExecutionHandler', () => { return Buffer.from(item.data).toString('utf-8'); } - test('creates execution and starts it', async () => { + test('creates execution, clears output, sets planning output, and ends successfully', async () => { const cell = createAgentCell('Analyze data'); 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 = createAgentCell('Analyze data'); - - 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 = createAgentCell('Hello world'); - - await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - expect(mockExecution.replaceOutput.calledOnce).to.be.true; const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; @@ -261,14 +244,17 @@ suite('AgentCellExecutionHandler', () => { const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); expect(text).to.include('[Agent] Planning next steps...'); + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.true; }); - test('streams events via appendOutputItems using onAgentEvent callback', async () => { + // Incremental deltas only — full transcript per event is O(n²) over the EH boundary. + test('streams text_delta events via appendOutputItems with incremental stdout chunks', async () => { executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { - await context.onAgentEvent?.({ type: 'text_delta', text: 'Hello ' }); - await context.onAgentEvent?.({ type: 'text_delta', text: 'world' }); + await context.onAgentEvent?.({ type: 'text_delta', text: 'first' }); + await context.onAgentEvent?.({ type: 'text_delta', text: ' second' }); - return { finalOutput: 'Hello world' } as AgentBlockResult; + return { finalOutput: 'first second' }; }); const cell = createAgentCell(); @@ -279,21 +265,6 @@ suite('AgentCellExecutionHandler', () => { const item = mockExecution.appendOutputItems.firstCall.args[0] as NotebookCellOutputItem; expect(item.mime).to.equal('application/vnd.code.notebook.stdout'); - }); - - // Incremental deltas only — full transcript per event is O(n²) over the EH boundary. - 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' }); - - return { finalOutput: 'first second' } as AgentBlockResult; - }); - - const cell = createAgentCell(); - - await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - expect(getStdoutChunkText(0)).to.equal('[Agent] Text:\nfirst'); expect(getStdoutChunkText(1)).to.equal(' second'); }); @@ -303,7 +274,7 @@ suite('AgentCellExecutionHandler', () => { await context.onAgentEvent?.({ type: 'text_delta', text: 'thinking...' }); await context.onAgentEvent?.({ type: 'tool_called', toolName: 'search' }); - return { finalOutput: '' } as AgentBlockResult; + return { finalOutput: '' }; }); const cell = createAgentCell(); @@ -315,17 +286,8 @@ suite('AgentCellExecutionHandler', () => { expect(chunk2).to.include('[Agent] Tool called: search'); }); - test('ends execution with success', async () => { - const cell = createAgentCell(); - - await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - - 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')); + test('fails execution and writes clearOutput error to stderr', async () => { + mockExecution.clearOutput.rejects(new Error('Something went wrong')); const cell = createAgentCell(); @@ -333,15 +295,6 @@ suite('AgentCellExecutionHandler', () => { 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 = createAgentCell(); - - await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - expect(mockExecution.appendOutput.calledOnce).to.be.true; const outputs = mockExecution.appendOutput.firstCall.args[0] as NotebookCellOutput[]; @@ -404,36 +357,23 @@ suite('AgentCellExecutionHandler', () => { expect(text).to.include('previous run'); }); - 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 () => { + test('inserts ephemeral markdown cells after the agent cell in order', 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; + return { finalOutput: '' }; }); await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect(cells.map((cell) => cell.document.getText())).to.deep.equal(['Test prompt', 'first', 'second']); + expect(cells[1].metadata?.is_ephemeral).to.be.true; + expect(cells[1].metadata?.agent_source_block_id).to.equal(agentCell.metadata?.id); + expect(cells[2].metadata?.is_ephemeral).to.be.true; + expect(cells[2].metadata?.agent_source_block_id).to.equal(agentCell.metadata?.id); }); // cellAt clamps — failed insert must not run an existing cell at that index. @@ -447,7 +387,7 @@ suite('AgentCellExecutionHandler', () => { executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { toolResult = await context.addAndExecuteCodeBlock({ code: 'print(1)' }); - return { finalOutput: '' } as AgentBlockResult; + return { finalOutput: '' }; }); await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts index c8463c4a34..faaad4604c 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -77,7 +77,7 @@ suite('AgentCellStatusBarProvider', () => { }); suite('Agent Block Indicator', () => { - test('Should display agent block label with icon', () => { + test('Should display agent block label with icon and no command', () => { const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); const items = provider.provideCellStatusBarItems(cell, mockToken)!; @@ -85,30 +85,29 @@ suite('AgentCellStatusBarProvider', () => { 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', () => { + test('Should display default model picker for agent cell without model metadata', () => { 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)'); + expect(items[1].command).to.not.be.undefined; + const cmd = items[1].command as any; + expect(cmd.command).to.equal('deepnote.switchAgentModel'); + expect(items[1].priority).to.equal(90); }); test('Should display configured model from metadata', () => { const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' }, - deepnote_agent_model: 'gpt-4o' + deepnote_agent_model: 'gpt-4o', + deepnote_max_iterations: 50 } }); const items = provider.provideCellStatusBarItems(cell, mockToken)!; @@ -116,18 +115,6 @@ suite('AgentCellStatusBarProvider', () => { expect(items[1].text).to.include('Model: gpt-4o'); }); - test('Should display gpt-5 model', () => { - const cell = createMockCell({ - metadata: { - __deepnotePocket: { type: 'agent' }, - deepnote_agent_model: 'gpt-5' - } - }); - const items = provider.provideCellStatusBarItems(cell, mockToken)!; - - expect(items[1].text).to.include('Model: gpt-5'); - }); - test('Should display "auto" when model is empty string', () => { const cell = createMockCell({ metadata: { @@ -139,38 +126,5 @@ suite('AgentCellStatusBarProvider', () => { 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('Combined metadata', () => { - test('Should ignore metadata keys the runtime does not consume', () => { - const cell = createMockCell({ - metadata: { - __deepnotePocket: { type: 'agent' }, - deepnote_agent_model: 'gpt-4o', - deepnote_max_iterations: 50 - } - }); - const items = provider.provideCellStatusBarItems(cell, mockToken)!; - - expect(items).to.have.lengthOf(2); - expect(items[0].text).to.include('Agent Block'); - expect(items[1].text).to.include('Model: gpt-4o'); - }); }); }); diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts index a3ce26acf8..1a2c46337f 100644 --- a/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts +++ b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts @@ -12,11 +12,8 @@ suite('AgentBlockConverter', () => { }); suite('canConvert', () => { - test('returns true for "agent" type', () => { + test('accepts agent type case-insensitively', () => { assert.strictEqual(converter.canConvert('agent'), true); - }); - - test('returns true for "Agent" type (case insensitive)', () => { assert.strictEqual(converter.canConvert('Agent'), true); }); @@ -53,8 +50,8 @@ suite('AgentBlockConverter', () => { assert.strictEqual(cell.languageId, 'plaintext'); }); - test('handles empty content', () => { - const block: DeepnoteBlock = { + test('normalizes missing or empty content to empty cell value', () => { + const emptyBlock: DeepnoteBlock = { blockGroup: 'test-group', content: '', id: 'agent-block-456', @@ -62,16 +59,7 @@ suite('AgentBlockConverter', () => { 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, 'plaintext'); - }); - - test('handles undefined content', () => { - const block: DeepnoteBlock = { + const undefinedContentBlock: DeepnoteBlock = { blockGroup: 'test-group', id: 'agent-block-789', sortingKey: 'a2', @@ -79,11 +67,8 @@ suite('AgentBlockConverter', () => { type: 'agent' }; - const cell = converter.convertToCell(block); - - assert.strictEqual(cell.kind, NotebookCellKind.Code); - assert.strictEqual(cell.value, ''); - assert.strictEqual(cell.languageId, 'plaintext'); + assert.strictEqual(converter.convertToCell(emptyBlock).value, ''); + assert.strictEqual(converter.convertToCell(undefinedContentBlock).value, ''); }); test('preserves multiline prompt', () => { @@ -111,46 +96,33 @@ suite('AgentBlockConverter', () => { assert.strictEqual(cell.value, prompt); assert.strictEqual(cell.languageId, 'plaintext'); }); - - 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, 'plaintext'); - }); }); suite('applyChangesToBlock', () => { - test('updates block content from cell value', () => { + test('updates block content without modifying other block properties', () => { const block: DeepnoteBlock = { blockGroup: 'test-group', content: 'Old prompt', - id: 'agent-block-123', - sortingKey: 'a0', - metadata: { deepnote_agent_model: 'auto' }, + id: 'agent-block-789', + metadata: { + deepnote_agent_model: 'gpt-4o', + custom: 'value' + }, + sortingKey: 'a2', type: 'agent' }; - const cell = new NotebookCellData( - NotebookCellKind.Code, - 'New prompt with updated instructions', - 'plaintext' - ); + const cell = new NotebookCellData(NotebookCellKind.Code, 'New prompt', 'plaintext'); converter.applyChangesToBlock(block, cell); - assert.strictEqual(block.content, 'New prompt with updated instructions'); + 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' + }); }); test('handles empty cell value', () => { @@ -168,31 +140,5 @@ suite('AgentBlockConverter', () => { 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', 'plaintext'); - - 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/dataConversionUtils.unit.test.ts b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts index 1b68c61595..ee90e0c0b9 100644 --- a/src/notebooks/deepnote/dataConversionUtils.unit.test.ts +++ b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts @@ -11,28 +11,20 @@ suite('DataConversionUtils', () => { expect(isAgentCell(cell)).to.be.true; }); - test('returns false for cell with code pocket type', () => { - const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } }); + test('returns false for non-agent pocket types', () => { + const codeCell = createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } }); + const markdownCell = createMockCell({ metadata: { __deepnotePocket: { type: 'markdown' } } }); - expect(isAgentCell(cell)).to.be.false; + expect(isAgentCell(codeCell)).to.be.false; + expect(isAgentCell(markdownCell)).to.be.false; }); - test('returns false for cell with markdown pocket type', () => { - const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'markdown' } } }); + test('returns false when pocket type is not agent', () => { + const noPocketCell = createMockCell({ metadata: {} }); + const noMetadataCell = createMockCell({ metadata: undefined }); - 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; + expect(isAgentCell(noPocketCell)).to.be.false; + expect(isAgentCell(noMetadataCell)).to.be.false; }); }); @@ -77,16 +69,12 @@ suite('DataConversionUtils', () => { }); // agent_source_block_id alone does not mark a cell for agent cleanup. - 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: {} }); + test('returns undefined when the cell is not ephemeral or ordinary', () => { + const withSourceOnly = createMockCell({ metadata: { agent_source_block_id: 'agent-block-1' } }); + const ordinaryCell = createMockCell({ metadata: {} }); - expect(getEphemeralCellAgentSourceBlockId(cell)).to.be.undefined; + expect(getEphemeralCellAgentSourceBlockId(withSourceOnly)).to.be.undefined; + expect(getEphemeralCellAgentSourceBlockId(ordinaryCell)).to.be.undefined; }); }); }); diff --git a/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts index 3ba41edd7f..f0c339a29a 100644 --- a/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts @@ -7,15 +7,11 @@ import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; 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'; @@ -118,124 +114,55 @@ suite('deepnoteSecretStore', () => { assert.isUndefined(value); }); - }); - - 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'); - - assert.strictEqual(value, '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'); - assert.fail('Should have thrown'); - } catch (e) { - assert.strictEqual((e as Error).message, 'Value is required'); - } - }); - }); - - suite('getOpenAiApiKey', () => { - test('returns key when stored', async () => { - secretStorage.set('openAiApiKey', 'test-key'); - - const key = await getOpenAiApiKey(); - - assert.strictEqual(key, 'test-key'); - }); - - test('returns undefined when not set', async () => { - const key = await getOpenAiApiKey(); - - assert.isUndefined(key); - }); - - test('returns undefined when key is empty string', async () => { - secretStorage.set('openAiApiKey', ''); - - const key = await getOpenAiApiKey(); - - assert.isUndefined(key); - }); - }); - - suite('setOpenAiApiKey', () => { - test('stores key in secrets', async () => { - await setOpenAiApiKey('my-api-key'); - - assert.strictEqual(secretStorage.get('openAiApiKey'), 'my-api-key'); - }); - }); - - suite('clearOpenAiApiKey', () => { - test('deletes key from secrets', async () => { - secretStorage.set('openAiApiKey', 'test-key'); - - await clearOpenAiApiKey(); - - assert.isFalse(secretStorage.has('openAiApiKey')); - }); - }); - - 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(); - - assert.strictEqual(key, 'sk-abc123'); - assert.strictEqual(secretStorage.get('openAiApiKey'), 'sk-abc123'); - }); - - test('returns undefined when user cancels', async () => { - when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); - - const key = await promptForOpenAiApiKey(); - - assert.isUndefined(key); - }); test('returns undefined when user enters empty string', async () => { when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(' ')); - const key = await promptForOpenAiApiKey(); + const value = await promptForSecret('customKey', { prompt: 'Enter value' }); - assert.isUndefined(key); + assert.isUndefined(value); }); }); - suite('getOrPromptOpenAiApiKey', () => { - test('returns key when present in store', async () => { - secretStorage.set('openAiApiKey', 'stored-key'); + suite('generic getOrPromptSecret', () => { + test('returns value when present in store', async () => { + secretStorage.set('customKey', 'stored-value'); - const key = await getOrPromptOpenAiApiKey(); + const value = await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); - assert.strictEqual(key, 'stored-key'); + assert.strictEqual(value, 'stored-value'); }); - test('prompts and returns key when missing', async () => { - when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('prompted-key')); + test('prompts and returns value when missing', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('prompted-value')); - const key = await getOrPromptOpenAiApiKey(); + const value = await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); - assert.strictEqual(key, 'prompted-key'); + assert.strictEqual(value, 'prompted-value'); }); - test('throws when key missing and user cancels prompt', async () => { - when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); - - try { - await getOrPromptOpenAiApiKey(); - assert.fail('Should have thrown'); - } catch (e) { - assert.include((e as Error).message, 'OpenAI API key is not set'); + for (const scenario of [ + { + label: 'generic', + run: () => getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'), + assertError: (e: Error) => assert.strictEqual(e.message, 'Value is required') + }, + { + label: 'openAi', + run: () => getOrPromptOpenAiApiKey(), + assertError: (e: Error) => assert.include(e.message, 'OpenAI API key is not set') } - }); + ]) { + test(`throws when value missing and user cancels prompt (${scenario.label})`, async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + try { + await scenario.run(); + assert.fail('Should have thrown'); + } catch (e) { + scenario.assertError(e as Error); + } + }); + } }); }); diff --git a/src/notebooks/deepnote/deepnoteTestHelpers.ts b/src/notebooks/deepnote/deepnoteTestHelpers.ts index b93e0bae59..6c6b766bcf 100644 --- a/src/notebooks/deepnote/deepnoteTestHelpers.ts +++ b/src/notebooks/deepnote/deepnoteTestHelpers.ts @@ -4,7 +4,10 @@ import { NotebookCellKind, NotebookCellOutput, NotebookDocument, + Position, TextDocument, + Range, + TextLine, Uri, WorkspaceFolder } from 'vscode'; @@ -29,6 +32,7 @@ export interface CreateMockCellOptions { notebookUri?: Uri; notebookMetadata?: Record; index?: number; + mime?: string; } /** @@ -118,11 +122,12 @@ export function createMockCell(options?: CreateMockCellOptions): NotebookCell { outputs = [], notebookType = 'deepnote', notebookUri = Uri.file('/test/notebook.deepnote'), - index = 0 + index = 0, + mime = 'text/plain' } = opts; // Preserve explicit undefined for metadata fields - const metadata = Object.prototype.hasOwnProperty.call(opts, 'metadata') ? opts.metadata : {}; + const metadata = 'metadata' in opts ? opts.metadata ?? {} : {}; const notebookMetadata = Object.prototype.hasOwnProperty.call(opts, 'notebookMetadata') ? opts.notebookMetadata : {}; @@ -135,7 +140,7 @@ export function createMockCell(options?: CreateMockCellOptions): NotebookCell { const cellPath = `${notebookUri.path}#cell${index}`; - const document = { + const document: TextDocument = { uri: Uri.file(cellPath), fileName: cellPath, isUntitled: false, @@ -147,24 +152,25 @@ export function createMockCell(options?: CreateMockCellOptions): NotebookCell { save: async () => true, eol: 1, lineCount: 1, - lineAt: () => ({ text: '' }) as unknown, + lineAt: () => ({ text: '' }) as unknown as TextLine, offsetAt: () => 0, - positionAt: () => ({}) as unknown, - validateRange: () => ({}) as unknown, - validatePosition: () => ({}) as unknown, + positionAt: () => new Position(0, 0), + validateRange: () => new Range(new Position(0, 0), new Position(0, 0)), + validatePosition: () => new Position(0, 0), getWordRangeAtPosition: () => undefined, encoding: 'utf-8' - } as unknown as TextDocument; + }; return { index, + mime, notebook, kind, document, metadata, outputs, executionSummary: undefined - } as unknown as NotebookCell; + }; } /** A Deepnote code block (whole-file YAML shape); override any field. */ diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts index 27e53dcdf2..4785815bcc 100644 --- a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts @@ -21,13 +21,6 @@ suite('EphemeralCellStatusBarProvider', () => { }); 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); @@ -69,42 +62,25 @@ suite('EphemeralCellStatusBarProvider', () => { }); suite('Status Bar Item Properties', () => { - test('Should display sparkle icon with Ephemeral label', () => { + test('Should set ephemeral status bar item properties', () => { 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', () => { + test('Should describe ephemeral tooltip without source block when agent_source_block_id is absent', () => { const cell = createMockCell({ metadata: { is_ephemeral: true } }); const item = provider.provideCellStatusBarItems(cell, mockToken)!; expect(item.tooltip).to.include('Auto-generated ephemeral block'); + expect(item.tooltip).to.not.include('Source agent block'); }); test('Should include agent source block ID in tooltip when present', () => { @@ -119,51 +95,24 @@ suite('EphemeralCellStatusBarProvider', () => { 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; + test('Should return item for ephemeral cells regardless of pocket type', () => { + const pocketTypes = ['agent', 'code', 'markdown'] as const; + + for (const type of pocketTypes) { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type }, + is_ephemeral: true, + ...(type === 'agent' ? { agent_source_block_id: 'source-id' } : {}) + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + } }); }); }); diff --git a/src/platform/deepnote/pocket.unit.test.ts b/src/platform/deepnote/pocket.unit.test.ts index 73b317561a..7671412e38 100644 --- a/src/platform/deepnote/pocket.unit.test.ts +++ b/src/platform/deepnote/pocket.unit.test.ts @@ -121,7 +121,8 @@ suite('Pocket', () => { sortingKey: 'a0', executionCount: 5 }, - id: 'block-123', + __deepnoteBlockId: 'block-123', + id: 'rewritten-by-vscode', custom: 'value' }; @@ -135,20 +136,6 @@ suite('Pocket', () => { assert.strictEqual((block as any).outputs, undefined); }); - 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 16198a2fc343e62728ce8f711b89fd25ba17c616 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 6 Aug 2026 16:17:36 +0000 Subject: [PATCH 43/80] test(agent-block): use suite hooks for ephemeral cell test cleanup Replace try/finally disposal of tokens and fake timers with nested suites and setup/teardown. Co-authored-by: Cursor --- .../agentCellExecutionHandler.unit.test.ts | 129 +++++++++++------- 1 file changed, 76 insertions(+), 53 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index b0598893ea..f94fccf653 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -523,81 +523,106 @@ suite('AgentCellExecutionHandler', () => { }); suite('executeEphemeralCell', () => { - let tokenSource: CancellationTokenSource; + suite('with active cancellation token', () => { + let tokenSource: CancellationTokenSource; - setup(() => { - tokenSource = new CancellationTokenSource(); - }); + setup(() => { + tokenSource = new CancellationTokenSource(); + }); - teardown(() => { - tokenSource.dispose(); - reset(mockedVSCodeNamespaces.commands); - }); + teardown(() => { + tokenSource.dispose(); + reset(mockedVSCodeNamespaces.commands); + }); - test('uses current cell index, not stale index from insertion time', async () => { - const staleIndex = 5; - const currentIndex = 6; + test('uses current cell index, not stale index from insertion time', async () => { + const staleIndex = 5; + const currentIndex = 6; - const cell = createMockCell({ index: staleIndex }); + const cell = createMockCell({ index: staleIndex }); - (cell as { index: number }).index = currentIndex; + (cell as { index: number }).index = currentIndex; + + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenCall(async () => { + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Idle); + }); - when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenCall(async () => { - notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Idle); + await executeEphemeralCell(cell, tokenSource.token); + + 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 + }); }); - await executeEphemeralCell(cell, tokenSource.token); + test('reports the failure reason instead of swallowing it', async () => { + const cell = createMockCell({ index: 0 }); - const [commandName, commandArg] = capture( - mockedVSCodeNamespaces.commands.executeCommand as (cmd: string, arg: unknown) => Thenable - ).last(); + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenReject( + new Error('kernel is dead') + ); - expect(commandName).to.equal('notebook.cell.execute'); - expect(commandArg).to.deep.equal({ - ranges: [{ start: currentIndex, end: currentIndex + 1 }], - document: cell.notebook.uri + const result = await executeEphemeralCell(cell, tokenSource.token); + + expect(result.success).to.be.false; + expect(result.error).to.equal('kernel is dead'); }); }); // Pre-cancelled token must skip executeCommand, not only the idle wait. - 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(); + suite('with pre-cancelled token', () => { + let tokenSource: CancellationTokenSource; - when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); + setup(() => { + tokenSource = new CancellationTokenSource(); + tokenSource.cancel(); + }); - try { - await executeEphemeralCell(cell, tokenSource.token); - expect.fail('Should have thrown'); - } catch (e) { - expect(e).to.be.instanceOf(CancellationError); - } finally { + teardown(() => { tokenSource.dispose(); - } - - verify(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).never(); - }); + reset(mockedVSCodeNamespaces.commands); + }); - test('reports the failure reason instead of swallowing it', async () => { - const cell = createMockCell({ index: 0 }); + test('throws without dispatching to the kernel when the token is pre-cancelled', async () => { + const cell = createMockCell({ index: 0 }); - when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenReject( - new Error('kernel is dead') - ); + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); - const result = await executeEphemeralCell(cell, tokenSource.token); + try { + await executeEphemeralCell(cell, tokenSource.token); + expect.fail('Should have thrown'); + } catch (e) { + expect(e).to.be.instanceOf(CancellationError); + } - expect(result.success).to.be.false; - expect(result.error).to.equal('kernel is dead'); + verify(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).never(); + }); }); // Timeout must fire even when executeCommand never resolves. - test('times out while the dispatch is still pending', async () => { - const cell = createMockCell({ index: 0 }); - const clock = sinon.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }); + suite('with fake timers', () => { + let clock: sinon.SinonFakeTimers; + let tokenSource: CancellationTokenSource; + + setup(() => { + clock = sinon.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }); + tokenSource = new CancellationTokenSource(); + }); + + teardown(() => { + clock.restore(); + tokenSource.dispose(); + reset(mockedVSCodeNamespaces.commands); + }); + + test('times out while the dispatch is still pending', async () => { + const cell = createMockCell({ index: 0 }); - try { when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenCall( () => new Promise(() => undefined) ); @@ -609,9 +634,7 @@ suite('AgentCellExecutionHandler', () => { expect(result.success).to.be.false; expect(result.error).to.equal('Ephemeral cell execution timed out'); - } finally { - clock.restore(); - } + }); }); }); }); From d535cd6a97ad37bbad7315faf885b90c058fd3ad Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 6 Aug 2026 17:34:53 +0000 Subject: [PATCH 44/80] feat(agent-block): enhance execution flow for agent cells Refactor the execution logic in VSCodeNotebookController to ensure that agent cells are executed correctly and notify the completion of the execution queue when an agent cell is run. This change improves the handling of pending kernel cells and integrates the notebookCellExecutions notification for better state management. Additionally, update unit tests to support the new execution flow and introduce a helper function for creating mock notebooks with cells, enhancing test coverage and maintainability. Co-authored-by: Cursor --- .../controllers/vscodeNotebookController.ts | 31 +- .../vscodeNotebookController.unit.test.ts | 265 +++++++++++++++--- src/notebooks/deepnote/deepnoteTestHelpers.ts | 32 ++- .../ephemeralCellDecorationProvider.ts | 8 +- 4 files changed, 281 insertions(+), 55 deletions(-) diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index 502a3b1009..824f54d9e1 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -75,6 +75,7 @@ import { IExtensionContext } from '../../platform/common/types'; import { getNotebookMetadata, isJupyterNotebook, updateNotebookMetadata } from '../../platform/common/utils'; +import { notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; import { createDeferred } from '../../platform/common/utils/async'; import { DisposableStore, dispose } from '../../platform/common/utils/lifecycle'; import { Common, DataScience } from '../../platform/common/utils/localize'; @@ -633,21 +634,31 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont const cellsToExecute = await removeEphemeralCellsForAgentBlocks(doc, queuedCells); let pendingKernelCells: NotebookCell[] = []; + let ranAgentCell = false; - for (const cell of cellsToExecute) { - if (!isAgentCell(cell)) { - pendingKernelCells.push(cell); - continue; + try { + for (const cell of cellsToExecute) { + if (!isAgentCell(cell)) { + pendingKernelCells.push(cell); + continue; + } + + ranAgentCell = true; + 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); - pendingKernelCells = []; - - logger.trace(`Executing agent cell ${cell.index} for ${getDisplayPath(doc.uri)} without kernel`); - await executeAgentCell(cell, this.controller).catch(noop); + } finally { + // Batches without an agent cell (including reentrant ephemeral-cell runs) already notify via + // CellExecutionQueue. Explicit notify covers agent-only and agent-then-kernel batches. + if (ranAgentCell) { + notebookCellExecutions.notifyQueueComplete(doc.uri.toString()); + } } - - await this.executeKernelCells(doc, pendingKernelCells); } private async executeKernelCells(doc: NotebookDocument, cells: NotebookCell[]) { diff --git a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts index c03296abb6..3369230a41 100644 --- a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts +++ b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts @@ -7,7 +7,17 @@ /* eslint-disable @typescript-eslint/no-require-imports, @typescript-eslint/no-var-requires */ import { assert } from 'chai'; import * as fakeTimers from '@sinonjs/fake-timers'; -import { NotebookDocument, EventEmitter, NotebookController, Uri, Disposable } from 'vscode'; +import * as sinon from 'sinon'; +import { + Disposable, + EventEmitter, + ExtensionMode, + NotebookController, + NotebookDocument, + SecretStorage, + SecretStorageChangeEvent, + Uri +} from 'vscode'; import { VSCodeNotebookController, warnWhenUsingOutdatedPython } from './vscodeNotebookController'; import { IKernel, @@ -26,6 +36,7 @@ import { IWatchableJupyterSettings } from '../../platform/common/types'; import { dispose } from '../../platform/common/utils/lifecycle'; +import { ServiceContainer } from '../../platform/ioc/container'; import { NotebookCellLanguageService } from '../languages/cellLanguageService'; import { IServiceContainer } from '../../platform/ioc/types'; import { IJupyterServerProviderRegistry } from '../../kernels/jupyter/types'; @@ -43,6 +54,64 @@ import { mockedVSCode, mockedVSCodeNamespaces, resetVSCodeMocks } from '../../te import { Environment, PythonExtension } from '@vscode/python-extension'; import { crateMockedPythonApi, whenResolveEnvironment } from '../../kernels/helpers.unit.test'; import { IJupyterVariablesProvider } from '../../kernels/variables/types'; +import { notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; +import { createMockNotebookWithCells } from '../deepnote/deepnoteTestHelpers'; + +function stubSecretStorageForAgentTests(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(); + }); +} + +function installMockedCreateNotebookController( + onDidChangeSelectedNotebooksEvent: EventEmitter<{ + readonly notebook: NotebookDocument; + readonly selected: boolean; + }>['event'], + createNotebookCellExecution: NotebookController['createNotebookCellExecution'] = () => ({}) as any +): void { + (mockedVSCode as any).notebooks.createNotebookController = ( + _id: string, + _view: string, + _label: string, + executeHandler: any, + _rendererScripts: any + ) => { + return { + id: _id, + label: _label, + description: '', + detail: '', + supportedLanguages: [], + supportsExecutionOrder: false, + interruptHandler: undefined, + executeHandler, + onDidChangeSelectedNotebooks: onDidChangeSelectedNotebooksEvent, + onDidReceiveMessage: new EventEmitter().event, + dispose: () => {}, + asWebviewUri: (uri: Uri) => uri, + postMessage: () => Promise.resolve(true), + updateNotebookAffinity: () => {}, + createNotebookCellExecution, + createNotebookExecution: () => ({}) as any, + notebookType: _view, + rendererScripts: _rendererScripts || [] + } as NotebookController; + }; +} suite(`Notebook Controller`, function () { let controller: NotebookController; @@ -103,42 +172,7 @@ suite(`Notebook Controller`, function () { when(controller.label).thenReturn('Test Controller'); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); when(mockedVSCodeNamespaces.workspace.onDidCloseNotebookDocument).thenReturn(onDidCloseNotebookDocument.event); - // Override just the createNotebookController method on the existing notebooks object - (mockedVSCode as any).notebooks.createNotebookController = ( - _id: string, - _view: string, - _label: string, - _handler: any, - _rendererScripts: any - ) => { - console.log('MOCK createNotebookController CALLED with id:', _id); - const mockControllerObject: any = { - id: _id, - label: _label, - description: '', - detail: '', - supportedLanguages: [], - supportsExecutionOrder: false, - interruptHandler: undefined, - executeHandler: _handler, - onDidChangeSelectedNotebooks: onDidChangeSelectedNotebooks.event, - onDidReceiveMessage: new EventEmitter().event, - dispose: () => {}, - asWebviewUri: (uri: Uri) => uri, - postMessage: () => Promise.resolve(true), - updateNotebookAffinity: () => {}, - createNotebookCellExecution: () => ({}) as any, - createNotebookExecution: () => ({}) as any, - notebookType: _view, - rendererScripts: _rendererScripts || [] - }; - console.log('MOCK createNotebookController RETURNING controller with id:', mockControllerObject.id); - return mockControllerObject; - }; - console.log( - 'mockedVSCode.notebooks.createNotebookController:', - typeof (mockedVSCode as any).notebooks.createNotebookController - ); + installMockedCreateNotebookController(onDidChangeSelectedNotebooks.event); when(languageService.getSupportedLanguages(anything())).thenReturn([PYTHON_LANGUAGE]); when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(true); when(mockedVSCodeNamespaces.workspace.onDidCloseNotebookDocument).thenReturn(onDidCloseNotebookDocument.event); @@ -826,4 +860,161 @@ suite(`Notebook Controller`, function () { assert.isDefined(result); }); }); + + suite('executeQueuedCells', function () { + const secretStorage = new Map(); + + let vscodeController: VSCodeNotebookController; + let notifyQueueCompleteSpy: sinon.SinonSpy; + let createNotebookCellExecutionStub: sinon.SinonStub; + let mockExecution: { + appendOutput: sinon.SinonStub; + clearOutput: sinon.SinonStub; + end: sinon.SinonStub; + replaceOutput: sinon.SinonStub; + appendOutputItems: sinon.SinonStub; + start: sinon.SinonStub; + token: { isCancellationRequested: boolean }; + }; + + setup(function () { + crateMockedPythonApi(disposables); + secretStorage.clear(); + secretStorage.set('openAiApiKey', 'test-key'); + stubSecretStorageForAgentTests(secretStorage); + when(serviceContainer.tryGet(anything())).thenReturn(undefined); + + mockExecution = { + appendOutput: sinon.stub().resolves(), + clearOutput: sinon.stub().resolves(), + end: sinon.stub(), + replaceOutput: sinon.stub().resolves(), + appendOutputItems: sinon.stub().resolves(), + start: sinon.stub(), + token: { isCancellationRequested: false } + }; + createNotebookCellExecutionStub = sinon.stub().callsFake(() => { + mockExecution.end = sinon.stub(); + + return mockExecution; + }); + + installMockedCreateNotebookController(onDidChangeSelectedNotebooks.event, createNotebookCellExecutionStub); + + notifyQueueCompleteSpy = sinon.spy(notebookCellExecutions, 'notifyQueueComplete'); + + vscodeController = new VSCodeNotebookController( + instance(kernelConnection), + 'test-controller-id', + 'jupyter-notebook', + instance(kernelProvider), + instance(context), + disposables, + instance(languageService), + instance(configService), + instance(extensionChecker), + instance(serviceContainer), + displayDataProvider + ); + }); + + teardown(function () { + notifyQueueCompleteSpy.restore(); + sinon.restore(); + }); + + test('agent-only batch fires notifyQueueComplete (arms deferred snapshot save)', async function () { + // Catches: agent-only runs never reach CellExecutionQueue, so snapshot save never arms. + const { + notebook: agentNotebook, + cells: [agentCell] + } = createMockNotebookWithCells([ + { + metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, + text: 'Test prompt' + } + ]); + + const notebookUri = agentNotebook.uri.toString(); + + let queueCompletionUri: string | undefined; + const queueListener = notebookCellExecutions.onDidCompleteQueueExecution((event) => { + queueCompletionUri = event.notebookUri; + }); + disposables.push(new Disposable(() => queueListener.dispose())); + + const executeHandler = vscodeController.controller.executeHandler; + assert.isDefined(executeHandler); + + await executeHandler([agentCell], agentNotebook, vscodeController.controller); + + assert.isTrue(notifyQueueCompleteSpy.calledOnce, 'notifyQueueComplete must run after agent-only execution'); + assert.strictEqual(notifyQueueCompleteSpy.firstCall.args[0], notebookUri); + assert.strictEqual(queueCompletionUri, notebookUri); + assert.isTrue(createNotebookCellExecutionStub.calledOnce, 'agent cell should run through executeAgentCell'); + }); + + test('kernel-only batch after agent batch does not fire a second explicit queue completion notify', async function () { + // Catches: explicit notify on kernel-only executeQueuedCells (e.g. reentrant ephemeral runs) when ranAgentCell is false. + const { + notebook: agentNotebook, + cells: [agentCell, codeCell] + } = createMockNotebookWithCells([ + { + metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, + text: 'Test prompt' + }, + { + metadata: { id: 'ephemeral-code-1' }, + text: 'print(1)' + } + ]); + + const executeHandler = vscodeController.controller.executeHandler; + assert.isDefined(executeHandler); + + notifyQueueCompleteSpy.resetHistory(); + + await executeHandler([agentCell], agentNotebook, vscodeController.controller); + + try { + await executeHandler([codeCell], agentNotebook, vscodeController.controller); + } catch { + // Kernel harness may not fully mock cell execution startup. + } + + assert.strictEqual( + notifyQueueCompleteSpy.callCount, + 1, + 'only the agent batch should fire explicit queue completion' + ); + }); + + test('kernel-only batch does not fire explicit queue completion notify', async function () { + // Catches: explicit notify on pure kernel batches that already signal via CellExecutionQueue. + const { + notebook: codeNotebook, + cells: [codeCell] + } = createMockNotebookWithCells([ + { + metadata: { id: 'code-block-1' }, + text: 'x = 1' + } + ]); + + const executeHandler = vscodeController.controller.executeHandler; + assert.isDefined(executeHandler); + + try { + await executeHandler([codeCell], codeNotebook, vscodeController.controller); + } catch { + // Kernel harness may not fully mock cell execution startup. + } + + assert.isFalse( + notifyQueueCompleteSpy.called, + 'batches without agent cells must rely on CellExecutionQueue for completion notify' + ); + }); + }); }); diff --git a/src/notebooks/deepnote/deepnoteTestHelpers.ts b/src/notebooks/deepnote/deepnoteTestHelpers.ts index 6c6b766bcf..4abbc16b4c 100644 --- a/src/notebooks/deepnote/deepnoteTestHelpers.ts +++ b/src/notebooks/deepnote/deepnoteTestHelpers.ts @@ -33,6 +33,7 @@ export interface CreateMockCellOptions { notebookMetadata?: Record; index?: number; mime?: string; + notebook?: NotebookDocument; } /** @@ -86,6 +87,22 @@ export function createMockNotebook(options?: CreateMockNotebookOptions): Noteboo } satisfies NotebookDocument; } +/** + * Builds one mock notebook and cells that share it (correct `index` and `notebook` references). + */ +export function createMockNotebookWithCells( + cellOptions: Omit[] +): { cells: NotebookCell[]; notebook: NotebookDocument } { + const cells: NotebookCell[] = []; + const notebook = createMockNotebook({ cells }); + + for (let index = 0; index < cellOptions.length; index++) { + cells.push(createMockCell({ ...cellOptions[index], index, notebook })); + } + + return { cells, notebook }; +} + /** * Creates a mock NotebookCellOutput for testing. * @@ -132,13 +149,16 @@ export function createMockCell(options?: CreateMockCellOptions): NotebookCell { ? opts.notebookMetadata : {}; - const notebook = createMockNotebook({ - notebookType, - uri: notebookUri, - metadata: notebookMetadata - }); + const notebook = + opts.notebook ?? + createMockNotebook({ + notebookType, + uri: notebookUri, + metadata: notebookMetadata + }); + const resolvedUri = notebook.uri; - const cellPath = `${notebookUri.path}#cell${index}`; + const cellPath = `${resolvedUri.path}#cell${index}`; const document: TextDocument = { uri: Uri.file(cellPath), diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts index 30878e73a7..19c3aa8470 100644 --- a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts @@ -13,6 +13,7 @@ import { import { injectable } from 'inversify'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { logger } from '../../platform/logging'; import { isEphemeralCell } from './dataConversionUtils'; const NOTEBOOK_CELL_SCHEME = 'vscode-notebook-cell'; @@ -113,8 +114,11 @@ export class EphemeralCellDecorationProvider implements IExtensionSyncActivation } editor.setDecorations(this.ephemeralDecorationType, lineRanges); - } catch { - continue; + } catch (error) { + logger.warn( + `EphemeralCellDecorationProvider: Failed to update decorations for ${editor.document.uri.path}`, + error + ); } } } From 85b36e6ee4bda55c8306940ee3c78588a5719881 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 6 Aug 2026 18:31:47 +0000 Subject: [PATCH 45/80] fix(agent-block): resolve snapshot blocks by file index, not live cell index MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit executeSnapshotOutputUpdate fell back to originalBlocks[i] when a cell had lost its block-id metadata, indexing the main file's block list by the live cell index. Ephemeral agent cells exist in the document but are stripped from the file, so every cell below one is offset: the metadata-less cell adopted the wrong block's outputs, and blockIdFromFallback wrote that id back into the cell — leaving two cells claiming one block, which the next save persists as duplicate block ids. Track a separate cursor into the file's block list, advanced only by persisted cells. The cursor increments before any continue so a skipped cell cannot desync it. cellIndex stays the live index: it feeds NotebookEdit ranges, which are live-document coordinates. Ephemeral cells still resolve by their own id when they have one — they are excluded only from the index fallback. contentActuallyChanged already filtered ephemeral cells; this loop did not. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- .../deepnote/deepnoteFileChangeWatcher.ts | 11 +- .../deepnoteFileChangeWatcher.unit.test.ts | 152 +++++++++++++++++- 2 files changed, 159 insertions(+), 4 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index 102b4dfaf9..5a304abc81 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -375,15 +375,20 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic blockIdFromFallback: boolean; }> = []; + // Ephemeral cells live in the document but are stripped from the file, so a live index cannot + // address originalBlocks. Track the file's own cursor, advanced only by persisted cells. + let fileBlockIndex = 0; + for (let i = 0; i < liveCells.length; i++) { try { const cell = liveCells[i]; + const originalBlock = isEphemeralCell(cell) ? undefined : originalBlocks?.[fileBlockIndex++]; let blockId = getBlockId(cell); let blockIdFromFallback = false; // Fallback to original project blocks when metadata was lost - if (!blockId && originalBlocks) { - blockId = originalBlocks[i]?.id; + if (!blockId && originalBlock) { + blockId = originalBlock.id; blockIdFromFallback = true; } @@ -391,7 +396,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic continue; } - const fallbackType = originalBlocks?.[i]?.type; + const fallbackType = originalBlock?.type; const blockType = ((cell.metadata?.type as string) ?? fallbackType ?? 'code') as DeepnoteBlock['type']; const newOutputs = this.converter.transformOutputsForVsCode( snapshotOutputs.get(blockId)!, diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index e45ca19899..58a096d1db 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -2,7 +2,15 @@ import type { DeepnoteFile } from '@deepnote/blocks'; import { assert } from 'chai'; import * as sinon from 'sinon'; import { anything, instance, mock, when } from 'ts-mockito'; -import { Disposable, EventEmitter, FileSystemWatcher, NotebookCellKind, NotebookDocument, Uri } from 'vscode'; +import { + Disposable, + EventEmitter, + FileSystemWatcher, + NotebookCellKind, + NotebookDocument, + NotebookEdit, + Uri +} from 'vscode'; import type { IControllerRegistration } from '../controllers/types'; import type { IDisposableRegistry } from '../../platform/common/types'; @@ -1169,6 +1177,148 @@ project: fallbackOnDidCreate.dispose(); }); + test('should resolve the block for a metadata-less cell that sits below an ephemeral cell', async () => { + // Catches: indexing the main file's block list by live cell index. Ephemeral agent cells + // exist in the document but are stripped from the file, so every cell below one is offset — + // the metadata-less cell would adopt the wrong block's id and outputs, and that id gets + // written back, leaving two cells claiming one block. + const mockedManager = mock(); + when(mockedManager.getProjectForNotebook('e132b172-b114-410e-8331-011517db664f', 'notebook-1')).thenReturn({ + version: '1.0', + metadata: { createdAt: '2025-01-01T00:00:00Z' }, + project: { + id: 'e132b172-b114-410e-8331-011517db664f', + name: 'Test Project', + notebooks: [ + { + id: 'notebook-1', + name: 'Notebook 1', + blocks: [ + { id: 'block-1', type: 'code', sortingKey: 'a0' }, + { id: 'block-2', type: 'code', sortingKey: 'a1' } + ] + } + ] + } + } as DeepnoteFile); + + const offsetDisposables: IDisposableRegistry = []; + const offsetOnDidChange = new EventEmitter(); + const offsetOnDidCreate = new EventEmitter(); + const offsetFsWatcher = mock(); + when(offsetFsWatcher.onDidChange).thenReturn(offsetOnDidChange.event); + when(offsetFsWatcher.onDidCreate).thenReturn(offsetOnDidCreate.event); + when(offsetFsWatcher.dispose()).thenReturn(); + when(mockedVSCodeNamespaces.workspace.createFileSystemWatcher(anything())).thenReturn( + instance(offsetFsWatcher) + ); + + let offsetApplyEditCount = 0; + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { + offsetApplyEditCount++; + return Promise.resolve(true); + }); + + // The vscode mock's NotebookEdit.updateCellMetadata drops its metadata argument, so the + // WorkspaceEdit cannot reveal which block id was written. Stub the static to capture it. + const metadataWrites: Array<{ index: number; metadata: Record }> = []; + sinon.stub(NotebookEdit, 'updateCellMetadata').callsFake((index: number, metadata) => { + metadataWrites.push({ index, metadata: metadata as Record }); + + return {} as NotebookEdit; + }); + + const offsetWatcher = new DeepnoteFileChangeWatcher( + offsetDisposables, + instance(mockedManager), + instance(mockSnapshotService) + ); + offsetWatcher.activate(); + + const notebook = createMockNotebook({ + uri: Uri.file('/workspace/test.deepnote'), + metadata: { + deepnoteProjectId: 'e132b172-b114-410e-8331-011517db664f', + deepnoteNotebookId: 'notebook-1' + }, + cells: [ + { + metadata: { id: 'block-1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("first")' } + }, + { + // Agent scratch cell: lives in the document, never persisted to the file + metadata: { id: 'eph-1', type: 'code', is_ephemeral: true, agent_source_block_id: 'agent-1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("scratch")' } + }, + { + metadata: { type: 'code' }, // No id — VS Code lost it + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("second")' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + const newOutputs = new Map([ + [ + 'block-1', + [ + { + output_type: 'execute_result', + data: { 'text/plain': 'First Output' }, + execution_count: 1 + } as DeepnoteOutput + ] + ], + [ + 'block-2', + [ + { + output_type: 'execute_result', + data: { 'text/plain': 'Second Output' }, + execution_count: 2 + } as DeepnoteOutput + ] + ] + ]); + when(mockSnapshotService.readSnapshot(anything(), anything())).thenReturn(Promise.resolve(newOutputs)); + + offsetOnDidChange.fire( + Uri.file( + '/workspace/snapshots/my-project_e132b172-b114-410e-8331-011517db664f_latest.snapshot.deepnote' + ) + ); + + await waitFor(() => offsetApplyEditCount > 0); + await waitFor(() => metadataWrites.length > 0); + + const writeForLastCell = metadataWrites.find((w) => w.index === 2); + assert.isDefined(writeForLastCell, 'the metadata-less cell at live index 2 should receive a block id'); + assert.strictEqual( + writeForLastCell!.metadata.__deepnoteBlockId, + 'block-2', + 'live index 2 is the second *persisted* cell, so it must resolve to block-2' + ); + assert.notStrictEqual( + writeForLastCell!.metadata.__deepnoteBlockId, + 'block-1', + 'block-1 already belongs to live index 0 — two cells must never claim one block' + ); + + for (const d of offsetDisposables) { + d.dispose(); + } + offsetOnDidChange.dispose(); + offsetOnDidCreate.dispose(); + }); + test('should only update cells whose outputs changed (per-cell updates)', async () => { const snapshotUri = Uri.file( '/workspace/snapshots/my-project_e132b172-b114-410e-8331-011517db664f_latest.snapshot.deepnote' From de94d14e010936183d6f7fc9124749b773b83b89 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 6 Aug 2026 18:31:58 +0000 Subject: [PATCH 46/80] test(agent-block): cover the agent model switch write path MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit switchModel/updateCellMetadata is the only path that persists deepnote_agent_model, and it had no coverage in either the unit suite or E2E — the existing tests all exercise provideCellStatusBarItems, the read side. Five tests, each verified to fail against a deliberate mutation of the code it guards: inverted metadata spread, each half of the early-return guard dropped, the applyEdit failure branch removed, and the isAgentCell guard removed. Every mutation killed exactly one test. The inverted spread is the reason this matters: it turns the model switch into a silent no-op while still calling applyEdit, so a call-count assertion would not notice. Asserting the written metadata is what catches it. Metadata is captured by stubbing the NotebookEdit.updateCellMetadata static rather than WorkspaceEdit.prototype.set — the vscode mock's implementation discards its metadata argument, so the WorkspaceEdit cannot show what was written. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- .../agentCellStatusBarProvider.unit.test.ts | 111 +++++++++++++++++- 1 file changed, 110 insertions(+), 1 deletion(-) diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts index faaad4604c..90bd9e17db 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -1,6 +1,9 @@ import { expect } from 'chai'; -import { CancellationToken } from 'vscode'; +import * as sinon from 'sinon'; +import { anything, verify, when } from 'ts-mockito'; +import { CancellationToken, NotebookCell, NotebookEdit } from 'vscode'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; import { AgentCellStatusBarProvider } from './agentCellStatusBarProvider'; import { createMockCell } from './deepnoteTestHelpers'; @@ -127,4 +130,110 @@ suite('AgentCellStatusBarProvider', () => { expect(items[1].text).to.include('Model: auto'); }); }); + + suite('Model Switching', () => { + let capturedEdit: { index: number; metadata: Record } | undefined; + + setup(() => { + resetVSCodeMocks(); + capturedEdit = undefined; + + // The vscode mock's NotebookEdit.updateCellMetadata discards its metadata argument, so + // capturing the WorkspaceEdit cannot show what was written. Stub the static instead. + sinon.stub(NotebookEdit, 'updateCellMetadata').callsFake((index: number, metadata) => { + capturedEdit = { index, metadata: metadata as Record }; + + return {} as NotebookEdit; + }); + }); + + teardown(() => { + sinon.restore(); + resetVSCodeMocks(); + }); + + function agentCell(): NotebookCell { + return createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent', id: 'pocket-1' }, + id: 'block-1', + deepnote_agent_model: 'gpt-4o' + }, + index: 2 + }); + } + + function pick(label: string | undefined) { + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( + Promise.resolve(label === undefined ? undefined : ({ label } as any)) + ); + } + + function switchModel(cell: NotebookCell): Promise { + return (provider as unknown as { switchModel(cell: NotebookCell): Promise }).switchModel(cell); + } + + test('Should write the picked model without dropping the cell’s other metadata', async () => { + // Catches: an inverted spread in updateCellMetadata, which makes the switch a silent + // no-op while still calling applyEdit — so a call-count assertion would not notice. + pick('gpt-5'); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(true)); + + await switchModel(agentCell()); + + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).once(); + expect(capturedEdit!.index).to.equal(2); + expect(capturedEdit!.metadata).to.deep.equal({ + __deepnotePocket: { type: 'agent', id: 'pocket-1' }, + id: 'block-1', + deepnote_agent_model: 'gpt-5' + }); + }); + + test('Should not edit the notebook when the current model is re-picked', async () => { + // Catches: losing the `selected.label === currentModel` guard, which dirties the + // document on a no-op selection. + pick('gpt-4o'); + + await switchModel(agentCell()); + + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('Should not edit the notebook when the picker is dismissed', async () => { + // Catches: losing the `!selected` guard, which throws on `selected.label` when the + // user presses Escape. + pick(undefined); + + await switchModel(agentCell()); + + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('Should report an error when the workspace edit is rejected', async () => { + // Catches: dropping the `if (!success)` branch, which loses the model change silently. + pick('gpt-5'); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(false)); + + let statusBarRefreshed = false; + provider.onDidChangeCellStatusBarItems(() => { + statusBarRefreshed = true; + }); + + await switchModel(agentCell()); + + verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); + expect(statusBarRefreshed, 'a rejected edit must not refresh the status bar').to.be.false; + }); + + test('Should ignore a non-agent cell', async () => { + // Catches: dropping the isAgentCell guard, which would offer the model picker on any cell. + pick('gpt-5'); + + await switchModel(createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } })); + + verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + }); }); From bba37d39d5b995b5e7cb80f8ed66b7fd2f5895e3 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 7 Aug 2026 10:26:46 +0000 Subject: [PATCH 47/80] chore(agent-block): update the agent model slugs Replaces gpt-4o/gpt-5 with the gpt-5.6 sol/terra/luna variants in the model picker, its tests, and the E2E fixture. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- .../deepnote/agentCellStatusBarProvider.ts | 2 +- .../agentCellStatusBarProvider.unit.test.ts | 16 ++++++++-------- .../converters/agentBlockConverter.unit.test.ts | 4 ++-- test/e2e/fixtures/agent-block.deepnote | 2 +- 4 files changed, 12 insertions(+), 12 deletions(-) diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index 96d6aeac55..4a4b3bed2f 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -24,7 +24,7 @@ const AGENT_MODEL_METADATA_KEY = 'deepnote_agent_model'; /** Persisted default — absent key becomes `undefined` and breaks openai() model selection. */ const AGENT_MODEL_AUTO = 'auto'; -const AGENT_MODEL_OPTIONS = [AGENT_MODEL_AUTO, 'gpt-4o', 'gpt-5']; +const AGENT_MODEL_OPTIONS = [AGENT_MODEL_AUTO, 'gpt-5.6-sol', 'gpt-5.6-terra', 'gpt-5.6-luna']; const AGENT_INDICATOR_PRIORITY = 100; const MODEL_PICKER_PRIORITY = 90; diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts index 90bd9e17db..46a0cc4c47 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -109,13 +109,13 @@ suite('AgentCellStatusBarProvider', () => { const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' }, - deepnote_agent_model: 'gpt-4o', + deepnote_agent_model: 'gpt-5.6-sol', deepnote_max_iterations: 50 } }); const items = provider.provideCellStatusBarItems(cell, mockToken)!; - expect(items[1].text).to.include('Model: gpt-4o'); + expect(items[1].text).to.include('Model: gpt-5.6-sol'); }); test('Should display "auto" when model is empty string', () => { @@ -157,7 +157,7 @@ suite('AgentCellStatusBarProvider', () => { metadata: { __deepnotePocket: { type: 'agent', id: 'pocket-1' }, id: 'block-1', - deepnote_agent_model: 'gpt-4o' + deepnote_agent_model: 'gpt-5.6-sol' }, index: 2 }); @@ -176,7 +176,7 @@ suite('AgentCellStatusBarProvider', () => { test('Should write the picked model without dropping the cell’s other metadata', async () => { // Catches: an inverted spread in updateCellMetadata, which makes the switch a silent // no-op while still calling applyEdit — so a call-count assertion would not notice. - pick('gpt-5'); + pick('gpt-5.6-terra'); when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(true)); await switchModel(agentCell()); @@ -186,14 +186,14 @@ suite('AgentCellStatusBarProvider', () => { expect(capturedEdit!.metadata).to.deep.equal({ __deepnotePocket: { type: 'agent', id: 'pocket-1' }, id: 'block-1', - deepnote_agent_model: 'gpt-5' + deepnote_agent_model: 'gpt-5.6-terra' }); }); test('Should not edit the notebook when the current model is re-picked', async () => { // Catches: losing the `selected.label === currentModel` guard, which dirties the // document on a no-op selection. - pick('gpt-4o'); + pick('gpt-5.6-sol'); await switchModel(agentCell()); @@ -212,7 +212,7 @@ suite('AgentCellStatusBarProvider', () => { test('Should report an error when the workspace edit is rejected', async () => { // Catches: dropping the `if (!success)` branch, which loses the model change silently. - pick('gpt-5'); + pick('gpt-5.6-luna'); when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(false)); let statusBarRefreshed = false; @@ -228,7 +228,7 @@ suite('AgentCellStatusBarProvider', () => { test('Should ignore a non-agent cell', async () => { // Catches: dropping the isAgentCell guard, which would offer the model picker on any cell. - pick('gpt-5'); + pick('gpt-5.6-luna'); await switchModel(createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } })); diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts index 1a2c46337f..65cc4cc488 100644 --- a/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts +++ b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts @@ -105,7 +105,7 @@ suite('AgentBlockConverter', () => { content: 'Old prompt', id: 'agent-block-789', metadata: { - deepnote_agent_model: 'gpt-4o', + deepnote_agent_model: 'gpt-5.6-sol', custom: 'value' }, sortingKey: 'a2', @@ -120,7 +120,7 @@ suite('AgentBlockConverter', () => { assert.strictEqual(block.type, 'agent'); assert.strictEqual(block.sortingKey, 'a2'); assert.deepStrictEqual(block.metadata, { - deepnote_agent_model: 'gpt-4o', + deepnote_agent_model: 'gpt-5.6-sol', custom: 'value' }); }); diff --git a/test/e2e/fixtures/agent-block.deepnote b/test/e2e/fixtures/agent-block.deepnote index 334562ffec..c5fed3056c 100644 --- a/test/e2e/fixtures/agent-block.deepnote +++ b/test/e2e/fixtures/agent-block.deepnote @@ -16,7 +16,7 @@ project: Run some Python, then add a markdown block summarising this notebook. sortingKey: a0 metadata: - deepnote_agent_model: gpt-5 + deepnote_agent_model: 'gpt-5.6-sol' executionMode: block isModule: false settings: {} From 418c57ae979616dfed1ffeec118eea8cdd9e306e Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 7 Aug 2026 10:27:00 +0000 Subject: [PATCH 48/80] refactor(agent-block): keep the OpenAI key in IEncryptedStorage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The secret store reached for ServiceContainer.instance because its callers are free functions with no injection point. IEncryptedStorage already wraps extensionContext.secrets and is registered in both the node and web registries, so the generic getSecret/setSecret/clearSecret layer was a duplicate of it with exactly one key ever passed through. Collapse the module to the OpenAI-key functions, which take the storage as a parameter. The command handler injects it; executeAgentCell receives it from the controller's already-injected IServiceContainer, so nothing new is threaded through controllerRegistration. Two behaviour changes fall out. Secrets are namespaced as `deepnote-agent.openAiApiKey` now that EncryptedStorage composes the key — no migration, this has never shipped. And the ExtensionMode.Test bail is gone: it silently turned writes into no-ops under the in-process test harness, where EncryptedStorage keeps a working in-memory store instead. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01LaXoas2nAryMghEk3FQS2o --- .../controllers/vscodeNotebookController.ts | 7 +- .../vscodeNotebookController.unit.test.ts | 44 +--- .../deepnote/agentCellExecutionHandler.ts | 4 +- .../agentCellExecutionHandler.unit.test.ts | 89 +++++--- .../agentOpenAiApiKeyCommandHandler.ts | 10 +- src/notebooks/deepnote/deepnoteSecretStore.ts | 119 ++-------- .../deepnote/deepnoteSecretStore.unit.test.ts | 205 +++++++----------- 7 files changed, 187 insertions(+), 291 deletions(-) diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index 824f54d9e1..47e1051ca4 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -57,6 +57,7 @@ import { import { IJupyterVariablesProvider } from '../../kernels/variables/types'; import { IPyWidgetMessages } from '../../messageTypes'; import { IPythonExtensionChecker } from '../../platform/api/types'; +import { IEncryptedStorage } from '../../platform/common/application/types'; import { isCancellationError } from '../../platform/common/cancellation'; import { Commands, @@ -648,7 +649,11 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont pendingKernelCells = []; logger.trace(`Executing agent cell ${cell.index} for ${getDisplayPath(doc.uri)} without kernel`); - await executeAgentCell(cell, this.controller).catch(noop); + await executeAgentCell( + cell, + this.controller, + this.serviceContainer.get(IEncryptedStorage) + ).catch(noop); } await this.executeKernelCells(doc, pendingKernelCells); diff --git a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts index 3369230a41..669a20752a 100644 --- a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts +++ b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts @@ -8,16 +8,7 @@ import { assert } from 'chai'; import * as fakeTimers from '@sinonjs/fake-timers'; import * as sinon from 'sinon'; -import { - Disposable, - EventEmitter, - ExtensionMode, - NotebookController, - NotebookDocument, - SecretStorage, - SecretStorageChangeEvent, - Uri -} from 'vscode'; +import { Disposable, EventEmitter, NotebookController, NotebookDocument, Uri } from 'vscode'; import { VSCodeNotebookController, warnWhenUsingOutdatedPython } from './vscodeNotebookController'; import { IKernel, @@ -29,6 +20,7 @@ import { RemoteKernelSpecConnectionMetadata } from '../../kernels/types'; import { anything, instance, mock, verify, when } from 'ts-mockito'; +import { IEncryptedStorage } from '../../platform/common/application/types'; import { IConfigurationService, IDisposable, @@ -57,23 +49,15 @@ import { IJupyterVariablesProvider } from '../../kernels/variables/types'; import { notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; import { createMockNotebookWithCells } from '../deepnote/deepnoteTestHelpers'; -function stubSecretStorageForAgentTests(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(); - }); +// executeAgentCell takes IEncryptedStorage from the controller's container; getProjectAgentContext +// still resolves the notebook manager off the static one. +function stubAgentDependencies(serviceContainer: IServiceContainer, openAiApiKey: string): void { + const encryptedStorage = mock(); + const staticServiceContainer = instance(mock()); + + when(encryptedStorage.retrieve(anything(), anything())).thenResolve(openAiApiKey); + when(serviceContainer.get(IEncryptedStorage)).thenReturn(instance(encryptedStorage)); + sinon.stub(ServiceContainer, 'instance').get(() => staticServiceContainer); } function installMockedCreateNotebookController( @@ -862,8 +846,6 @@ suite(`Notebook Controller`, function () { }); suite('executeQueuedCells', function () { - const secretStorage = new Map(); - let vscodeController: VSCodeNotebookController; let notifyQueueCompleteSpy: sinon.SinonSpy; let createNotebookCellExecutionStub: sinon.SinonStub; @@ -879,9 +861,7 @@ suite(`Notebook Controller`, function () { setup(function () { crateMockedPythonApi(disposables); - secretStorage.clear(); - secretStorage.set('openAiApiKey', 'test-key'); - stubSecretStorageForAgentTests(secretStorage); + stubAgentDependencies(serviceContainer, 'test-key'); when(serviceContainer.tryGet(anything())).thenReturn(undefined); mockExecution = { diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index 58d1d562a8..748aecbc79 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -23,6 +23,7 @@ import { } from '@deepnote/runtime-core'; import { translateCellDisplayOutput } from '../../kernels/execution/helpers'; +import { IEncryptedStorage } from '../../platform/common/application/types'; import type { IDisposable } from '../../platform/common/types'; import { createDeferred } from '../../platform/common/utils/async'; import { dispose } from '../../platform/common/utils/lifecycle'; @@ -159,6 +160,7 @@ export interface ExecuteAgentCellOptions { export async function executeAgentCell( cell: NotebookCell, controller: NotebookController, + encryptedStorage: IEncryptedStorage, options?: ExecuteAgentCellOptions ): Promise { const executeAgentBlockFn = options?.executeAgentBlockFn ?? executeAgentBlock; @@ -192,7 +194,7 @@ export async function executeAgentCell( ); } - const openAiToken = await getOrPromptOpenAiApiKey(); + const openAiToken = await getOrPromptOpenAiApiKey(encryptedStorage); let lastAgentEventType: AgentStreamEvent['type'] | undefined; diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index f94fccf653..18c0b96ad7 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -5,24 +5,20 @@ import { CancellationError, CancellationTokenSource, Disposable, - EventEmitter, - ExtensionMode, NotebookCell, NotebookCellData, NotebookCellOutput, NotebookCellOutputItem, NotebookController, NotebookDocument, - SecretStorage, - SecretStorageChangeEvent, WorkspaceEdit } from 'vscode'; import type { AgentBlock } from '@deepnote/blocks'; import type { AgentBlockContext } from '@deepnote/runtime-core'; +import { IEncryptedStorage } from '../../platform/common/application/types'; import type { IDisposable } from '../../platform/common/types'; -import { IExtensionContext } from '../../platform/common/types'; import { dispose } from '../../platform/common/utils/lifecycle'; import { ServiceContainer } from '../../platform/ioc/container'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; @@ -37,24 +33,33 @@ import { import { IDeepnoteNotebookManager } from '../types'; import { createDeepnoteFile, createDeepnoteProject, createMockCell, createMockNotebook } from './deepnoteTestHelpers'; -// ExtensionMode.Test skips secrets; Production + in-memory store exercises real paths. -function stubSecretStorage(secretStorage: Map): ServiceContainer { - const context = mock(); - const secrets = mock(); - const onDidChangeSecrets = new EventEmitter(); +// Key namespacing is EncryptedStorage's job and is covered in deepnoteSecretStore.unit.test.ts. +function createEncryptedStorageFake(secretStorage: Map): IEncryptedStorage { + const encryptedStorage = mock(); + + when(encryptedStorage.store(anything(), anything(), anything())).thenCall( + (_service: string, key: string, value: string | undefined) => { + if (value === undefined) { + secretStorage.delete(key); + } else { + secretStorage.set(key, value); + } + + return Promise.resolve(); + } + ); + when(encryptedStorage.retrieve(anything(), anything())).thenCall((_service: string, key: string) => + Promise.resolve(secretStorage.get(key)) + ); + + return instance(encryptedStorage); +} + +// getProjectAgentContext still resolves the notebook manager off the static container. +function stubServiceContainerInstance(): ServiceContainer { 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(); - }); return serviceContainer; } @@ -167,11 +172,13 @@ suite('AgentCellExecutionHandler', () => { let mockController: NotebookController; let executeAgentBlockStub: sinon.SinonStub; let mockServiceContainer: ServiceContainer; + let encryptedStorage: IEncryptedStorage; setup(() => { secretStorage.clear(); secretStorage.set('openAiApiKey', 'test-key'); - mockServiceContainer = stubSecretStorage(secretStorage); + encryptedStorage = createEncryptedStorageFake(secretStorage); + mockServiceContainer = stubServiceContainerInstance(); disposables.push(new Disposable(() => sinon.restore())); mockExecution = { @@ -230,7 +237,9 @@ suite('AgentCellExecutionHandler', () => { test('creates execution, clears output, sets planning output, and ends successfully', async () => { const cell = createAgentCell('Analyze data'); - await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + await executeAgentCell(cell, mockController, encryptedStorage, { + executeAgentBlockFn: executeAgentBlockStub + }); expect((mockController.createNotebookCellExecution as sinon.SinonStub).calledOnceWith(cell)).to.be.true; expect(mockExecution.start.calledOnce).to.be.true; @@ -259,7 +268,9 @@ suite('AgentCellExecutionHandler', () => { const cell = createAgentCell(); - await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + await executeAgentCell(cell, mockController, encryptedStorage, { + executeAgentBlockFn: executeAgentBlockStub + }); expect(mockExecution.appendOutputItems.callCount).to.equal(2); @@ -279,7 +290,9 @@ suite('AgentCellExecutionHandler', () => { const cell = createAgentCell(); - await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + await executeAgentCell(cell, mockController, encryptedStorage, { + executeAgentBlockFn: executeAgentBlockStub + }); const chunk2 = getStdoutChunkText(1); expect(chunk2).to.include('\n\n'); @@ -291,7 +304,9 @@ suite('AgentCellExecutionHandler', () => { const cell = createAgentCell(); - await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + await executeAgentCell(cell, mockController, encryptedStorage, { + executeAgentBlockFn: executeAgentBlockStub + }); expect(mockExecution.end.calledOnce).to.be.true; expect(mockExecution.end.firstCall.args[0]).to.be.false; @@ -310,7 +325,9 @@ suite('AgentCellExecutionHandler', () => { test('handles empty prompt', async () => { const cell = createAgentCell(''); - await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + await executeAgentCell(cell, mockController, encryptedStorage, { + executeAgentBlockFn: executeAgentBlockStub + }); expect(mockExecution.end.calledOnce).to.be.true; expect(mockExecution.end.firstCall.args[0]).to.be.true; @@ -326,7 +343,9 @@ suite('AgentCellExecutionHandler', () => { const cell = createAgentCell(); - await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + await executeAgentCell(cell, mockController, encryptedStorage, { + executeAgentBlockFn: executeAgentBlockStub + }); expect(mockExecution.end.calledOnce).to.be.true; expect(mockExecution.end.firstCall.args[0]).to.be.false; @@ -346,7 +365,9 @@ suite('AgentCellExecutionHandler', () => { }); const { agentCell, cells } = createAgentCellInMutableNotebook([previousResult]); - await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + await executeAgentCell(agentCell, mockController, encryptedStorage, { + executeAgentBlockFn: executeAgentBlockStub + }); expect(executeAgentBlockStub.called).to.be.false; expect(mockExecution.end.firstCall.args[0]).to.be.false; @@ -367,7 +388,9 @@ suite('AgentCellExecutionHandler', () => { return { finalOutput: '' }; }); - await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + await executeAgentCell(agentCell, mockController, encryptedStorage, { + executeAgentBlockFn: executeAgentBlockStub + }); expect(cells.map((cell) => cell.document.getText())).to.deep.equal(['Test prompt', 'first', 'second']); expect(cells[1].metadata?.is_ephemeral).to.be.true; @@ -390,7 +413,9 @@ suite('AgentCellExecutionHandler', () => { return { finalOutput: '' }; }); - await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + await executeAgentCell(agentCell, mockController, encryptedStorage, { + executeAgentBlockFn: executeAgentBlockStub + }); expect(toolResult).to.include('Execution error'); expect(toolResult).to.include('Failed to insert ephemeral code cell'); @@ -415,7 +440,9 @@ suite('AgentCellExecutionHandler', () => { notebookMetadata: { deepnoteProjectId: 'project-1', deepnoteNotebookId: 'notebook-1' } }); - await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + await executeAgentCell(cell, mockController, encryptedStorage, { + executeAgentBlockFn: executeAgentBlockStub + }); const context = executeAgentBlockStub.firstCall.args[1] as AgentBlockContext; expect(context.mcpServers).to.deep.equal(mcpServers); diff --git a/src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts b/src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts index 09ada758df..8ed762b1b1 100644 --- a/src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts +++ b/src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts @@ -2,12 +2,16 @@ import { inject, injectable } from 'inversify'; import { commands, l10n, window } from 'vscode'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { IEncryptedStorage } from '../../platform/common/application/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) {} + constructor( + @inject(IEncryptedStorage) private readonly encryptedStorage: IEncryptedStorage, + @inject(IExtensionContext) private readonly extensionContext: IExtensionContext + ) {} public activate(): void { this.extensionContext.subscriptions.push( @@ -17,14 +21,14 @@ export class AgentOpenAiApiKeyCommandHandler implements IExtensionSyncActivation } private async setApiKey(): Promise { - const key = await promptForOpenAiApiKey(); + const key = await promptForOpenAiApiKey(this.encryptedStorage); if (key) { void window.showInformationMessage(l10n.t('OpenAI API key has been saved.')); } } private async clearApiKey(): Promise { - await clearOpenAiApiKey(); + await clearOpenAiApiKey(this.encryptedStorage); void window.showInformationMessage(l10n.t('OpenAI API key has been cleared.')); } } diff --git a/src/notebooks/deepnote/deepnoteSecretStore.ts b/src/notebooks/deepnote/deepnoteSecretStore.ts index 9e940557f6..7a4844b51f 100644 --- a/src/notebooks/deepnote/deepnoteSecretStore.ts +++ b/src/notebooks/deepnote/deepnoteSecretStore.ts @@ -1,120 +1,47 @@ -import { ExtensionMode, l10n, window } from 'vscode'; +import { l10n, window } from 'vscode'; -import { ServiceContainer } from '../../platform/ioc/container'; -import { IExtensionContext } from '../../platform/common/types'; +import { IEncryptedStorage } from '../../platform/common/application/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 AGENT_SERVICE_NAME = 'deepnote-agent'; +const OPENAI_API_KEY = 'openAiApiKey'; - const value = await context.secrets.get(key); +async function retrieveOpenAiApiKey(storage: IEncryptedStorage): Promise { + const value = await storage.retrieve(AGENT_SERVICE_NAME, OPENAI_API_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 clearOpenAiApiKey(storage: IEncryptedStorage): Promise { + await storage.store(AGENT_SERVICE_NAME, OPENAI_API_KEY, undefined); } -export async function clearSecret(key: string): Promise { - const context = getContext(); +export async function getOrPromptOpenAiApiKey(storage: IEncryptedStorage): Promise { + const value = (await retrieveOpenAiApiKey(storage)) ?? (await promptForOpenAiApiKey(storage)); - if (!context) { - return; + if (!value) { + throw new Error( + l10n.t('OpenAI API key is not set. Use the command "Deepnote: Set OpenAI API Key" to configure it.') + ); } - await context.secrets.delete(key); + return value; } -export async function promptForSecret(key: string, options: SecretPromptOptions): Promise { +export async function promptForOpenAiApiKey(storage: IEncryptedStorage): Promise { const input = await window.showInputBox({ - prompt: options.prompt, - placeHolder: options.placeHolder, - password: options.password ?? true, + prompt: l10n.t('Enter your OpenAI API key'), + placeHolder: l10n.t('sk-...'), + password: true, ignoreFocusOut: true }); - if (!input || input.trim().length === 0) { + const trimmed = input?.trim(); + + if (!trimmed) { return undefined; } - const trimmed = input.trim(); - await setSecret(key, trimmed); + await storage.store(AGENT_SERVICE_NAME, OPENAI_API_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; -} - -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 index f0c339a29a..d935661c12 100644 --- a/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts @@ -1,168 +1,119 @@ 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 { IEncryptedStorage } from '../../platform/common/application/types'; import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; -import { IExtensionContext } from '../../platform/common/types'; -import { ServiceContainer } from '../../platform/ioc/container'; -import { - clearSecret, - getOrPromptOpenAiApiKey, - getOrPromptSecret, - getSecret, - promptForSecret, - setSecret -} from './deepnoteSecretStore'; +import { clearOpenAiApiKey, getOrPromptOpenAiApiKey, promptForOpenAiApiKey } from './deepnoteSecretStore'; -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'); +// The real EncryptedStorage namespaces secrets as `${service}.${key}`; the fake mirrors that so a +// store/retrieve mismatch between the two service names would surface here. +const STORED_KEY = 'deepnote-agent.openAiApiKey'; - const value = await getSecret('customKey'); - - assert.strictEqual(value, 'custom-value'); - }); +suite('deepnoteSecretStore', () => { + let storageData: Map; + let encryptedStorage: IEncryptedStorage; + let promptCount: number; - test('returns undefined when not set', async () => { - const value = await getSecret('customKey'); + function whenPromptReturns(value: string | undefined) { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenCall(() => { + promptCount++; - assert.isUndefined(value); + return Promise.resolve(value); }); + } - test('returns undefined when value is empty string', async () => { - secretStorage.set('customKey', ''); - - const value = await getSecret('customKey'); + setup(() => { + storageData = new Map(); + promptCount = 0; + encryptedStorage = mock(); + + when(encryptedStorage.store(anything(), anything(), anything())).thenCall( + (service: string, key: string, value: string | undefined) => { + if (value === undefined) { + storageData.delete(`${service}.${key}`); + } else { + storageData.set(`${service}.${key}`, value); + } - assert.isUndefined(value); - }); - }); + return Promise.resolve(); + } + ); - suite('generic setSecret', () => { - test('stores value in secrets', async () => { - await setSecret('customKey', 'custom-value'); + when(encryptedStorage.retrieve(anything(), anything())).thenCall((service: string, key: string) => + Promise.resolve(storageData.get(`${service}.${key}`)) + ); - assert.strictEqual(secretStorage.get('customKey'), 'custom-value'); - }); + whenPromptReturns(undefined); }); - suite('generic clearSecret', () => { - test('deletes value from secrets', async () => { - secretStorage.set('customKey', 'custom-value'); + suite('promptForOpenAiApiKey', () => { + test('trims, stores and returns the entered key', async () => { + whenPromptReturns(' sk-entered '); - await clearSecret('customKey'); + const value = await promptForOpenAiApiKey(instance(encryptedStorage)); - assert.isFalse(secretStorage.has('customKey')); + assert.strictEqual(value, 'sk-entered'); + assert.strictEqual(storageData.get(STORED_KEY), 'sk-entered'); }); - }); - suite('generic promptForSecret', () => { - test('stores and returns value when user enters input', async () => { - when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('user-input')); + for (const input of [undefined, ' ']) { + test(`stores nothing and returns undefined when the user enters ${JSON.stringify(input)}`, async () => { + whenPromptReturns(input); - const value = await promptForSecret('customKey', { - prompt: 'Enter value', - placeHolder: 'placeholder', - password: false - }); + const value = await promptForOpenAiApiKey(instance(encryptedStorage)); - assert.strictEqual(value, 'user-input'); - assert.strictEqual(secretStorage.get('customKey'), 'user-input'); - }); + assert.isUndefined(value); + assert.isFalse(storageData.has(STORED_KEY)); + }); + } + }); - test('returns undefined when user cancels', async () => { - when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + suite('clearOpenAiApiKey', () => { + test('deletes the stored key', async () => { + storageData.set(STORED_KEY, 'sk-stored'); - const value = await promptForSecret('customKey', { prompt: 'Enter value' }); + await clearOpenAiApiKey(instance(encryptedStorage)); - assert.isUndefined(value); + assert.isFalse(storageData.has(STORED_KEY)); }); + }); - test('returns undefined when user enters empty string', async () => { - when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(' ')); + suite('getOrPromptOpenAiApiKey', () => { + test('returns the stored key without prompting', async () => { + storageData.set(STORED_KEY, 'sk-stored'); - const value = await promptForSecret('customKey', { prompt: 'Enter value' }); + const value = await getOrPromptOpenAiApiKey(instance(encryptedStorage)); - assert.isUndefined(value); + assert.strictEqual(value, 'sk-stored'); + assert.strictEqual(promptCount, 0); }); - }); - suite('generic getOrPromptSecret', () => { - test('returns value when present in store', async () => { - secretStorage.set('customKey', 'stored-value'); + test('prompts and returns the entered key when nothing is stored', async () => { + whenPromptReturns('sk-prompted'); - const value = await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); + const value = await getOrPromptOpenAiApiKey(instance(encryptedStorage)); - assert.strictEqual(value, 'stored-value'); + assert.strictEqual(value, 'sk-prompted'); + assert.strictEqual(promptCount, 1); }); - test('prompts and returns value when missing', async () => { - when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('prompted-value')); + test('treats an empty stored value as missing and prompts', async () => { + storageData.set(STORED_KEY, ''); + whenPromptReturns('sk-prompted'); - const value = await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); + const value = await getOrPromptOpenAiApiKey(instance(encryptedStorage)); - assert.strictEqual(value, 'prompted-value'); + assert.strictEqual(value, 'sk-prompted'); + assert.strictEqual(promptCount, 1); }); - for (const scenario of [ - { - label: 'generic', - run: () => getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'), - assertError: (e: Error) => assert.strictEqual(e.message, 'Value is required') - }, - { - label: 'openAi', - run: () => getOrPromptOpenAiApiKey(), - assertError: (e: Error) => assert.include(e.message, 'OpenAI API key is not set') + test('throws when nothing is stored and the user cancels the prompt', async () => { + try { + await getOrPromptOpenAiApiKey(instance(encryptedStorage)); + assert.fail('Should have thrown'); + } catch (e) { + assert.include((e as Error).message, 'OpenAI API key is not set'); } - ]) { - test(`throws when value missing and user cancels prompt (${scenario.label})`, async () => { - when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); - - try { - await scenario.run(); - assert.fail('Should have thrown'); - } catch (e) { - scenario.assertError(e as Error); - } - }); - } + }); }); }); From 0c201093d346b1c92bcc62e0fda168aaa59d98a2 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 7 Aug 2026 11:23:10 +0000 Subject: [PATCH 49/80] Better typing for incorrect deepnote file --- .../deepnoteFileChangeWatcher.unit.test.ts | 26 ++++++++++++------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 58a096d1db..ff95daad98 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -1,4 +1,4 @@ -import type { DeepnoteFile } from '@deepnote/blocks'; +import type { DeepnoteBlock, DeepnoteFile } from '@deepnote/blocks'; import { assert } from 'chai'; import * as sinon from 'sinon'; import { anything, instance, mock, when } from 'ts-mockito'; @@ -1182,8 +1182,7 @@ project: // exist in the document but are stripped from the file, so every cell below one is offset — // the metadata-less cell would adopt the wrong block's id and outputs, and that id gets // written back, leaving two cells claiming one block. - const mockedManager = mock(); - when(mockedManager.getProjectForNotebook('e132b172-b114-410e-8331-011517db664f', 'notebook-1')).thenReturn({ + const deepnoteFile: DeepnoteFile = { version: '1.0', metadata: { createdAt: '2025-01-01T00:00:00Z' }, project: { @@ -1193,14 +1192,21 @@ project: { id: 'notebook-1', name: 'Notebook 1', - blocks: [ - { id: 'block-1', type: 'code', sortingKey: 'a0' }, - { id: 'block-2', type: 'code', sortingKey: 'a1' } - ] + blocks: [] } ] } - } as DeepnoteFile); + }; + // Force casting wihtout metadata + deepnoteFile.project.notebooks[0].blocks = [ + { id: 'block-1', type: 'code', sortingKey: 'a0' } as DeepnoteBlock, + { id: 'block-2', type: 'code', sortingKey: 'a1' } as DeepnoteBlock + ]; + + const mockedManager = mock(); + when(mockedManager.getProjectForNotebook('e132b172-b114-410e-8331-011517db664f', 'notebook-1')).thenReturn( + deepnoteFile + ); const offsetDisposables: IDisposableRegistry = []; const offsetOnDidChange = new EventEmitter(); @@ -1274,7 +1280,7 @@ project: output_type: 'execute_result', data: { 'text/plain': 'First Output' }, execution_count: 1 - } as DeepnoteOutput + } ] ], [ @@ -1284,7 +1290,7 @@ project: output_type: 'execute_result', data: { 'text/plain': 'Second Output' }, execution_count: 2 - } as DeepnoteOutput + } ] ] ]); From 28dd2fc1b4faca12068ef7197bc1d017e14263cc Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 7 Aug 2026 11:49:17 +0000 Subject: [PATCH 50/80] refactor(agent-block): change switchModel method visibility to public Updated the visibility of the switchModel method in AgentCellStatusBarProvider from private to public, allowing it to be accessed directly in unit tests. Adjusted corresponding test cases to call the method on the provider instance instead of a local function. --- .../deepnote/agentCellStatusBarProvider.ts | 2 +- .../agentCellStatusBarProvider.unit.test.ts | 14 +++++--------- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index 4a4b3bed2f..bcf34dc9dc 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -122,7 +122,7 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv return AGENT_MODEL_AUTO; } - private async switchModel(cell: NotebookCell): Promise { + public async switchModel(cell: NotebookCell): Promise { if (!isAgentCell(cell)) { return; } diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts index 46a0cc4c47..4259655279 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -169,17 +169,13 @@ suite('AgentCellStatusBarProvider', () => { ); } - function switchModel(cell: NotebookCell): Promise { - return (provider as unknown as { switchModel(cell: NotebookCell): Promise }).switchModel(cell); - } - test('Should write the picked model without dropping the cell’s other metadata', async () => { // Catches: an inverted spread in updateCellMetadata, which makes the switch a silent // no-op while still calling applyEdit — so a call-count assertion would not notice. pick('gpt-5.6-terra'); when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(true)); - await switchModel(agentCell()); + await provider.switchModel(agentCell()); verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).once(); expect(capturedEdit!.index).to.equal(2); @@ -195,7 +191,7 @@ suite('AgentCellStatusBarProvider', () => { // document on a no-op selection. pick('gpt-5.6-sol'); - await switchModel(agentCell()); + await provider.switchModel(agentCell()); verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); }); @@ -205,7 +201,7 @@ suite('AgentCellStatusBarProvider', () => { // user presses Escape. pick(undefined); - await switchModel(agentCell()); + await provider.switchModel(agentCell()); verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); }); @@ -220,7 +216,7 @@ suite('AgentCellStatusBarProvider', () => { statusBarRefreshed = true; }); - await switchModel(agentCell()); + await provider.switchModel(agentCell()); verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); expect(statusBarRefreshed, 'a rejected edit must not refresh the status bar').to.be.false; @@ -230,7 +226,7 @@ suite('AgentCellStatusBarProvider', () => { // Catches: dropping the isAgentCell guard, which would offer the model picker on any cell. pick('gpt-5.6-luna'); - await switchModel(createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } })); + await provider.switchModel(createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } })); verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).never(); verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); From de88b1c7fff858a5afffdc1a6e75232950c568e3 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 7 Aug 2026 11:53:10 +0000 Subject: [PATCH 51/80] refactor(deepnote): utilize createMockCell for test cell creation Replaced direct object definitions for agent and code cells in unit tests with the createMockCell helper function. This change enhances code readability and maintainability by standardizing cell creation across tests. --- .../deepnoteKernelAutoSelector.node.unit.test.ts | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts index 159124e7a0..72c5227374 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -3,6 +3,7 @@ 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.node'; +import { createMockCell } from './deepnoteTestHelpers'; import { ServerHandleRegistry } from '../../kernels/deepnote/deepnoteServerHandleRegistry.node'; import { IDeepnoteEnvironmentManager, @@ -20,7 +21,7 @@ import { IConfigurationService } from '../../platform/common/types'; import { IDeepnoteNotebookManager } from '../types'; import { IKernelProvider, IKernel, IJupyterKernelSpec } from '../../kernels/types'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; -import { EventEmitter, NotebookCell, NotebookDocument, Uri, NotebookController, CancellationToken } from 'vscode'; +import { EventEmitter, 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'; @@ -1063,11 +1064,8 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { return placeholder; } - const agentCell = { - index: 0, - metadata: { __deepnotePocket: { type: 'agent' } } - } as unknown as NotebookCell; - const codeCell = { index: 1, metadata: {} } as unknown as NotebookCell; + const agentCell = createMockCell({ index: 0, metadata: { __deepnotePocket: { type: 'agent' } } }); + const codeCell = createMockCell({ index: 1 }); test('configures the environment and executes nothing', async () => { when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(true); From 912a2e7604287d5fb03efaf4b9d04f21880fa764 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 7 Aug 2026 15:29:42 +0000 Subject: [PATCH 52/80] Remove different visualization for ephemeral blocks --- .../ephemeralCellDecorationProvider.ts | 125 ------------------ src/notebooks/serviceRegistry.node.ts | 5 - src/notebooks/serviceRegistry.web.ts | 5 - src/renderers/client/markdown.ts | 75 +---------- 4 files changed, 1 insertion(+), 209 deletions(-) delete mode 100644 src/notebooks/deepnote/ephemeralCellDecorationProvider.ts diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts deleted file mode 100644 index 19c3aa8470..0000000000 --- a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts +++ /dev/null @@ -1,125 +0,0 @@ -import { - Disposable, - NotebookCell, - NotebookDocument, - OverviewRulerLane, - Range, - TextEditor, - TextEditorDecorationType, - ThemeColor, - window, - workspace -} from 'vscode'; -import { injectable } from 'inversify'; - -import { IExtensionSyncActivationService } from '../../platform/activation/types'; -import { logger } from '../../platform/logging'; -import { isEphemeralCell } from './dataConversionUtils'; - -const NOTEBOOK_CELL_SCHEME = 'vscode-notebook-cell'; - -/** Ephemeral styling in code-cell editors; markdown cells use `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 { - for (const disposable of this.disposables) { - disposable.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) { - 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 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); - } catch (error) { - logger.warn( - `EphemeralCellDecorationProvider: Failed to update decorations for ${editor.document.uri.path}`, - error - ); - } - } - } -} diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index a70e84ac1c..f8a775783f 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -97,7 +97,6 @@ import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnote 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'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIntegrationStartupCodeProvider'; @@ -277,10 +276,6 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea 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 8e7b5e665a..fe6a8dbe90 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/src/notebooks/serviceRegistry.web.ts @@ -53,7 +53,6 @@ import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnote 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'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; @@ -143,10 +142,6 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea 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 4374444265..b5399de2df 100644 --- a/src/renderers/client/markdown.ts +++ b/src/renderers/client/markdown.ts @@ -1,30 +1,3 @@ -import type { ActivationFunction } from 'vscode-notebook-renderer'; - -// Local markdown-it shape — transitive dep, no types. -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; @@ -58,59 +31,13 @@ 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 const activate: ActivationFunction = async (ctx) => { +export async function activate() { 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'); - const extendMarkdownIt = markdownRenderer?.extendMarkdownIt as ExtendMarkdownIt | undefined; - - if (typeof extendMarkdownIt === 'function') { - extendMarkdownIt((md) => { - addEphemeralCellWrapper(md); - }); - } - - return undefined; -}; - -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; - } - - 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 8320959efb7b710a983fb28337004b2028f6e235 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 7 Aug 2026 19:05:34 +0000 Subject: [PATCH 53/80] feat(agent-block): add a Clear ephemeral blocks button MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The ephemeral cell status bar now carries a button next to the Ephemeral indicator. It deletes every ephemeral cell sharing the clicked cell's agent_source_block_id — the whole generated run, not just the cell that was clicked — behind a modal confirmation showing the count. An ephemeral cell that records no source block clears only itself. Deletions go into one WorkspaceEdit in descending index order, so each range still addresses the cell it was computed from. The E2E test drives the real status bar item and generates the run it clears, so it holds up under --grep and a Mocha retry. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- .../ephemeralCellStatusBarProvider.ts | 90 +++++++- ...phemeralCellStatusBarProvider.unit.test.ts | 211 +++++++++++++++--- test/e2e/helpers/notebook.ts | 41 ++++ test/e2e/suite/agentBlock.e2e.test.ts | 77 ++++++- 4 files changed, 386 insertions(+), 33 deletions(-) diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts index d3da9be862..4f9d1a8f36 100644 --- a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts @@ -5,16 +5,24 @@ import { NotebookCell, NotebookCellStatusBarItem, NotebookCellStatusBarItemProvider, + NotebookEdit, + NotebookRange, + WorkspaceEdit, + commands, l10n, notebooks, + window, workspace } from 'vscode'; import { injectable } from 'inversify'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; -import { isEphemeralCell } from './dataConversionUtils'; +import { getEphemeralCellAgentSourceBlockId, isEphemeralCell } from './dataConversionUtils'; + +const CLEAR_EPHEMERAL_BLOCKS_COMMAND = 'deepnote.clearEphemeralBlocks'; const EPHEMERAL_INDICATOR_PRIORITY = 1000; +const CLEAR_EPHEMERAL_PRIORITY = 990; @injectable() export class EphemeralCellStatusBarProvider @@ -36,9 +44,49 @@ export class EphemeralCellStatusBarProvider }) ); + this.disposables.push( + commands.registerCommand(CLEAR_EPHEMERAL_BLOCKS_COMMAND, async (cell?: NotebookCell) => { + const activeCell = cell || this.getActiveCell(); + if (activeCell) { + await this.clearEphemeralBlocks(activeCell); + } + }) + ); + this.disposables.push(this._onDidChangeCellStatusBarItems); } + /** Deletes the ephemeral cells the same agent block generated, after a modal confirmation. */ + public async clearEphemeralBlocks(cell: NotebookCell): Promise { + if (!isEphemeralCell(cell)) { + return; + } + + const cellsToClear = this.getCellsToClear(cell); + + const confirmation = await window.showWarningMessage( + l10n.t('Clear {0} ephemeral block(s) from this notebook?', cellsToClear.length), + { modal: true }, + l10n.t('Clear') + ); + + if (confirmation !== l10n.t('Clear')) { + return; + } + + // Descending so each deletion's index still addresses the cell it was computed from. + const deletions = [...cellsToClear] + .sort((a, b) => b.index - a.index) + .map((target) => NotebookEdit.deleteCells(new NotebookRange(target.index, target.index + 1))); + + const edit = new WorkspaceEdit(); + edit.set(cell.notebook.uri, deletions); + + if (!(await workspace.applyEdit(edit))) { + void window.showErrorMessage(l10n.t('Failed to clear ephemeral blocks')); + } + } + public dispose(): void { this.disposables.forEach((d) => d.dispose()); } @@ -46,7 +94,7 @@ export class EphemeralCellStatusBarProvider public provideCellStatusBarItems( cell: NotebookCell, token: CancellationToken - ): NotebookCellStatusBarItem | undefined { + ): NotebookCellStatusBarItem[] | undefined { if (token.isCancellationRequested) { return undefined; } @@ -57,7 +105,21 @@ export class EphemeralCellStatusBarProvider const agentSourceBlockId = cell.metadata?.agent_source_block_id as string | undefined; - return this.createEphemeralIndicatorItem(agentSourceBlockId ?? null); + return [this.createEphemeralIndicatorItem(agentSourceBlockId ?? null), this.createClearEphemeralItem(cell)]; + } + + private createClearEphemeralItem(cell: NotebookCell): NotebookCellStatusBarItem { + return { + text: `$(trash) ${l10n.t('Clear ephemeral blocks')}`, + alignment: 1, + priority: CLEAR_EPHEMERAL_PRIORITY, + tooltip: l10n.t('Remove the ephemeral blocks generated by this agent block'), + command: { + title: l10n.t('Clear ephemeral blocks'), + command: CLEAR_EPHEMERAL_BLOCKS_COMMAND, + arguments: [cell] + } + }; } private createEphemeralIndicatorItem(agentSourceBlockId: string | null): NotebookCellStatusBarItem { @@ -73,4 +135,26 @@ export class EphemeralCellStatusBarProvider tooltip: tooltipLines.join('\n') }; } + + private getActiveCell(): NotebookCell | undefined { + const activeEditor = window.activeNotebookEditor; + if (activeEditor && activeEditor.selection) { + return activeEditor.notebook.cellAt(activeEditor.selection.start); + } + + return undefined; + } + + /** Siblings of the same agent run; just the cell itself when it records no source block. */ + private getCellsToClear(cell: NotebookCell): NotebookCell[] { + const agentSourceBlockId = getEphemeralCellAgentSourceBlockId(cell); + + if (!agentSourceBlockId) { + return [cell]; + } + + return cell.notebook + .getCells() + .filter((candidate) => getEphemeralCellAgentSourceBlockId(candidate) === agentSourceBlockId); + } } diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts index 4785815bcc..73bd80414e 100644 --- a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts @@ -1,8 +1,11 @@ import { expect } from 'chai'; -import { CancellationToken } from 'vscode'; +import * as sinon from 'sinon'; +import { anything, verify, when } from 'ts-mockito'; +import { CancellationToken, NotebookCell, WorkspaceEdit } from 'vscode'; +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; import { EphemeralCellStatusBarProvider } from './ephemeralCellStatusBarProvider'; -import { createMockCell } from './deepnoteTestHelpers'; +import { createMockCell, createMockNotebookWithCells } from './deepnoteTestHelpers'; suite('EphemeralCellStatusBarProvider', () => { let provider: EphemeralCellStatusBarProvider; @@ -23,30 +26,30 @@ suite('EphemeralCellStatusBarProvider', () => { suite('Ephemeral Cell Detection', () => { test('Should return undefined when is_ephemeral is false', () => { const cell = createMockCell({ metadata: { is_ephemeral: false } }); - const item = provider.provideCellStatusBarItems(cell, mockToken); + const items = provider.provideCellStatusBarItems(cell, mockToken); - expect(item).to.be.undefined; + expect(items).to.be.undefined; }); test('Should return undefined when is_ephemeral is not set', () => { const cell = createMockCell({ metadata: {} }); - const item = provider.provideCellStatusBarItems(cell, mockToken); + const items = provider.provideCellStatusBarItems(cell, mockToken); - expect(item).to.be.undefined; + expect(items).to.be.undefined; }); test('Should return undefined for cell without metadata', () => { const cell = createMockCell({ metadata: undefined }); - const item = provider.provideCellStatusBarItems(cell, mockToken); + const items = provider.provideCellStatusBarItems(cell, mockToken); - expect(item).to.be.undefined; + expect(items).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); + const items = provider.provideCellStatusBarItems(cell, mockToken); - expect(item).to.be.undefined; + expect(items).to.be.undefined; }); test('Should return undefined when cancellation is requested', () => { @@ -55,32 +58,47 @@ suite('EphemeralCellStatusBarProvider', () => { onCancellationRequested: () => ({ dispose: () => undefined }) } as any; const cell = createMockCell({ metadata: { is_ephemeral: true } }); - const item = provider.provideCellStatusBarItems(cell, cancelledToken); + const items = provider.provideCellStatusBarItems(cell, cancelledToken); - expect(item).to.be.undefined; + expect(items).to.be.undefined; }); }); suite('Status Bar Item Properties', () => { test('Should set ephemeral status bar item properties', () => { const cell = createMockCell({ metadata: { is_ephemeral: true } }); - const item = provider.provideCellStatusBarItems(cell, mockToken)!; + const items = provider.provideCellStatusBarItems(cell, mockToken)!; - expect(item.text).to.include('$(sparkle)'); - expect(item.text).to.include('Ephemeral'); - expect(item.alignment).to.equal(1); - expect(item.priority).to.equal(1000); - expect(item.command).to.be.undefined; + expect(items).to.have.lengthOf(2); + expect(items[0].text).to.include('$(sparkle)'); + expect(items[0].text).to.include('Ephemeral'); + expect(items[0].alignment).to.equal(1); + expect(items[0].priority).to.equal(1000); + expect(items[0].command).to.be.undefined; + }); + + test('Should set clear button properties and pass the cell to its command', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('$(trash)'); + expect(items[1].text).to.include('Clear ephemeral blocks'); + expect(items[1].alignment).to.equal(1); + expect(items[1].priority).to.equal(990); + + const command = items[1].command as { command: string; arguments: unknown[] }; + expect(command.command).to.equal('deepnote.clearEphemeralBlocks'); + expect(command.arguments).to.deep.equal([cell]); }); }); suite('Tooltip', () => { test('Should describe ephemeral tooltip without source block when agent_source_block_id is absent', () => { const cell = createMockCell({ metadata: { is_ephemeral: true } }); - const item = provider.provideCellStatusBarItems(cell, mockToken)!; + const items = provider.provideCellStatusBarItems(cell, mockToken)!; - expect(item.tooltip).to.include('Auto-generated ephemeral block'); - expect(item.tooltip).to.not.include('Source agent block'); + expect(items[0].tooltip).to.include('Auto-generated ephemeral block'); + expect(items[0].tooltip).to.not.include('Source agent block'); }); test('Should include agent source block ID in tooltip when present', () => { @@ -90,15 +108,15 @@ suite('EphemeralCellStatusBarProvider', () => { agent_source_block_id: 'a0000000000000000000000000000004' } }); - const item = provider.provideCellStatusBarItems(cell, mockToken)!; + const items = provider.provideCellStatusBarItems(cell, mockToken)!; - expect(item.tooltip).to.include('a0000000000000000000000000000004'); - expect(item.tooltip).to.include('Source agent block'); + expect(items[0].tooltip).to.include('a0000000000000000000000000000004'); + expect(items[0].tooltip).to.include('Source agent block'); }); }); suite('Coexistence with other cell types', () => { - test('Should return item for ephemeral cells regardless of pocket type', () => { + test('Should return items for ephemeral cells regardless of pocket type', () => { const pocketTypes = ['agent', 'code', 'markdown'] as const; for (const type of pocketTypes) { @@ -109,10 +127,149 @@ suite('EphemeralCellStatusBarProvider', () => { ...(type === 'agent' ? { agent_source_block_id: 'source-id' } : {}) } }); - const item = provider.provideCellStatusBarItems(cell, mockToken); + const items = provider.provideCellStatusBarItems(cell, mockToken); - expect(item).to.not.be.undefined; + expect(items).to.not.be.undefined; } }); }); + + suite('Clearing ephemeral blocks', () => { + setup(() => { + resetVSCodeMocks(); + }); + + teardown(() => { + sinon.restore(); + resetVSCodeMocks(); + }); + + // Mocked WorkspaceEdit.set drops the edits, so capture them on the prototype and replay the + // deletions against `cells` — an ascending delete order corrupts the survivors, not the count. + function applyDeletionsTo(cells: NotebookCell[]): void { + let recordedEdits: { range: { start: number; end: number } }[] = []; + + sinon.stub(WorkspaceEdit.prototype, 'set').callsFake((_uri, edits) => { + recordedEdits = edits as unknown as { range: { start: number; end: number } }[]; + }); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { + for (const { range } of recordedEdits) { + cells.splice(range.start, range.end - range.start); + } + + return Promise.resolve(true); + }); + } + + function confirmWith(label: string | undefined): void { + when(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).thenReturn( + Promise.resolve(label as any) + ); + } + + function ephemeralCell(text: string, agentSourceBlockId?: string) { + return { + text, + metadata: { + is_ephemeral: true, + ...(agentSourceBlockId ? { agent_source_block_id: agentSourceBlockId } : {}) + } + }; + } + + function agentCell(text: string, blockId: string) { + return { text, metadata: { __deepnotePocket: { type: 'agent' }, id: blockId } }; + } + + test('Should delete every cell generated by the clicked cell’s agent block and nothing else', async () => { + const { cells } = createMockNotebookWithCells([ + agentCell('agent A', 'agent-block-1'), + ephemeralCell('eph A1', 'agent-block-1'), + ephemeralCell('eph A2', 'agent-block-1'), + { text: 'user code', metadata: {} }, + agentCell('agent B', 'agent-block-2'), + ephemeralCell('eph B1', 'agent-block-2') + ]); + const clickedCell = cells[1]; + applyDeletionsTo(cells); + confirmWith('Clear'); + + await provider.clearEphemeralBlocks(clickedCell); + + expect(cells.map((cell) => cell.document.getText())).to.deep.equal([ + 'agent A', + 'user code', + 'agent B', + 'eph B1' + ]); + }); + + test('Should delete only the clicked cell when it records no source agent block', async () => { + const { cells } = createMockNotebookWithCells([ + ephemeralCell('orphan'), + ephemeralCell('eph A1', 'agent-block-1') + ]); + applyDeletionsTo(cells); + confirmWith('Clear'); + + await provider.clearEphemeralBlocks(cells[0]); + + expect(cells.map((cell) => cell.document.getText())).to.deep.equal(['eph A1']); + }); + + test('Should report how many blocks the clear removes', async () => { + const { cells } = createMockNotebookWithCells([ + ephemeralCell('eph A1', 'agent-block-1'), + ephemeralCell('eph A2', 'agent-block-1') + ]); + applyDeletionsTo(cells); + confirmWith('Clear'); + + await provider.clearEphemeralBlocks(cells[0]); + + verify( + mockedVSCodeNamespaces.window.showWarningMessage( + 'Clear 2 ephemeral block(s) from this notebook?', + anything(), + anything() + ) + ).once(); + }); + + test('Should not edit the notebook when the confirmation is dismissed', async () => { + // Catches: applying the edit before the modal is answered, which deletes on a cancel. + const { cells } = createMockNotebookWithCells([ephemeralCell('eph A1', 'agent-block-1')]); + applyDeletionsTo(cells); + confirmWith(undefined); + + await provider.clearEphemeralBlocks(cells[0]); + + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + expect(cells).to.have.lengthOf(1); + }); + + test('Should ignore a cell that is not ephemeral', async () => { + // Catches: dropping the isEphemeralCell guard, which would offer to clear any cell. + const { cells } = createMockNotebookWithCells([{ text: 'user code', metadata: {} }]); + applyDeletionsTo(cells); + confirmWith('Clear'); + + await provider.clearEphemeralBlocks(cells[0]); + + verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('Should report an error when the workspace edit is rejected', async () => { + // Catches: dropping the `if (!applyEdit)` branch, which loses the clear silently. + const { cells } = createMockNotebookWithCells([ephemeralCell('eph A1', 'agent-block-1')]); + confirmWith('Clear'); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(false)); + + await provider.clearEphemeralBlocks(cells[0]); + + verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); + }); + }); }); diff --git a/test/e2e/helpers/notebook.ts b/test/e2e/helpers/notebook.ts index 8f6af68872..1945608472 100644 --- a/test/e2e/helpers/notebook.ts +++ b/test/e2e/helpers/notebook.ts @@ -41,6 +41,47 @@ export async function clickRunAll(notebookFileName: string): Promise { ); } +/** + * Clicks the notebook cell status bar item whose text contains `label`. Cell chrome lives in the + * main window DOM (not the output iframe), so this switches out of the webview first and matches on + * `textContent` — Selenium's `getText()` is empty for items scrolled out of view. + */ +export async function clickCellStatusBarItem(label: string): Promise { + const driver = VSBrowser.instance.driver; + + await new WebView().switchBack().catch((error) => { + console.warn('[deepnote-e2e] switch back before clicking a cell status bar item:', error); + }); + + // Locate AND click in the same wait loop: the status bar re-renders as cells execute, which + // would otherwise surface as a StaleElementReferenceError between finding and clicking. + await driver.wait( + async () => { + try { + for (const item of await driver.findElements(By.css('.cell-statusbar-container .cell-status-item'))) { + const text = (await item.getAttribute('textContent')) ?? ''; + if (!text.includes(label)) { + continue; + } + + await driver.executeScript('arguments[0].scrollIntoView({block: "center"})', item); + await item.click(); + + return true; + } + + return false; + } catch (error) { + console.warn('[deepnote-e2e] locate/click cell status bar item (retrying):', error); + + return false; + } + }, + WORKBENCH_TIMEOUT, + `notebook cell status bar item "${label}" did not appear or could not be clicked` + ); +} + /** Run `read` in the notebook output webview; '' if the frame is missing. */ async function readInsideNotebookWebview(read: (webView: WebView) => Promise): Promise { const driver = VSBrowser.instance.driver; diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts index 4e3551c91f..c89a058ad0 100644 --- a/test/e2e/suite/agentBlock.e2e.test.ts +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -9,6 +9,7 @@ import { QUICK_PICK_TIMEOUT, SUITE_TIMEOUT, WORKBENCH_TIMEOUT, + clickCellStatusBarItem, clickRunAll, confirmModalDialog, copyFixtureToTempDir, @@ -45,6 +46,16 @@ const RERUN_MARKDOWN_TEXT = 'Second-run markdown from the E2E agent'; const RERUN_FINAL_AGENT_TEXT = 'Re-run summary added as a markdown block.'; // executeAgentCell stale-run error substring. const STALE_CELLS_ERROR_TEXT = 'from its previous run'; +// Third run, disjoint from both prior runs so the clear test stands on its own. +const CLEAR_RUN_PYTHON_OUTPUT_MARKER = 'clear-run-python-ran'; +const CLEAR_RUN_GENERATED_PYTHON = `print("${CLEAR_RUN_PYTHON_OUTPUT_MARKER}")`; +const CLEAR_RUN_MARKDOWN_TEXT = 'Third-run markdown from the E2E agent'; +const CLEAR_RUN_FINAL_AGENT_TEXT = 'Clear-run summary added as a markdown block.'; +// EphemeralCellStatusBarProvider button and its confirmation. +const CLEAR_EPHEMERAL_BUTTON = 'Clear ephemeral blocks'; +const CLEAR_EPHEMERAL_CONFIRM = 'Clear'; +const CLEAR_EPHEMERAL_CONFIRM_TEXT = 'ephemeral block'; +const CLEAR_EPHEMERAL_TIMEOUT = 30_000; 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'; @@ -52,7 +63,12 @@ const REVERT_FILE_COMMAND = 'File: Revert File'; const DISCARD_CHANGES_BUTTON = "Don't Save"; const API_KEY_SAVED_NOTIFICATION = /OpenAI API key has been saved/; -async function awaitWebviewMarkers(markers: string[], timeout: number, context: string): Promise { +async function awaitWebviewMarkers( + markers: string[], + timeout: number, + context: string, + absentMarkers: string[] = [] +): Promise { const driver = VSBrowser.instance.driver; const deadline = Date.now() + timeout; let text = ''; @@ -60,7 +76,8 @@ async function awaitWebviewMarkers(markers: string[], timeout: number, context: while (Date.now() < deadline) { text = await readNotebookWebviewText(); const missing = markers.filter((marker) => !text.includes(marker)); - if (missing.length === 0) { + const lingering = absentMarkers.filter((marker) => text.includes(marker)); + if (missing.length === 0 && lingering.length === 0) { return text; } @@ -68,9 +85,10 @@ async function awaitWebviewMarkers(markers: string[], timeout: number, context: } const missing = markers.filter((marker) => !text.includes(marker)); + const lingering = absentMarkers.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)}` + `Lingering: ${JSON.stringify(lingering)}. Last text: ${JSON.stringify(text)}` ); } @@ -278,4 +296,57 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu assertOccurrences(rendered, RERUN_MARKDOWN_TEXT, 1); assertOccurrences(rendered, STALE_CELLS_ERROR_TEXT, 0); }); + + // Self-contained: generates the run it clears, so it survives --grep and a Mocha retry (Run All + // drops any stale generated cells first). + it('clears the whole generated run from the ephemeral cell status bar button', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: CLEAR_RUN_GENERATED_PYTHON }), + id: 'call_e2e_clear_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: CLEAR_RUN_PYTHON_OUTPUT_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: CLEAR_RUN_MARKDOWN_TEXT }), + id: 'call_e2e_clear_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: CLEAR_RUN_FINAL_AGENT_TEXT } + } + ]); + + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); + + await awaitWebviewMarkers( + [CLEAR_RUN_PYTHON_OUTPUT_MARKER, CLEAR_RUN_MARKDOWN_TEXT, CLEAR_RUN_FINAL_AGENT_TEXT], + AGENT_RUN_TIMEOUT, + 'agent run whose cells the button clears' + ); + + await clickCellStatusBarItem(CLEAR_EPHEMERAL_BUTTON); + await confirmModalDialog(CLEAR_EPHEMERAL_CONFIRM, { messageIncludes: CLEAR_EPHEMERAL_CONFIRM_TEXT }); + + // The clicked cell is the generated code cell and the markdown cell of the same run goes with + // it. Requiring the agent's own transcript to survive keeps an unreadable webview (which reads + // as '') from passing this as "the generated cells are gone". + await awaitWebviewMarkers([CLEAR_RUN_FINAL_AGENT_TEXT], CLEAR_EPHEMERAL_TIMEOUT, 'ephemeral cells cleared', [ + CLEAR_RUN_PYTHON_OUTPUT_MARKER, + CLEAR_RUN_MARKDOWN_TEXT + ]); + + await screenshot('agent-ephemeral-cleared'); + }); }); From 1a55453a14fce16b26b377b61493df5e68d699e7 Mon Sep 17 00:00:00 2001 From: tomas Date: Sat, 8 Aug 2026 16:01:46 +0000 Subject: [PATCH 54/80] refactor(agent-block): move the clear button onto the agent block MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The agent block owns the cells it generates, so it should own the button that clears them. The target flips accordingly: from siblings of the clicked ephemeral cell to children of this agent block, matched on getBlockId(agentCell) — the same derivation removeEphemeralCellsForAgentBlocks already uses. The button stays hidden while the block owns no ephemeral cells, so it can never prompt to clear zero of them. The command handler now throws when invoked without a cell rather than silently doing nothing. Only the status bar item reaches it and that always passes the cell, but registerCommand publishes the id globally, so a missing argument is a wiring bug worth surfacing. Ephemeral cells keep just their label. Orphans (is_ephemeral with no agent_source_block_id) lose their self-clear button; the serializer strips them from the file either way. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- .../deepnote/agentCellStatusBarProvider.ts | 90 ++++++- .../agentCellStatusBarProvider.unit.test.ts | 237 +++++++++++++++++- .../ephemeralCellStatusBarProvider.ts | 88 +------ ...phemeralCellStatusBarProvider.unit.test.ts | 164 +----------- test/e2e/suite/agentBlock.e2e.test.ts | 10 +- 5 files changed, 332 insertions(+), 257 deletions(-) diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index bcf34dc9dc..b4ead0ba8d 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -6,6 +6,7 @@ import { NotebookCellStatusBarItem, NotebookCellStatusBarItemProvider, NotebookEdit, + NotebookRange, WorkspaceEdit, commands, l10n, @@ -16,18 +17,21 @@ import { import { injectable } from 'inversify'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; -import { isAgentCell } from './dataConversionUtils'; +import { getBlockId, getEphemeralCellAgentSourceBlockId, isAgentCell } from './dataConversionUtils'; /** Same key as `agentBlockSchema` / `executeAgentBlock`. */ -const AGENT_MODEL_METADATA_KEY = 'deepnote_agent_model'; +export const AGENT_MODEL_METADATA_KEY = 'deepnote_agent_model'; /** Persisted default — absent key becomes `undefined` and breaks openai() model selection. */ -const AGENT_MODEL_AUTO = 'auto'; +export const AGENT_MODEL_AUTO = 'auto'; const AGENT_MODEL_OPTIONS = [AGENT_MODEL_AUTO, 'gpt-5.6-sol', 'gpt-5.6-terra', 'gpt-5.6-luna']; +const CLEAR_EPHEMERAL_BLOCKS_COMMAND = 'deepnote.clearEphemeralBlocks'; + const AGENT_INDICATOR_PRIORITY = 100; const MODEL_PICKER_PRIORITY = 90; +const CLEAR_EPHEMERAL_PRIORITY = 80; @injectable() export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService { @@ -56,9 +60,54 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv }) ); + this.disposables.push( + commands.registerCommand(CLEAR_EPHEMERAL_BLOCKS_COMMAND, async (cell?: NotebookCell) => { + if (!cell) { + throw new Error(`${CLEAR_EPHEMERAL_BLOCKS_COMMAND} requires the cell it was invoked from`); + } + + await this.clearEphemeralBlocks(cell); + }) + ); + this.disposables.push(this._onDidChangeCellStatusBarItems); } + /** Deletes the ephemeral cells this agent block generated, after a modal confirmation. */ + public async clearEphemeralBlocks(cell: NotebookCell): Promise { + if (!isAgentCell(cell)) { + return; + } + + const cellsToClear = this.getCellsToClear(cell); + + if (cellsToClear.length === 0) { + return; + } + + const confirmation = await window.showWarningMessage( + l10n.t('Clear {0} ephemeral block(s) from this notebook?', cellsToClear.length), + { modal: true }, + l10n.t('Clear') + ); + + if (confirmation !== l10n.t('Clear')) { + return; + } + + // Descending so each deletion's index still addresses the cell it was computed from. + const deletions = [...cellsToClear] + .sort((a, b) => b.index - a.index) + .map((target) => NotebookEdit.deleteCells(new NotebookRange(target.index, target.index + 1))); + + const edit = new WorkspaceEdit(); + edit.set(cell.notebook.uri, deletions); + + if (!(await workspace.applyEdit(edit))) { + void window.showErrorMessage(l10n.t('Failed to clear ephemeral blocks')); + } + } + public dispose(): void { this.disposables.forEach((disposable) => disposable.dispose()); } @@ -78,7 +127,13 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv const metadata = cell.metadata as Record | undefined; const model = this.getModel(metadata); - return [this.createAgentIndicatorItem(), this.createModelPickerItem(cell, model)]; + const items = [this.createAgentIndicatorItem(), this.createModelPickerItem(cell, model)]; + + if (this.getCellsToClear(cell).length > 0) { + items.push(this.createClearEphemeralItem(cell)); + } + + return items; } private createAgentIndicatorItem(): NotebookCellStatusBarItem { @@ -90,6 +145,20 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv }; } + private createClearEphemeralItem(cell: NotebookCell): NotebookCellStatusBarItem { + return { + text: `$(trash) ${l10n.t('Clear ephemeral blocks')}`, + alignment: 1, + priority: CLEAR_EPHEMERAL_PRIORITY, + tooltip: l10n.t('Remove the ephemeral blocks generated by this agent block'), + command: { + title: l10n.t('Clear ephemeral blocks'), + command: CLEAR_EPHEMERAL_BLOCKS_COMMAND, + arguments: [cell] + } + }; + } + private createModelPickerItem(cell: NotebookCell, model: string): NotebookCellStatusBarItem { return { text: `$(symbol-enum) ${l10n.t('Model: {0}', model)}`, @@ -113,6 +182,19 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv return undefined; } + /** Ephemeral cells this agent block generated; empty when it has no block id or has not run. */ + private getCellsToClear(cell: NotebookCell): NotebookCell[] { + const agentBlockId = getBlockId(cell); + + if (!agentBlockId) { + return []; + } + + return cell.notebook + .getCells() + .filter((candidate) => getEphemeralCellAgentSourceBlockId(candidate) === agentBlockId); + } + private getModel(metadata: Record | undefined): string { const value = metadata?.[AGENT_MODEL_METADATA_KEY]; if (typeof value === 'string' && value) { diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts index 4259655279..04b6f9234e 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -1,11 +1,11 @@ -import { expect } from 'chai'; +import { assert, expect } from 'chai'; import * as sinon from 'sinon'; import { anything, verify, when } from 'ts-mockito'; -import { CancellationToken, NotebookCell, NotebookEdit } from 'vscode'; +import { CancellationToken, NotebookCell, NotebookEdit, WorkspaceEdit } from 'vscode'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; import { AgentCellStatusBarProvider } from './agentCellStatusBarProvider'; -import { createMockCell } from './deepnoteTestHelpers'; +import { createMockCell, createMockNotebookWithCells } from './deepnoteTestHelpers'; suite('AgentCellStatusBarProvider', () => { let provider: AgentCellStatusBarProvider; @@ -232,4 +232,235 @@ suite('AgentCellStatusBarProvider', () => { verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); }); }); + + suite('Clearing ephemeral blocks', () => { + setup(() => { + resetVSCodeMocks(); + }); + + teardown(() => { + sinon.restore(); + resetVSCodeMocks(); + }); + + // Mocked WorkspaceEdit.set drops the edits, so capture them on the prototype and replay the + // deletions against `cells` — an ascending delete order corrupts the survivors, not the count. + function applyDeletionsTo(cells: NotebookCell[]): void { + let recordedEdits: { range: { start: number; end: number } }[] = []; + + sinon.stub(WorkspaceEdit.prototype, 'set').callsFake((_uri, edits) => { + recordedEdits = edits as unknown as { range: { start: number; end: number } }[]; + }); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { + for (const { range } of recordedEdits) { + cells.splice(range.start, range.end - range.start); + } + + return Promise.resolve(true); + }); + } + + function confirmWith(label: string | undefined): void { + when(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).thenReturn( + Promise.resolve(label as any) + ); + } + + function agentBlock(text: string, blockId: string) { + return { text, metadata: { __deepnotePocket: { type: 'agent' }, id: blockId } }; + } + + function ephemeralCell(text: string, agentSourceBlockId?: string) { + return { + text, + metadata: { + is_ephemeral: true, + ...(agentSourceBlockId ? { agent_source_block_id: agentSourceBlockId } : {}) + } + }; + } + + test('Should delete every cell this agent block generated and nothing else', async () => { + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph A1', 'agent-block-1'), + ephemeralCell('eph A2', 'agent-block-1'), + { text: 'user code', metadata: {} }, + agentBlock('agent B', 'agent-block-2'), + ephemeralCell('eph B1', 'agent-block-2'), + ephemeralCell('orphan') + ]); + applyDeletionsTo(cells); + confirmWith('Clear'); + + await provider.clearEphemeralBlocks(cells[0]); + + expect(cells.map((cell) => cell.document.getText())).to.deep.equal([ + 'agent A', + 'user code', + 'agent B', + 'eph B1', + 'orphan' + ]); + }); + + test('Should report how many blocks the clear removes', async () => { + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph A1', 'agent-block-1'), + ephemeralCell('eph A2', 'agent-block-1') + ]); + applyDeletionsTo(cells); + confirmWith('Clear'); + + await provider.clearEphemeralBlocks(cells[0]); + + verify( + mockedVSCodeNamespaces.window.showWarningMessage( + 'Clear 2 ephemeral block(s) from this notebook?', + anything(), + anything() + ) + ).once(); + }); + + test('Should not edit the notebook when the confirmation is dismissed', async () => { + // Catches: applying the edit before the modal is answered, which deletes on a cancel. + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph A1', 'agent-block-1') + ]); + applyDeletionsTo(cells); + confirmWith(undefined); + + await provider.clearEphemeralBlocks(cells[0]); + + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + expect(cells).to.have.lengthOf(2); + }); + + test('Should not prompt for an agent block that generated nothing', async () => { + // Catches: prompting to clear 0 blocks when the agent has not run. + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph B1', 'agent-block-2') + ]); + applyDeletionsTo(cells); + confirmWith('Clear'); + + await provider.clearEphemeralBlocks(cells[0]); + + verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('Should ignore a cell that is not an agent block', async () => { + // Catches: dropping the isAgentCell guard, which would clear from any cell whose id + // happens to own ephemeral children. + const { cells } = createMockNotebookWithCells([ + { text: 'user code', metadata: { id: 'agent-block-1' } }, + ephemeralCell('eph A1', 'agent-block-1') + ]); + applyDeletionsTo(cells); + confirmWith('Clear'); + + await provider.clearEphemeralBlocks(cells[0]); + + verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('Should report an error when the workspace edit is rejected', async () => { + // Catches: dropping the `if (!applyEdit)` branch, which loses the clear silently. + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph A1', 'agent-block-1') + ]); + confirmWith('Clear'); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(false)); + + await provider.clearEphemeralBlocks(cells[0]); + + verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); + }); + + suite('Clear button', () => { + test('Should offer the button on an agent block that generated ephemeral cells', () => { + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph A1', 'agent-block-1') + ]); + const items = provider.provideCellStatusBarItems(cells[0], mockToken)!; + + expect(items).to.have.lengthOf(3); + expect(items[2].text).to.include('$(trash)'); + expect(items[2].text).to.include('Clear ephemeral blocks'); + expect(items[2].alignment).to.equal(1); + expect(items[2].priority).to.equal(80); + + const command = items[2].command as { command: string; arguments: unknown[] }; + expect(command.command).to.equal('deepnote.clearEphemeralBlocks'); + expect(command.arguments).to.deep.equal([cells[0]]); + }); + + test('Should hide the button on an agent block that owns no ephemeral cells', () => { + // Catches: an always-visible button, which prompts to clear 0 blocks. + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph B1', 'agent-block-2') + ]); + const items = provider.provideCellStatusBarItems(cells[0], mockToken)!; + + expect(items).to.have.lengthOf(2); + }); + }); + + suite('Command handler', () => { + let invokeCommand: (cell?: NotebookCell) => Promise; + + setup(() => { + when( + mockedVSCodeNamespaces.notebooks.registerNotebookCellStatusBarItemProvider(anything(), anything()) + ).thenReturn({ dispose: () => undefined }); + when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument).thenReturn(() => ({ + dispose: () => undefined + })); + when(mockedVSCodeNamespaces.commands.registerCommand(anything(), anything())).thenCall( + (id: string, callback: (cell?: NotebookCell) => Promise) => { + if (id === 'deepnote.clearEphemeralBlocks') { + invokeCommand = callback; + } + + return { dispose: () => undefined }; + } + ); + + provider.activate(); + }); + + test('Should clear the blocks of the cell the command is given', async () => { + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph A1', 'agent-block-1') + ]); + applyDeletionsTo(cells); + confirmWith('Clear'); + + await invokeCommand(cells[0]); + + expect(cells.map((cell) => cell.document.getText())).to.deep.equal(['agent A']); + }); + + test('Should reject when invoked without a cell', async () => { + // Catches: falling back to the selected cell, which clears a run the user never clicked. + confirmWith('Clear'); + + await assert.isRejected(invokeCommand(undefined), /requires the cell it was invoked from/); + + verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + }); + }); }); diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts index 4f9d1a8f36..f336454226 100644 --- a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts @@ -5,24 +5,16 @@ import { NotebookCell, NotebookCellStatusBarItem, NotebookCellStatusBarItemProvider, - NotebookEdit, - NotebookRange, - WorkspaceEdit, - commands, l10n, notebooks, - window, workspace } from 'vscode'; import { injectable } from 'inversify'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; -import { getEphemeralCellAgentSourceBlockId, isEphemeralCell } from './dataConversionUtils'; - -const CLEAR_EPHEMERAL_BLOCKS_COMMAND = 'deepnote.clearEphemeralBlocks'; +import { isEphemeralCell } from './dataConversionUtils'; const EPHEMERAL_INDICATOR_PRIORITY = 1000; -const CLEAR_EPHEMERAL_PRIORITY = 990; @injectable() export class EphemeralCellStatusBarProvider @@ -44,49 +36,9 @@ export class EphemeralCellStatusBarProvider }) ); - this.disposables.push( - commands.registerCommand(CLEAR_EPHEMERAL_BLOCKS_COMMAND, async (cell?: NotebookCell) => { - const activeCell = cell || this.getActiveCell(); - if (activeCell) { - await this.clearEphemeralBlocks(activeCell); - } - }) - ); - this.disposables.push(this._onDidChangeCellStatusBarItems); } - /** Deletes the ephemeral cells the same agent block generated, after a modal confirmation. */ - public async clearEphemeralBlocks(cell: NotebookCell): Promise { - if (!isEphemeralCell(cell)) { - return; - } - - const cellsToClear = this.getCellsToClear(cell); - - const confirmation = await window.showWarningMessage( - l10n.t('Clear {0} ephemeral block(s) from this notebook?', cellsToClear.length), - { modal: true }, - l10n.t('Clear') - ); - - if (confirmation !== l10n.t('Clear')) { - return; - } - - // Descending so each deletion's index still addresses the cell it was computed from. - const deletions = [...cellsToClear] - .sort((a, b) => b.index - a.index) - .map((target) => NotebookEdit.deleteCells(new NotebookRange(target.index, target.index + 1))); - - const edit = new WorkspaceEdit(); - edit.set(cell.notebook.uri, deletions); - - if (!(await workspace.applyEdit(edit))) { - void window.showErrorMessage(l10n.t('Failed to clear ephemeral blocks')); - } - } - public dispose(): void { this.disposables.forEach((d) => d.dispose()); } @@ -105,21 +57,7 @@ export class EphemeralCellStatusBarProvider const agentSourceBlockId = cell.metadata?.agent_source_block_id as string | undefined; - return [this.createEphemeralIndicatorItem(agentSourceBlockId ?? null), this.createClearEphemeralItem(cell)]; - } - - private createClearEphemeralItem(cell: NotebookCell): NotebookCellStatusBarItem { - return { - text: `$(trash) ${l10n.t('Clear ephemeral blocks')}`, - alignment: 1, - priority: CLEAR_EPHEMERAL_PRIORITY, - tooltip: l10n.t('Remove the ephemeral blocks generated by this agent block'), - command: { - title: l10n.t('Clear ephemeral blocks'), - command: CLEAR_EPHEMERAL_BLOCKS_COMMAND, - arguments: [cell] - } - }; + return [this.createEphemeralIndicatorItem(agentSourceBlockId ?? null)]; } private createEphemeralIndicatorItem(agentSourceBlockId: string | null): NotebookCellStatusBarItem { @@ -135,26 +73,4 @@ export class EphemeralCellStatusBarProvider tooltip: tooltipLines.join('\n') }; } - - private getActiveCell(): NotebookCell | undefined { - const activeEditor = window.activeNotebookEditor; - if (activeEditor && activeEditor.selection) { - return activeEditor.notebook.cellAt(activeEditor.selection.start); - } - - return undefined; - } - - /** Siblings of the same agent run; just the cell itself when it records no source block. */ - private getCellsToClear(cell: NotebookCell): NotebookCell[] { - const agentSourceBlockId = getEphemeralCellAgentSourceBlockId(cell); - - if (!agentSourceBlockId) { - return [cell]; - } - - return cell.notebook - .getCells() - .filter((candidate) => getEphemeralCellAgentSourceBlockId(candidate) === agentSourceBlockId); - } } diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts index 73bd80414e..42b38a4617 100644 --- a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts @@ -1,11 +1,8 @@ import { expect } from 'chai'; -import * as sinon from 'sinon'; -import { anything, verify, when } from 'ts-mockito'; -import { CancellationToken, NotebookCell, WorkspaceEdit } from 'vscode'; +import { CancellationToken } from 'vscode'; -import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; import { EphemeralCellStatusBarProvider } from './ephemeralCellStatusBarProvider'; -import { createMockCell, createMockNotebookWithCells } from './deepnoteTestHelpers'; +import { createMockCell } from './deepnoteTestHelpers'; suite('EphemeralCellStatusBarProvider', () => { let provider: EphemeralCellStatusBarProvider; @@ -69,27 +66,15 @@ suite('EphemeralCellStatusBarProvider', () => { const cell = createMockCell({ metadata: { is_ephemeral: true } }); const items = provider.provideCellStatusBarItems(cell, mockToken)!; - expect(items).to.have.lengthOf(2); + // Catches: an actionable item creeping back in — clearing is the agent block's button, + // and an ephemeral cell only labels itself. + expect(items).to.have.lengthOf(1); expect(items[0].text).to.include('$(sparkle)'); expect(items[0].text).to.include('Ephemeral'); expect(items[0].alignment).to.equal(1); expect(items[0].priority).to.equal(1000); expect(items[0].command).to.be.undefined; }); - - test('Should set clear button properties and pass the cell to its command', () => { - const cell = createMockCell({ metadata: { is_ephemeral: true } }); - const items = provider.provideCellStatusBarItems(cell, mockToken)!; - - expect(items[1].text).to.include('$(trash)'); - expect(items[1].text).to.include('Clear ephemeral blocks'); - expect(items[1].alignment).to.equal(1); - expect(items[1].priority).to.equal(990); - - const command = items[1].command as { command: string; arguments: unknown[] }; - expect(command.command).to.equal('deepnote.clearEphemeralBlocks'); - expect(command.arguments).to.deep.equal([cell]); - }); }); suite('Tooltip', () => { @@ -133,143 +118,4 @@ suite('EphemeralCellStatusBarProvider', () => { } }); }); - - suite('Clearing ephemeral blocks', () => { - setup(() => { - resetVSCodeMocks(); - }); - - teardown(() => { - sinon.restore(); - resetVSCodeMocks(); - }); - - // Mocked WorkspaceEdit.set drops the edits, so capture them on the prototype and replay the - // deletions against `cells` — an ascending delete order corrupts the survivors, not the count. - function applyDeletionsTo(cells: NotebookCell[]): void { - let recordedEdits: { range: { start: number; end: number } }[] = []; - - sinon.stub(WorkspaceEdit.prototype, 'set').callsFake((_uri, edits) => { - recordedEdits = edits as unknown as { range: { start: number; end: number } }[]; - }); - - when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { - for (const { range } of recordedEdits) { - cells.splice(range.start, range.end - range.start); - } - - return Promise.resolve(true); - }); - } - - function confirmWith(label: string | undefined): void { - when(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).thenReturn( - Promise.resolve(label as any) - ); - } - - function ephemeralCell(text: string, agentSourceBlockId?: string) { - return { - text, - metadata: { - is_ephemeral: true, - ...(agentSourceBlockId ? { agent_source_block_id: agentSourceBlockId } : {}) - } - }; - } - - function agentCell(text: string, blockId: string) { - return { text, metadata: { __deepnotePocket: { type: 'agent' }, id: blockId } }; - } - - test('Should delete every cell generated by the clicked cell’s agent block and nothing else', async () => { - const { cells } = createMockNotebookWithCells([ - agentCell('agent A', 'agent-block-1'), - ephemeralCell('eph A1', 'agent-block-1'), - ephemeralCell('eph A2', 'agent-block-1'), - { text: 'user code', metadata: {} }, - agentCell('agent B', 'agent-block-2'), - ephemeralCell('eph B1', 'agent-block-2') - ]); - const clickedCell = cells[1]; - applyDeletionsTo(cells); - confirmWith('Clear'); - - await provider.clearEphemeralBlocks(clickedCell); - - expect(cells.map((cell) => cell.document.getText())).to.deep.equal([ - 'agent A', - 'user code', - 'agent B', - 'eph B1' - ]); - }); - - test('Should delete only the clicked cell when it records no source agent block', async () => { - const { cells } = createMockNotebookWithCells([ - ephemeralCell('orphan'), - ephemeralCell('eph A1', 'agent-block-1') - ]); - applyDeletionsTo(cells); - confirmWith('Clear'); - - await provider.clearEphemeralBlocks(cells[0]); - - expect(cells.map((cell) => cell.document.getText())).to.deep.equal(['eph A1']); - }); - - test('Should report how many blocks the clear removes', async () => { - const { cells } = createMockNotebookWithCells([ - ephemeralCell('eph A1', 'agent-block-1'), - ephemeralCell('eph A2', 'agent-block-1') - ]); - applyDeletionsTo(cells); - confirmWith('Clear'); - - await provider.clearEphemeralBlocks(cells[0]); - - verify( - mockedVSCodeNamespaces.window.showWarningMessage( - 'Clear 2 ephemeral block(s) from this notebook?', - anything(), - anything() - ) - ).once(); - }); - - test('Should not edit the notebook when the confirmation is dismissed', async () => { - // Catches: applying the edit before the modal is answered, which deletes on a cancel. - const { cells } = createMockNotebookWithCells([ephemeralCell('eph A1', 'agent-block-1')]); - applyDeletionsTo(cells); - confirmWith(undefined); - - await provider.clearEphemeralBlocks(cells[0]); - - verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); - expect(cells).to.have.lengthOf(1); - }); - - test('Should ignore a cell that is not ephemeral', async () => { - // Catches: dropping the isEphemeralCell guard, which would offer to clear any cell. - const { cells } = createMockNotebookWithCells([{ text: 'user code', metadata: {} }]); - applyDeletionsTo(cells); - confirmWith('Clear'); - - await provider.clearEphemeralBlocks(cells[0]); - - verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); - verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); - }); - - test('Should report an error when the workspace edit is rejected', async () => { - // Catches: dropping the `if (!applyEdit)` branch, which loses the clear silently. - const { cells } = createMockNotebookWithCells([ephemeralCell('eph A1', 'agent-block-1')]); - confirmWith('Clear'); - when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(false)); - - await provider.clearEphemeralBlocks(cells[0]); - - verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); - }); - }); }); diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts index c89a058ad0..54a3e6d9e4 100644 --- a/test/e2e/suite/agentBlock.e2e.test.ts +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -51,7 +51,7 @@ const CLEAR_RUN_PYTHON_OUTPUT_MARKER = 'clear-run-python-ran'; const CLEAR_RUN_GENERATED_PYTHON = `print("${CLEAR_RUN_PYTHON_OUTPUT_MARKER}")`; const CLEAR_RUN_MARKDOWN_TEXT = 'Third-run markdown from the E2E agent'; const CLEAR_RUN_FINAL_AGENT_TEXT = 'Clear-run summary added as a markdown block.'; -// EphemeralCellStatusBarProvider button and its confirmation. +// AgentCellStatusBarProvider button and its confirmation. const CLEAR_EPHEMERAL_BUTTON = 'Clear ephemeral blocks'; const CLEAR_EPHEMERAL_CONFIRM = 'Clear'; const CLEAR_EPHEMERAL_CONFIRM_TEXT = 'ephemeral block'; @@ -299,7 +299,7 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu // Self-contained: generates the run it clears, so it survives --grep and a Mocha retry (Run All // drops any stale generated cells first). - it('clears the whole generated run from the ephemeral cell status bar button', async function () { + it('clears the whole generated run from the agent block status bar button', async function () { mockServer = await startMockOpenAiServer([ { match: { hasToolResult: false }, @@ -339,9 +339,9 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu await clickCellStatusBarItem(CLEAR_EPHEMERAL_BUTTON); await confirmModalDialog(CLEAR_EPHEMERAL_CONFIRM, { messageIncludes: CLEAR_EPHEMERAL_CONFIRM_TEXT }); - // The clicked cell is the generated code cell and the markdown cell of the same run goes with - // it. Requiring the agent's own transcript to survive keeps an unreadable webview (which reads - // as '') from passing this as "the generated cells are gone". + // The button lives on the agent block and takes both cells its run generated. Requiring the + // agent's own transcript to survive keeps an unreadable webview (which reads as '') from + // passing this as "the generated cells are gone". await awaitWebviewMarkers([CLEAR_RUN_FINAL_AGENT_TEXT], CLEAR_EPHEMERAL_TIMEOUT, 'ephemeral cells cleared', [ CLEAR_RUN_PYTHON_OUTPUT_MARKER, CLEAR_RUN_MARKDOWN_TEXT From 075ab44fc3addbbdd65c3e0105d1b9295f2c0807 Mon Sep 17 00:00:00 2001 From: tomas Date: Sat, 8 Aug 2026 16:01:53 +0000 Subject: [PATCH 55/80] feat(agent-block): add a command and toolbar button to create agent blocks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit There was no way to create an agent block from the extension — one could only arrive by opening a .deepnote file that already contained it. Unlike the sibling add*Block commands this mints the block id up front. createBlockFromPocket hands an id-less block a fresh random id on every call, so each run would stamp its generated cells with a different owner: the stale-run guard would never match, removeEphemeralCellsForAgentBlocks skips id-less agents outright, and scratch cells would accumulate on every run until the first save and reload. The toolbar button takes navigation@3, ahead of the other block buttons. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- package.json | 21 +++- package.nls.json | 1 + .../deepnoteNotebookCommandListener.ts | 44 ++++++++ ...epnoteNotebookCommandListener.unit.test.ts | 106 ++++++++++++++++++ src/platform/common/constants.ts | 1 + 5 files changed, 168 insertions(+), 5 deletions(-) diff --git a/package.json b/package.json index 8b89c5523d..876ffea3f9 100644 --- a/package.json +++ b/package.json @@ -173,6 +173,12 @@ "category": "Deepnote", "icon": "$(notebook)" }, + { + "command": "deepnote.addAgentBlock", + "title": "%deepnote.commands.addAgentBlock.title%", + "category": "Deepnote", + "icon": "$(hubot)" + }, { "command": "deepnote.addSqlBlock", "title": "%deepnote.commands.addSqlBlock.title%", @@ -1032,30 +1038,35 @@ "when": "notebookType == 'deepnote'" }, { - "command": "deepnote.addSqlBlock", + "command": "deepnote.addAgentBlock", "group": "navigation@3", "when": "notebookType == 'deepnote'" }, { - "command": "deepnote.addChartBlock", + "command": "deepnote.addSqlBlock", "group": "navigation@4", "when": "notebookType == 'deepnote'" }, { - "command": "deepnote.addBigNumberChartBlock", + "command": "deepnote.addChartBlock", "group": "navigation@5", "when": "notebookType == 'deepnote'" }, { - "command": "deepnote.addInputBlock", + "command": "deepnote.addBigNumberChartBlock", "group": "navigation@6", "when": "notebookType == 'deepnote'" }, { - "command": "deepnote.addTextBlock", + "command": "deepnote.addInputBlock", "group": "navigation@7", "when": "notebookType == 'deepnote'" }, + { + "command": "deepnote.addTextBlock", + "group": "navigation@8", + "when": "notebookType == 'deepnote'" + }, { "command": "deepnote.restartkernel", "group": "navigation/execute@5", diff --git a/package.nls.json b/package.nls.json index 7d98f46777..b22f804064 100644 --- a/package.nls.json +++ b/package.nls.json @@ -259,6 +259,7 @@ "deepnote.commands.newProject.title": "New Project", "deepnote.commands.importNotebook.title": "Import Notebook", "deepnote.commands.importJupyterNotebook.title": "Import Jupyter Notebook", + "deepnote.commands.addAgentBlock.title": "Add Agent Block", "deepnote.commands.addSqlBlock.title": "Add SQL Block", "deepnote.commands.addBigNumberChartBlock.title": "Add Big Number Block", "deepnote.commands.addChartBlock.title": "Add Chart Block", diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts index d91608c15a..2165cbb869 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts @@ -37,6 +37,8 @@ import { } from './deepnoteSchemas'; import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes'; import { Pocket } from '../../platform/deepnote/pocket'; +import { AGENT_MODEL_AUTO, AGENT_MODEL_METADATA_KEY } from './agentCellStatusBarProvider'; +import { generateBlockId } from './dataConversionUtils'; export const INPUT_BLOCK_TYPES = [ 'input-text', @@ -163,6 +165,7 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation } private registerCommands(): void { + this.disposableRegistry.push(commands.registerCommand(Commands.AddAgentBlock, () => this.addAgentBlock())); this.disposableRegistry.push(commands.registerCommand(Commands.AddSqlBlock, () => this.addSqlBlock())); this.disposableRegistry.push( commands.registerCommand(Commands.AddBigNumberChartBlock, () => this.addBigNumberChartBlock()) @@ -227,6 +230,47 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation ); } + /** + * Inserts an empty agent block below the selection. + * + * Unlike the other block commands this mints the block id up front: an agent block without one + * gets a fresh random id on every `convertCellToBlock`, so each run would stamp its generated + * cells with a different owner and the stale-run cleanup would never match them. + */ + public async addAgentBlock(): Promise { + const editor = window.activeNotebookEditor; + if (!editor) { + throw new Error(l10n.t('No active notebook editor found')); + } + const document = editor.notebook; + const selection = editor.selection; + + const insertIndex = selection ? selection.end : document.cellCount; + const blockId = generateBlockId(); + + const result = await notebookUpdaterUtils.chainWithPendingUpdates(document, (edit) => { + const newCell = new NotebookCellData(NotebookCellKind.Code, '', 'plaintext'); + newCell.metadata = { + __deepnotePocket: { + type: 'agent' + }, + id: blockId, + __deepnoteBlockId: blockId, + [AGENT_MODEL_METADATA_KEY]: AGENT_MODEL_AUTO + }; + const nbEdit = NotebookEdit.insertCells(insertIndex, [newCell]); + edit.set(document.uri, [nbEdit]); + }); + if (result !== true) { + throw new Error(l10n.t('Failed to insert agent block')); + } + + const notebookRange = new NotebookRange(insertIndex, insertIndex + 1); + editor.revealRange(notebookRange, NotebookEditorRevealType.Default); + editor.selection = notebookRange; + await commands.executeCommand('notebook.cell.edit'); + } + public async addSqlBlock(): Promise { const editor = window.activeNotebookEditor; if (!editor) { diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts index ff600a2bf9..d214b0b3b8 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts @@ -838,6 +838,112 @@ suite('DeepnoteNotebookCommandListener', () => { }); }); + suite('addAgentBlock', () => { + function insertedCell(getCapturedNotebookEdits: () => any[] | null) { + const edits = getCapturedNotebookEdits()!; + assert.equal(edits.length, 1, 'Should have one notebook edit'); + + const notebookEdit = edits[0] as any; + assert.equal(notebookEdit.newCells.length, 1, 'Should insert one cell'); + + return notebookEdit.newCells[0]; + } + + test('should add an empty plaintext agent block at the end when no selection exists', async () => { + const { editor, document } = createMockEditor([], undefined); + const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + assert.isTrue(chainStub.calledOnce, 'chainWithPendingUpdates should be called once'); + assert.equal(chainStub.firstCall.args[0], document, 'Should be called with correct document'); + + const newCell = insertedCell(getCapturedNotebookEdits); + assert.equal(newCell.kind, NotebookCellKind.Code, 'Should be a code cell'); + assert.equal(newCell.languageId, 'plaintext', 'Should have plaintext language'); + assert.equal(newCell.value, '', 'Should have empty content'); + assert.equal(newCell.metadata.__deepnotePocket.type, 'agent', 'Should have agent pocket type'); + + assert.isTrue((editor.revealRange as sinon.SinonStub).calledOnce, 'Should reveal the new cell range'); + const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + assert.equal(revealCall.args[0].start, 0, 'Should reveal correct range start'); + assert.equal(revealCall.args[0].end, 1, 'Should reveal correct range end'); + }); + + test('should add the agent block after the selection when one exists', async () => { + const existingCells = [createMockCell('{}'), createMockCell('{}')]; + const { editor, document } = createMockEditor(existingCells, new NotebookRange(1, 2)); + const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + insertedCell(getCapturedNotebookEdits); + assert.equal(document.uri.fsPath, '/test/notebook.ipynb', 'Should edit the active document'); + + const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + assert.equal(revealCall.args[0].start, 2, 'Should insert below the selection'); + assert.equal(revealCall.args[0].end, 3, 'Should select only the new cell'); + }); + + test('should mint a block id under both keys so runs keep a stable owner', async () => { + // Catches: an id-less agent block, which gets a fresh random id on every + // convertCellToBlock — its generated cells would never be matched back to it. + const { editor } = createMockEditor([], undefined); + const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + const { metadata } = insertedCell(getCapturedNotebookEdits); + assert.match(metadata.id, /^[0-9a-f]{32}$/, 'Should mint a 32-char hex block id'); + assert.equal(metadata.__deepnoteBlockId, metadata.id, 'Backup id key must match id'); + }); + + test('should give consecutive agent blocks distinct ids', async () => { + // Catches: a hoisted/constant id, which would make two agent blocks fight over the + // same generated cells. + const { editor } = createMockEditor([], undefined); + const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + const first = insertedCell(getCapturedNotebookEdits).metadata.id; + + await commandListener.addAgentBlock(); + const second = insertedCell(getCapturedNotebookEdits).metadata.id; + + assert.notEqual(first, second, 'Each agent block needs its own id'); + }); + + test('should persist the default model rather than leaving the key absent', async () => { + // Catches: omitting deepnote_agent_model, which reaches openai() as undefined. + const { editor } = createMockEditor([], undefined); + const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + const { metadata } = insertedCell(getCapturedNotebookEdits); + assert.equal(metadata.deepnote_agent_model, 'auto', 'Should persist the auto default'); + }); + + test('should throw error when no active editor exists', async () => { + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(undefined); + + await assert.isRejected(commandListener.addAgentBlock(), Error, 'No active notebook editor found'); + }); + + test('should throw error when chainWithPendingUpdates fails', async () => { + const { editor } = createMockEditor([], undefined); + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(editor); + + sandbox.replace( + notebookUpdater.notebookUpdaterUtils, + 'chainWithPendingUpdates', + sinon.stub().resolves(false) + ); + + await assert.isRejected(commandListener.addAgentBlock(), Error, 'Failed to insert agent block'); + }); + }); + suite('addBigNumberChartBlock', () => { test('should add big number block at the end when no selection exists', async () => { // Setup mocks diff --git a/src/platform/common/constants.ts b/src/platform/common/constants.ts index 39f3381a0c..fc87a3e1bb 100644 --- a/src/platform/common/constants.ts +++ b/src/platform/common/constants.ts @@ -228,6 +228,7 @@ export namespace Commands { export const DisableSnapshots = 'deepnote.disableSnapshots'; export const AuthenticateIntegration = 'deepnote.authenticateIntegration'; export const ManageIntegrations = 'deepnote.manageIntegrations'; + export const AddAgentBlock = 'deepnote.addAgentBlock'; export const AddSqlBlock = 'deepnote.addSqlBlock'; export const AddBigNumberChartBlock = 'deepnote.addBigNumberChartBlock'; export const AddChartBlock = 'deepnote.addChartBlock'; From adc2f50c0081a355c76ea8675e448c18f4ff4218 Mon Sep 17 00:00:00 2001 From: tomas Date: Sat, 8 Aug 2026 21:31:22 +0000 Subject: [PATCH 56/80] feat(agent-block): allow only one agent block per notebook A second request reports the limit and leaves the notebook untouched rather than failing, since asking twice is a reasonable thing to do. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- .../deepnoteNotebookCommandListener.ts | 11 +++- ...epnoteNotebookCommandListener.unit.test.ts | 50 ++++++++++++++++++- 2 files changed, 57 insertions(+), 4 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts index 2165cbb869..55d428616b 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts @@ -38,7 +38,7 @@ import { import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes'; import { Pocket } from '../../platform/deepnote/pocket'; import { AGENT_MODEL_AUTO, AGENT_MODEL_METADATA_KEY } from './agentCellStatusBarProvider'; -import { generateBlockId } from './dataConversionUtils'; +import { generateBlockId, isAgentCell } from './dataConversionUtils'; export const INPUT_BLOCK_TYPES = [ 'input-text', @@ -231,7 +231,8 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation } /** - * Inserts an empty agent block below the selection. + * Inserts an empty agent block below the selection. A notebook may hold at most one; a second + * request reports that and leaves the notebook untouched. * * Unlike the other block commands this mints the block id up front: an agent block without one * gets a fresh random id on every `convertCellToBlock`, so each run would stamp its generated @@ -245,6 +246,12 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation const document = editor.notebook; const selection = editor.selection; + if (document.getCells().some(isAgentCell)) { + void window.showInformationMessage(l10n.t('This notebook already contains an agent block.')); + + return; + } + const insertIndex = selection ? selection.end : document.cellCount; const blockId = generateBlockId(); diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts index d214b0b3b8..cf0c18ce92 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts @@ -1,6 +1,6 @@ import { assert } from 'chai'; import * as sinon from 'sinon'; -import { when, reset, anything } from 'ts-mockito'; +import { when, reset, anything, verify } from 'ts-mockito'; import { NotebookCell, NotebookDocument, @@ -24,7 +24,7 @@ import { createMockedNotebookDocument } from '../../test/datascience/editor-inte import { WrappedError } from '../../platform/errors/types'; import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes'; import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; -import { createMockCell } from './deepnoteTestHelpers'; +import { createMockCell, createMockNotebookWithCells } from './deepnoteTestHelpers'; suite('DeepnoteNotebookCommandListener', () => { let commandListener: DeepnoteNotebookCommandListener; @@ -839,6 +839,23 @@ suite('DeepnoteNotebookCommandListener', () => { }); suite('addAgentBlock', () => { + // createMockedNotebookDocument drops NotebookCellData.metadata, so an existing agent + // block has to come from a notebook mock that keeps it. + function createMockEditorWithMetadata( + cellOptions: Parameters[0] + ): NotebookEditor { + const { notebook } = createMockNotebookWithCells(cellOptions); + const selection = new NotebookRange(0, cellOptions.length > 0 ? 1 : 0); + + return { + notebook, + selection, + selections: [selection], + visibleRanges: [], + revealRange: sandbox.stub() + }; + } + function insertedCell(getCapturedNotebookEdits: () => any[] | null) { const edits = getCapturedNotebookEdits()!; assert.equal(edits.length, 1, 'Should have one notebook edit'); @@ -924,6 +941,35 @@ suite('DeepnoteNotebookCommandListener', () => { assert.equal(metadata.deepnote_agent_model, 'auto', 'Should persist the auto default'); }); + test('should refuse a second agent block and leave the notebook untouched', async () => { + const editor = createMockEditorWithMetadata([ + { text: 'existing agent', metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' } }, + { text: 'user code', metadata: {} } + ]); + const { chainStub } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + assert.isFalse(chainStub.called, 'Must not edit a notebook that already has an agent block'); + verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); + assert.isFalse((editor.revealRange as sinon.SinonStub).called, 'Must not reveal anything'); + }); + + test('should still add the block when other cells carry no agent pocket', async () => { + // Catches: a guard that trips on any cell, blocking the first agent block outright. + const editor = createMockEditorWithMetadata([ + { text: 'user code', metadata: { __deepnotePocket: { type: 'code' }, id: 'code-block-1' } }, + { text: 'scratch', metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' } } + ]); + const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + assert.isTrue(chainStub.calledOnce, 'Should insert the first agent block'); + assert.equal(insertedCell(getCapturedNotebookEdits).metadata.__deepnotePocket.type, 'agent'); + verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).never(); + }); + test('should throw error when no active editor exists', async () => { when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(undefined); From c12c0da3ff81038586f67e67e629149b13013cfb Mon Sep 17 00:00:00 2001 From: tomas Date: Sun, 9 Aug 2026 09:24:01 +0000 Subject: [PATCH 57/80] test(agent-block): fold the metadata-preserving mock into createMockEditor createMockedNotebookDocument drops NotebookCellData.metadata, which the block commands read, so the agent-block guard needed a second editor helper to express an existing agent cell. Two helpers where one is strictly lossier invites picking the wrong one, so createMockEditor now builds on createMockNotebookWithCells and the duplicate is gone. The mock notebook is a .deepnote file rather than an .ipynb, so the one test asserting the document path now checks the document identity that chainWithPendingUpdates received instead. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- ...epnoteNotebookCommandListener.unit.test.ts | 52 ++++++++----------- 1 file changed, 21 insertions(+), 31 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts index cf0c18ce92..0e9c7a591e 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts @@ -8,8 +8,7 @@ import { NotebookRange, NotebookCellKind, NotebookCellData, - WorkspaceEdit, - Uri + WorkspaceEdit } from 'vscode'; import { @@ -20,7 +19,6 @@ import { import { formatInputBlockCellContent, getInputBlockLanguage } from './inputBlockContentFormatter'; import { IConfigurationService, IDisposable } from '../../platform/common/types'; import * as notebookUpdater from '../../kernels/execution/notebookUpdater'; -import { createMockedNotebookDocument } from '../../test/datascience/editor-integration/helpers'; import { WrappedError } from '../../platform/errors/types'; import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes'; import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; @@ -362,7 +360,10 @@ suite('DeepnoteNotebookCommandListener', () => { } /** - * Helper to create mock NotebookEditor and NotebookDocument + * Helper to create mock NotebookEditor and NotebookDocument. + * + * Built on createMockNotebookWithCells rather than createMockedNotebookDocument because the + * latter drops NotebookCellData.metadata, which the block commands read. */ function createMockEditor( cellDataArray: NotebookCellData[], @@ -371,8 +372,14 @@ suite('DeepnoteNotebookCommandListener', () => { editor: NotebookEditor; document: NotebookDocument; } { - const uri = Uri.file('/test/notebook.ipynb'); - const document = createMockedNotebookDocument(cellDataArray, {}, uri); + const { notebook: document } = createMockNotebookWithCells( + cellDataArray.map((data) => ({ + kind: data.kind, + languageId: data.languageId, + text: data.value, + metadata: data.metadata + })) + ); const editorSelection = selection != null ? selection : new NotebookRange(0, cellDataArray.length > 0 ? 1 : 0); @@ -839,23 +846,6 @@ suite('DeepnoteNotebookCommandListener', () => { }); suite('addAgentBlock', () => { - // createMockedNotebookDocument drops NotebookCellData.metadata, so an existing agent - // block has to come from a notebook mock that keeps it. - function createMockEditorWithMetadata( - cellOptions: Parameters[0] - ): NotebookEditor { - const { notebook } = createMockNotebookWithCells(cellOptions); - const selection = new NotebookRange(0, cellOptions.length > 0 ? 1 : 0); - - return { - notebook, - selection, - selections: [selection], - visibleRanges: [], - revealRange: sandbox.stub() - }; - } - function insertedCell(getCapturedNotebookEdits: () => any[] | null) { const edits = getCapturedNotebookEdits()!; assert.equal(edits.length, 1, 'Should have one notebook edit'); @@ -890,12 +880,12 @@ suite('DeepnoteNotebookCommandListener', () => { test('should add the agent block after the selection when one exists', async () => { const existingCells = [createMockCell('{}'), createMockCell('{}')]; const { editor, document } = createMockEditor(existingCells, new NotebookRange(1, 2)); - const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); await commandListener.addAgentBlock(); insertedCell(getCapturedNotebookEdits); - assert.equal(document.uri.fsPath, '/test/notebook.ipynb', 'Should edit the active document'); + assert.equal(chainStub.firstCall.args[0], document, 'Should edit the active document'); const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; assert.equal(revealCall.args[0].start, 2, 'Should insert below the selection'); @@ -942,9 +932,9 @@ suite('DeepnoteNotebookCommandListener', () => { }); test('should refuse a second agent block and leave the notebook untouched', async () => { - const editor = createMockEditorWithMetadata([ - { text: 'existing agent', metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' } }, - { text: 'user code', metadata: {} } + const { editor } = createMockEditor([ + createMockCell('existing agent', { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }), + createMockCell('user code') ]); const { chainStub } = mockNotebookUpdateAndExecute(editor); @@ -957,9 +947,9 @@ suite('DeepnoteNotebookCommandListener', () => { test('should still add the block when other cells carry no agent pocket', async () => { // Catches: a guard that trips on any cell, blocking the first agent block outright. - const editor = createMockEditorWithMetadata([ - { text: 'user code', metadata: { __deepnotePocket: { type: 'code' }, id: 'code-block-1' } }, - { text: 'scratch', metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' } } + const { editor } = createMockEditor([ + createMockCell('user code', { __deepnotePocket: { type: 'code' }, id: 'code-block-1' }), + createMockCell('scratch', { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }) ]); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); From e64e652e93ecada0ff1c604e3c17d31d39c379fd Mon Sep 17 00:00:00 2001 From: tomas Date: Sun, 9 Aug 2026 15:11:54 +0000 Subject: [PATCH 58/80] test: carry cell metadata through createMockedNotebookDocument MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The helper takes NotebookCellData but never stubbed cell.metadata, and an unstubbed ts-mockito member is a function, so every metadata read came back undefined instead of throwing. A guard keyed on metadata could not trip and its test passed for the wrong reason. No caller depended on the old behaviour. dataframeController.ts matches cells by c.metadata.id and its tests already worked around this with their own cell builder, which is the second local workaround this gap has produced. cell.outputs has the same problem — data.outputs is dropped — but wiring it up shifts the output counts two Cell Execution Message Handler tests assert on, so that one needs its own change. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- src/test/datascience/editor-integration/helpers.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/src/test/datascience/editor-integration/helpers.ts b/src/test/datascience/editor-integration/helpers.ts index bfd90523ad..2a01e61f2d 100644 --- a/src/test/datascience/editor-integration/helpers.ts +++ b/src/test/datascience/editor-integration/helpers.ts @@ -191,6 +191,7 @@ export function createMockedNotebookDocument( when(cell.document).thenReturn(mockedDocument); when(cell.index).thenReturn(index); when(cell.kind).thenReturn(data.kind); + when(cell.metadata).thenReturn(data.metadata ?? {}); const cellOutput: NotebookCellOutput[] = []; when(cell.outputs).thenReturn(cellOutput); when(cell.notebook).thenReturn(instance(notebook)); From dd9547e2e1907843fb0fecf233afcc95d851d7ac Mon Sep 17 00:00:00 2001 From: tomas Date: Sun, 9 Aug 2026 15:11:54 +0000 Subject: [PATCH 59/80] test(agent-block): rename the local cell builder to createMockCellData MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit It shadowed the createMockCell imported from deepnoteTestHelpers with a different signature and a different return type — NotebookCellData rather than a mocked NotebookCell — so which one a call site meant depended on where in the file it sat. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- ...epnoteNotebookCommandListener.unit.test.ts | 52 +++++++++---------- 1 file changed, 26 insertions(+), 26 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts index 0e9c7a591e..978f2406b8 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts @@ -349,9 +349,9 @@ suite('DeepnoteNotebookCommandListener', () => { }); /** - * Helper to create mock NotebookCell with metadata + * Helper to create NotebookCellData with metadata, for seeding createMockEditor. */ - function createMockCell(content: string, metadata?: Record): NotebookCellData { + function createMockCellData(content: string, metadata?: Record): NotebookCellData { const cell = new NotebookCellData(NotebookCellKind.Code, content, 'json'); if (metadata != null) { cell.metadata = metadata; @@ -450,7 +450,7 @@ suite('DeepnoteNotebookCommandListener', () => { { description: 'should add input-text block after selection when selection exists', blockType: 'input-text', - existingCells: [createMockCell('{}')], + existingCells: [createMockCellData('{}')], selection: new NotebookRange(0, 1), expectedInsertIndex: 1, expectedVariableName: 'input_1', @@ -559,8 +559,8 @@ suite('DeepnoteNotebookCommandListener', () => { description: 'should generate correct variable name when existing inputs exist', blockType: 'input-text', existingCells: [ - createMockCell('{ "deepnote_variable_name": "input_1" }'), - createMockCell('{ "deepnote_variable_name": "input_2" }') + createMockCellData('{ "deepnote_variable_name": "input_1" }'), + createMockCellData('{ "deepnote_variable_name": "input_2" }') ], selection: new NotebookRange(1, 2), expectedInsertIndex: 2, @@ -570,7 +570,7 @@ suite('DeepnoteNotebookCommandListener', () => { { description: 'should insert at selection.end when selection is in the middle', blockType: 'input-text', - existingCells: [createMockCell('{}'), createMockCell('{}'), createMockCell('{}')], + existingCells: [createMockCellData('{}'), createMockCellData('{}'), createMockCellData('{}')], selection: new NotebookRange(1, 2), expectedInsertIndex: 2, expectedVariableName: 'input_1', @@ -579,7 +579,7 @@ suite('DeepnoteNotebookCommandListener', () => { { description: 'should handle large variable numbers correctly', blockType: 'input-text', - existingCells: [createMockCell('{ "deepnote_variable_name": "input_99" }')], + existingCells: [createMockCellData('{ "deepnote_variable_name": "input_99" }')], selection: undefined, expectedInsertIndex: 1, expectedVariableName: 'input_100', @@ -760,7 +760,7 @@ suite('DeepnoteNotebookCommandListener', () => { test('should add SQL block after selection when selection exists', async () => { // Setup mocks - const existingCells = [createMockCell('{}'), createMockCell('{}')]; + const existingCells = [createMockCellData('{}'), createMockCellData('{}')]; const selection = new NotebookRange(1, 2); const { editor } = createMockEditor(existingCells, selection); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -783,8 +783,8 @@ suite('DeepnoteNotebookCommandListener', () => { test('should generate correct variable name when existing df variables exist', async () => { // Setup mocks with existing df variables const existingCells = [ - createMockCell('{ "deepnote_variable_name": "df_1" }'), - createMockCell('{ "deepnote_variable_name": "df_2" }') + createMockCellData('{ "deepnote_variable_name": "df_1" }'), + createMockCellData('{ "deepnote_variable_name": "df_2" }') ]; const { editor } = createMockEditor(existingCells, undefined); const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -803,8 +803,8 @@ suite('DeepnoteNotebookCommandListener', () => { test('should ignore input variables when generating df variable name', async () => { // Setup mocks with input variables (should not affect df numbering) const existingCells = [ - createMockCell('{ "deepnote_variable_name": "input_10" }'), - createMockCell('{ "deepnote_variable_name": "df_2" }') + createMockCellData('{ "deepnote_variable_name": "input_10" }'), + createMockCellData('{ "deepnote_variable_name": "df_2" }') ]; const { editor } = createMockEditor(existingCells, undefined); const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -878,7 +878,7 @@ suite('DeepnoteNotebookCommandListener', () => { }); test('should add the agent block after the selection when one exists', async () => { - const existingCells = [createMockCell('{}'), createMockCell('{}')]; + const existingCells = [createMockCellData('{}'), createMockCellData('{}')]; const { editor, document } = createMockEditor(existingCells, new NotebookRange(1, 2)); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -933,8 +933,8 @@ suite('DeepnoteNotebookCommandListener', () => { test('should refuse a second agent block and leave the notebook untouched', async () => { const { editor } = createMockEditor([ - createMockCell('existing agent', { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }), - createMockCell('user code') + createMockCellData('existing agent', { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }), + createMockCellData('user code') ]); const { chainStub } = mockNotebookUpdateAndExecute(editor); @@ -948,8 +948,8 @@ suite('DeepnoteNotebookCommandListener', () => { test('should still add the block when other cells carry no agent pocket', async () => { // Catches: a guard that trips on any cell, blocking the first agent block outright. const { editor } = createMockEditor([ - createMockCell('user code', { __deepnotePocket: { type: 'code' }, id: 'code-block-1' }), - createMockCell('scratch', { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }) + createMockCellData('user code', { __deepnotePocket: { type: 'code' }, id: 'code-block-1' }), + createMockCellData('scratch', { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }) ]); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -1027,7 +1027,7 @@ suite('DeepnoteNotebookCommandListener', () => { test('should add big number block after selection when selection exists', async () => { // Setup mocks - const existingCells = [createMockCell('{}'), createMockCell('{}')]; + const existingCells = [createMockCellData('{}'), createMockCellData('{}')]; const selection = new NotebookRange(0, 1); const { editor } = createMockEditor(existingCells, selection); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -1049,7 +1049,7 @@ suite('DeepnoteNotebookCommandListener', () => { test('should insert at correct position in the middle of notebook', async () => { // Setup mocks - const existingCells = [createMockCell('{}'), createMockCell('{}'), createMockCell('{}')]; + const existingCells = [createMockCellData('{}'), createMockCellData('{}'), createMockCellData('{}')]; const selection = new NotebookRange(1, 2); const { editor } = createMockEditor(existingCells, selection); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -1163,7 +1163,7 @@ suite('DeepnoteNotebookCommandListener', () => { test('should add chart block after selection when selection exists', async () => { // Setup mocks - const existingCells = [createMockCell('{}'), createMockCell('{}')]; + const existingCells = [createMockCellData('{}'), createMockCellData('{}')]; const selection = new NotebookRange(0, 1); const { editor } = createMockEditor(existingCells, selection); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -1186,8 +1186,8 @@ suite('DeepnoteNotebookCommandListener', () => { test('should use hardcoded variable name df_1', async () => { // Setup mocks with existing df variables const existingCells = [ - createMockCell('{ "deepnote_variable_name": "df_1" }'), - createMockCell('{ "variable": "df_2" }') + createMockCellData('{ "deepnote_variable_name": "df_1" }'), + createMockCellData('{ "variable": "df_2" }') ]; const { editor } = createMockEditor(existingCells, undefined); const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -1207,9 +1207,9 @@ suite('DeepnoteNotebookCommandListener', () => { test('should always use df_1 regardless of existing variables', async () => { // Setup mocks with various existing variables const existingCells = [ - createMockCell('{ "deepnote_variable_name": "input_10" }'), - createMockCell('{ "deepnote_variable_name": "df_5" }'), - createMockCell('{ "variable": "df_2" }') + createMockCellData('{ "deepnote_variable_name": "input_10" }'), + createMockCellData('{ "deepnote_variable_name": "df_5" }'), + createMockCellData('{ "variable": "df_2" }') ]; const { editor } = createMockEditor(existingCells, undefined); const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -1228,7 +1228,7 @@ suite('DeepnoteNotebookCommandListener', () => { test('should insert at correct position in the middle of notebook', async () => { // Setup mocks - const existingCells = [createMockCell('{}'), createMockCell('{}'), createMockCell('{}')]; + const existingCells = [createMockCellData('{}'), createMockCellData('{}'), createMockCellData('{}')]; const selection = new NotebookRange(1, 2); const { editor } = createMockEditor(existingCells, selection); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); From ab7ba2b308fede50eeebcd0aa60fce462ba96962 Mon Sep 17 00:00:00 2001 From: tomas Date: Sun, 9 Aug 2026 18:33:25 +0000 Subject: [PATCH 60/80] refactor(agent-block): validate the switch-model command like the clear one MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Both commands in this provider are reached only from their status bar item, which always passes the cell, and neither is contributed to the palette — so the getActiveCell fallback could never fire and the guard around it turned a wiring bug into a silent no-op. Throw instead, matching the clear command, and drop the dead lookup. The command id becomes a constant so the registration, the button and the error message cannot drift apart. The handler-capture setup moves up to the suite level now that both commands need it, following the captureCommandHandlers pattern in deepnoteExplorerView.unit.test.ts. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- .../deepnote/agentCellStatusBarProvider.ts | 21 ++--- .../agentCellStatusBarProvider.unit.test.ts | 81 ++++++++++++++----- 2 files changed, 70 insertions(+), 32 deletions(-) diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index b4ead0ba8d..75a9f5224d 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -28,6 +28,7 @@ export const AGENT_MODEL_AUTO = 'auto'; const AGENT_MODEL_OPTIONS = [AGENT_MODEL_AUTO, 'gpt-5.6-sol', 'gpt-5.6-terra', 'gpt-5.6-luna']; const CLEAR_EPHEMERAL_BLOCKS_COMMAND = 'deepnote.clearEphemeralBlocks'; +const SWITCH_AGENT_MODEL_COMMAND = 'deepnote.switchAgentModel'; const AGENT_INDICATOR_PRIORITY = 100; const MODEL_PICKER_PRIORITY = 90; @@ -52,11 +53,12 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv ); this.disposables.push( - commands.registerCommand('deepnote.switchAgentModel', async (cell?: NotebookCell) => { - const activeCell = cell || this.getActiveCell(); - if (activeCell) { - await this.switchModel(activeCell); + commands.registerCommand(SWITCH_AGENT_MODEL_COMMAND, async (cell?: NotebookCell) => { + if (!cell) { + throw new Error(`${SWITCH_AGENT_MODEL_COMMAND} requires the cell it was invoked from`); } + + await this.switchModel(cell); }) ); @@ -167,21 +169,12 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv tooltip: l10n.t('AI Model: {0}\nClick to change', model), command: { title: l10n.t('Switch Model'), - command: 'deepnote.switchAgentModel', + command: SWITCH_AGENT_MODEL_COMMAND, arguments: [cell] } }; } - private getActiveCell(): NotebookCell | undefined { - const activeEditor = window.activeNotebookEditor; - if (activeEditor && activeEditor.selection) { - return activeEditor.notebook.cellAt(activeEditor.selection.start); - } - - return undefined; - } - /** Ephemeral cells this agent block generated; empty when it has no block id or has not run. */ private getCellsToClear(cell: NotebookCell): NotebookCell[] { const agentBlockId = getBlockId(cell); diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts index 04b6f9234e..565261e2d9 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -4,13 +4,45 @@ import { anything, verify, when } from 'ts-mockito'; import { CancellationToken, NotebookCell, NotebookEdit, WorkspaceEdit } from 'vscode'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; -import { AgentCellStatusBarProvider } from './agentCellStatusBarProvider'; +import { AGENT_MODEL_METADATA_KEY, AgentCellStatusBarProvider } from './agentCellStatusBarProvider'; import { createMockCell, createMockNotebookWithCells } from './deepnoteTestHelpers'; suite('AgentCellStatusBarProvider', () => { let provider: AgentCellStatusBarProvider; let mockToken: CancellationToken; + const commandHandlers = new Map Promise>(); + + // Records every command registration. Call AFTER resetVSCodeMocks(), which regenerates the mocks. + function activateCapturingCommands(): void { + commandHandlers.clear(); + when( + mockedVSCodeNamespaces.notebooks.registerNotebookCellStatusBarItemProvider(anything(), anything()) + ).thenReturn({ dispose: () => undefined }); + when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument).thenReturn(() => ({ + dispose: () => undefined + })); + when(mockedVSCodeNamespaces.commands.registerCommand(anything(), anything())).thenCall( + (id: string, callback: (cell?: NotebookCell) => Promise) => { + commandHandlers.set(id, callback); + + return { dispose: () => undefined }; + } + ); + + provider.activate(); + } + + function handlerFor(id: string): (cell?: NotebookCell) => Promise { + const handler = commandHandlers.get(id); + + if (!handler) { + throw new Error(`No handler captured for '${id}'; call activateCapturingCommands() first.`); + } + + return handler; + } + setup(() => { mockToken = { isCancellationRequested: false, @@ -231,6 +263,34 @@ suite('AgentCellStatusBarProvider', () => { verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).never(); verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); }); + + suite('Command handler', () => { + let invokeCommand: (cell?: NotebookCell) => Promise; + + setup(() => { + activateCapturingCommands(); + invokeCommand = handlerFor('deepnote.switchAgentModel'); + }); + + test('Should switch the model of the cell the command is given', async () => { + pick('gpt-5.6-terra'); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(true)); + + await invokeCommand(agentCell()); + + expect(capturedEdit!.metadata[AGENT_MODEL_METADATA_KEY]).to.equal('gpt-5.6-terra'); + }); + + test('Should reject when invoked without a cell', async () => { + // Catches: falling back to the selected cell, which rewrites a block the user never clicked. + pick('gpt-5.6-terra'); + + await assert.isRejected(invokeCommand(undefined), /requires the cell it was invoked from/); + + verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + }); }); suite('Clearing ephemeral blocks', () => { @@ -420,23 +480,8 @@ suite('AgentCellStatusBarProvider', () => { let invokeCommand: (cell?: NotebookCell) => Promise; setup(() => { - when( - mockedVSCodeNamespaces.notebooks.registerNotebookCellStatusBarItemProvider(anything(), anything()) - ).thenReturn({ dispose: () => undefined }); - when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument).thenReturn(() => ({ - dispose: () => undefined - })); - when(mockedVSCodeNamespaces.commands.registerCommand(anything(), anything())).thenCall( - (id: string, callback: (cell?: NotebookCell) => Promise) => { - if (id === 'deepnote.clearEphemeralBlocks') { - invokeCommand = callback; - } - - return { dispose: () => undefined }; - } - ); - - provider.activate(); + activateCapturingCommands(); + invokeCommand = handlerFor('deepnote.clearEphemeralBlocks'); }); test('Should clear the blocks of the cell the command is given', async () => { From 0aaeda94255503c0405f7df360df2910d661669b Mon Sep 17 00:00:00 2001 From: tomas Date: Sun, 9 Aug 2026 18:51:56 +0000 Subject: [PATCH 61/80] test(agent-block): drive the provider suites through their commands MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The switch-model tests already covered switchModel five ways; adding a parallel command-handler suite restated the weakest of them. Point the existing cases at the registered command instead, so each one covers the wiring the status bar item actually goes through, and keep only the missing-cell rejection as handler-specific. Same for the clear suite. The five detection tests differed only by input metadata, so they fold into one table. The "cell without metadata" case went with them: createMockCell turns an explicit undefined into {}, so it built the same cell as "cell without a pocket" and could not fail independently. switchModel and clearEphemeralBlocks go back to private — the commands are the only callers now, which is what 28dd2fc1b opened switchModel up for. clearEphemeralBlocks moves down to keep members ordered by accessibility then name. 28 tests to 22, 511 lines to 463, with the wiring better covered than before. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- .vscode/launch.json | 2 +- .../deepnote/agentCellStatusBarProvider.ts | 60 ++++---- .../agentCellStatusBarProvider.unit.test.ts | 142 ++++++------------ 3 files changed, 78 insertions(+), 126 deletions(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index 8c305723c6..67602ba26d 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -17,7 +17,7 @@ "${workspaceFolder}/dist/**/*", "!${workspaceFolder}/**/node_modules**/*" ], - "preLaunchTask": "Build", + // "preLaunchTask": "Build", "skipFiles": [ "/**" ], diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index 75a9f5224d..d922a78ca4 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -75,8 +75,36 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv this.disposables.push(this._onDidChangeCellStatusBarItems); } + public dispose(): void { + this.disposables.forEach((disposable) => disposable.dispose()); + } + + public provideCellStatusBarItems( + cell: NotebookCell, + token: CancellationToken + ): NotebookCellStatusBarItem[] | undefined { + if (token.isCancellationRequested) { + return undefined; + } + + if (!isAgentCell(cell)) { + return undefined; + } + + const metadata = cell.metadata as Record | undefined; + const model = this.getModel(metadata); + + const items = [this.createAgentIndicatorItem(), this.createModelPickerItem(cell, model)]; + + if (this.getCellsToClear(cell).length > 0) { + items.push(this.createClearEphemeralItem(cell)); + } + + return items; + } + /** Deletes the ephemeral cells this agent block generated, after a modal confirmation. */ - public async clearEphemeralBlocks(cell: NotebookCell): Promise { + private async clearEphemeralBlocks(cell: NotebookCell): Promise { if (!isAgentCell(cell)) { return; } @@ -110,34 +138,6 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv } } - public dispose(): void { - this.disposables.forEach((disposable) => disposable.dispose()); - } - - public provideCellStatusBarItems( - cell: NotebookCell, - token: CancellationToken - ): NotebookCellStatusBarItem[] | undefined { - if (token.isCancellationRequested) { - return undefined; - } - - if (!isAgentCell(cell)) { - return undefined; - } - - const metadata = cell.metadata as Record | undefined; - const model = this.getModel(metadata); - - const items = [this.createAgentIndicatorItem(), this.createModelPickerItem(cell, model)]; - - if (this.getCellsToClear(cell).length > 0) { - items.push(this.createClearEphemeralItem(cell)); - } - - return items; - } - private createAgentIndicatorItem(): NotebookCellStatusBarItem { return { text: `$(hubot) ${l10n.t('Agent Block')}`, @@ -197,7 +197,7 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv return AGENT_MODEL_AUTO; } - public async switchModel(cell: NotebookCell): Promise { + private async switchModel(cell: NotebookCell): Promise { if (!isAgentCell(cell)) { return; } diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts index 565261e2d9..7adb119af1 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -4,7 +4,7 @@ import { anything, verify, when } from 'ts-mockito'; import { CancellationToken, NotebookCell, NotebookEdit, WorkspaceEdit } from 'vscode'; import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; -import { AGENT_MODEL_METADATA_KEY, AgentCellStatusBarProvider } from './agentCellStatusBarProvider'; +import { AgentCellStatusBarProvider } from './agentCellStatusBarProvider'; import { createMockCell, createMockNotebookWithCells } from './deepnoteTestHelpers'; suite('AgentCellStatusBarProvider', () => { @@ -64,39 +64,19 @@ suite('AgentCellStatusBarProvider', () => { expect(items).to.have.lengthOf(2); }); - 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 any cell that is not an agent block', () => { + const nonAgentCells: Record> = { + 'code cell': { __deepnotePocket: { type: 'code' } }, + 'sql cell': { __deepnotePocket: { type: 'sql' } }, + 'markdown cell': { __deepnotePocket: { type: 'markdown' } }, + 'cell without a pocket': {} + }; - test('Should return undefined for cell without pocket', () => { - const cell = createMockCell({ metadata: {} }); - const items = provider.provideCellStatusBarItems(cell, mockToken); + for (const [description, metadata] of Object.entries(nonAgentCells)) { + const items = provider.provideCellStatusBarItems(createMockCell({ metadata }), 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; + expect(items, description).to.be.undefined; + } }); test('Should return undefined when cancellation is requested', () => { @@ -163,8 +143,11 @@ suite('AgentCellStatusBarProvider', () => { }); }); + // Driven through the registered command rather than switchModel directly, so each case also + // covers the wiring the status bar item actually goes through. suite('Model Switching', () => { let capturedEdit: { index: number; metadata: Record } | undefined; + let invokeCommand: (cell?: NotebookCell) => Promise; setup(() => { resetVSCodeMocks(); @@ -177,6 +160,9 @@ suite('AgentCellStatusBarProvider', () => { return {} as NotebookEdit; }); + + activateCapturingCommands(); + invokeCommand = handlerFor('deepnote.switchAgentModel'); }); teardown(() => { @@ -207,7 +193,7 @@ suite('AgentCellStatusBarProvider', () => { pick('gpt-5.6-terra'); when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(true)); - await provider.switchModel(agentCell()); + await invokeCommand(agentCell()); verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).once(); expect(capturedEdit!.index).to.equal(2); @@ -223,7 +209,7 @@ suite('AgentCellStatusBarProvider', () => { // document on a no-op selection. pick('gpt-5.6-sol'); - await provider.switchModel(agentCell()); + await invokeCommand(agentCell()); verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); }); @@ -233,7 +219,7 @@ suite('AgentCellStatusBarProvider', () => { // user presses Escape. pick(undefined); - await provider.switchModel(agentCell()); + await invokeCommand(agentCell()); verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); }); @@ -248,7 +234,7 @@ suite('AgentCellStatusBarProvider', () => { statusBarRefreshed = true; }); - await provider.switchModel(agentCell()); + await invokeCommand(agentCell()); verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); expect(statusBarRefreshed, 'a rejected edit must not refresh the status bar').to.be.false; @@ -258,44 +244,32 @@ suite('AgentCellStatusBarProvider', () => { // Catches: dropping the isAgentCell guard, which would offer the model picker on any cell. pick('gpt-5.6-luna'); - await provider.switchModel(createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } })); + await invokeCommand(createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } })); verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).never(); verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); }); - suite('Command handler', () => { - let invokeCommand: (cell?: NotebookCell) => Promise; - - setup(() => { - activateCapturingCommands(); - invokeCommand = handlerFor('deepnote.switchAgentModel'); - }); - - test('Should switch the model of the cell the command is given', async () => { - pick('gpt-5.6-terra'); - when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(true)); - - await invokeCommand(agentCell()); - - expect(capturedEdit!.metadata[AGENT_MODEL_METADATA_KEY]).to.equal('gpt-5.6-terra'); - }); - - test('Should reject when invoked without a cell', async () => { - // Catches: falling back to the selected cell, which rewrites a block the user never clicked. - pick('gpt-5.6-terra'); + test('Should reject when invoked without a cell', async () => { + // Catches: falling back to the selected cell, which rewrites a block the user never clicked. + pick('gpt-5.6-terra'); - await assert.isRejected(invokeCommand(undefined), /requires the cell it was invoked from/); + await assert.isRejected(invokeCommand(undefined), /requires the cell it was invoked from/); - verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).never(); - verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); - }); + verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); }); }); + // Driven through the registered command rather than clearEphemeralBlocks directly, so each case + // also covers the wiring the status bar item actually goes through. suite('Clearing ephemeral blocks', () => { + let invokeCommand: (cell?: NotebookCell) => Promise; + setup(() => { resetVSCodeMocks(); + activateCapturingCommands(); + invokeCommand = handlerFor('deepnote.clearEphemeralBlocks'); }); teardown(() => { @@ -354,7 +328,7 @@ suite('AgentCellStatusBarProvider', () => { applyDeletionsTo(cells); confirmWith('Clear'); - await provider.clearEphemeralBlocks(cells[0]); + await invokeCommand(cells[0]); expect(cells.map((cell) => cell.document.getText())).to.deep.equal([ 'agent A', @@ -374,7 +348,7 @@ suite('AgentCellStatusBarProvider', () => { applyDeletionsTo(cells); confirmWith('Clear'); - await provider.clearEphemeralBlocks(cells[0]); + await invokeCommand(cells[0]); verify( mockedVSCodeNamespaces.window.showWarningMessage( @@ -394,7 +368,7 @@ suite('AgentCellStatusBarProvider', () => { applyDeletionsTo(cells); confirmWith(undefined); - await provider.clearEphemeralBlocks(cells[0]); + await invokeCommand(cells[0]); verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); expect(cells).to.have.lengthOf(2); @@ -409,7 +383,7 @@ suite('AgentCellStatusBarProvider', () => { applyDeletionsTo(cells); confirmWith('Clear'); - await provider.clearEphemeralBlocks(cells[0]); + await invokeCommand(cells[0]); verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); @@ -425,7 +399,7 @@ suite('AgentCellStatusBarProvider', () => { applyDeletionsTo(cells); confirmWith('Clear'); - await provider.clearEphemeralBlocks(cells[0]); + await invokeCommand(cells[0]); verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); @@ -440,7 +414,7 @@ suite('AgentCellStatusBarProvider', () => { confirmWith('Clear'); when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(false)); - await provider.clearEphemeralBlocks(cells[0]); + await invokeCommand(cells[0]); verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); }); @@ -476,36 +450,14 @@ suite('AgentCellStatusBarProvider', () => { }); }); - suite('Command handler', () => { - let invokeCommand: (cell?: NotebookCell) => Promise; - - setup(() => { - activateCapturingCommands(); - invokeCommand = handlerFor('deepnote.clearEphemeralBlocks'); - }); - - test('Should clear the blocks of the cell the command is given', async () => { - const { cells } = createMockNotebookWithCells([ - agentBlock('agent A', 'agent-block-1'), - ephemeralCell('eph A1', 'agent-block-1') - ]); - applyDeletionsTo(cells); - confirmWith('Clear'); - - await invokeCommand(cells[0]); - - expect(cells.map((cell) => cell.document.getText())).to.deep.equal(['agent A']); - }); - - test('Should reject when invoked without a cell', async () => { - // Catches: falling back to the selected cell, which clears a run the user never clicked. - confirmWith('Clear'); + test('Should reject when invoked without a cell', async () => { + // Catches: falling back to the selected cell, which clears a run the user never clicked. + confirmWith('Clear'); - await assert.isRejected(invokeCommand(undefined), /requires the cell it was invoked from/); + await assert.isRejected(invokeCommand(undefined), /requires the cell it was invoked from/); - verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); - verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); - }); + verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); }); }); }); From 7eca40a30cff0d9bb6b1c35b54ae3d3d1f30a5c9 Mon Sep 17 00:00:00 2001 From: tomas Date: Sun, 9 Aug 2026 19:20:00 +0000 Subject: [PATCH 62/80] chore: restore launch.json, committed by mistake The commented-out preLaunchTask is a local dev tweak that slipped into 0aaeda942 via `git add -u`; it does not belong in the branch. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- .vscode/launch.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index 67602ba26d..8c305723c6 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -17,7 +17,7 @@ "${workspaceFolder}/dist/**/*", "!${workspaceFolder}/**/node_modules**/*" ], - // "preLaunchTask": "Build", + "preLaunchTask": "Build", "skipFiles": [ "/**" ], From 14c0b1c4378bd086dfb8e83999438b5420fc3ad4 Mon Sep 17 00:00:00 2001 From: tomas Date: Sun, 9 Aug 2026 19:35:49 +0000 Subject: [PATCH 63/80] test(agent-block): type createMockEditor's revealRange as the stub it is The helper declared its return as NotebookEditor, so revealRange carried the VS Code signature and every assertion on it had to cast to SinonStub. Widening the return type to the intersection removes all twelve casts and makes the stub-ness checked: swapping revealRange for a plain function is now a compile error, where the cast would have accepted it and failed at runtime. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- ...epnoteNotebookCommandListener.unit.test.ts | 32 +++++++++---------- 1 file changed, 15 insertions(+), 17 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts index 9048192d22..78c5d29427 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts @@ -380,7 +380,8 @@ suite('DeepnoteNotebookCommandListener', () => { cellDataArray: NotebookCellData[], selection?: NotebookRange ): { - editor: NotebookEditor; + // revealRange is narrowed to the stub it actually is, so assertions on it need no cast. + editor: NotebookEditor & { revealRange: sinon.SinonStub }; document: NotebookDocument; } { const { notebook: document } = createMockNotebookWithCells( @@ -395,7 +396,7 @@ suite('DeepnoteNotebookCommandListener', () => { const editorSelection = selection != null ? selection : new NotebookRange(0, cellDataArray.length > 0 ? 1 : 0); - const editor: NotebookEditor = { + const editor: NotebookEditor & { revealRange: sinon.SinonStub } = { notebook: document, selection: editorSelection, selections: [editorSelection], @@ -671,11 +672,8 @@ suite('DeepnoteNotebookCommandListener', () => { }); // Verify reveal and selection were set - assert.isTrue( - (editor.revealRange as sinon.SinonStub).calledOnce, - 'Should reveal the new cell range' - ); - const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + assert.isTrue(editor.revealRange.calledOnce, 'Should reveal the new cell range'); + const revealCall = editor.revealRange.firstCall; assert.equal(revealCall.args[0].start, expectedInsertIndex, 'Should reveal correct range start'); assert.equal(revealCall.args[0].end, expectedInsertIndex + 1, 'Should reveal correct range end'); }); @@ -762,8 +760,8 @@ suite('DeepnoteNotebookCommandListener', () => { ); // Verify reveal and selection were set - assert.isTrue((editor.revealRange as sinon.SinonStub).calledOnce, 'Should reveal the new cell range'); - const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + assert.isTrue(editor.revealRange.calledOnce, 'Should reveal the new cell range'); + const revealCall = editor.revealRange.firstCall; assert.equal(revealCall.args[0].start, 0, 'Should reveal correct range start'); assert.equal(revealCall.args[0].end, 1, 'Should reveal correct range end'); assert.equal(revealCall.args[1], 0, 'Should use NotebookEditorRevealType.Default (value 0)'); @@ -882,8 +880,8 @@ suite('DeepnoteNotebookCommandListener', () => { assert.equal(newCell.value, '', 'Should have empty content'); assert.equal(newCell.metadata.__deepnotePocket.type, 'agent', 'Should have agent pocket type'); - assert.isTrue((editor.revealRange as sinon.SinonStub).calledOnce, 'Should reveal the new cell range'); - const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + assert.isTrue(editor.revealRange.calledOnce, 'Should reveal the new cell range'); + const revealCall = editor.revealRange.firstCall; assert.equal(revealCall.args[0].start, 0, 'Should reveal correct range start'); assert.equal(revealCall.args[0].end, 1, 'Should reveal correct range end'); }); @@ -898,7 +896,7 @@ suite('DeepnoteNotebookCommandListener', () => { insertedCell(getCapturedNotebookEdits); assert.equal(chainStub.firstCall.args[0], document, 'Should edit the active document'); - const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + const revealCall = editor.revealRange.firstCall; assert.equal(revealCall.args[0].start, 2, 'Should insert below the selection'); assert.equal(revealCall.args[0].end, 3, 'Should select only the new cell'); }); @@ -953,7 +951,7 @@ suite('DeepnoteNotebookCommandListener', () => { assert.isFalse(chainStub.called, 'Must not edit a notebook that already has an agent block'); verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); - assert.isFalse((editor.revealRange as sinon.SinonStub).called, 'Must not reveal anything'); + assert.isFalse(editor.revealRange.called, 'Must not reveal anything'); }); test('should still add the block when other cells carry no agent pocket', async () => { @@ -1029,8 +1027,8 @@ suite('DeepnoteNotebookCommandListener', () => { assert.equal(newCell.metadata.__deepnotePocket.type, 'big-number', 'Should have big-number type'); // Verify reveal and selection were set - assert.isTrue((editor.revealRange as sinon.SinonStub).calledOnce, 'Should reveal the new cell range'); - const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + assert.isTrue(editor.revealRange.calledOnce, 'Should reveal the new cell range'); + const revealCall = editor.revealRange.firstCall; assert.equal(revealCall.args[0].start, 0, 'Should reveal correct range start'); assert.equal(revealCall.args[0].end, 1, 'Should reveal correct range end'); assert.equal(revealCall.args[1], 0, 'Should use NotebookEditorRevealType.Default (value 0)'); @@ -1165,8 +1163,8 @@ suite('DeepnoteNotebookCommandListener', () => { assert.equal(newCell.metadata.__deepnotePocket.type, 'visualization', 'Should have visualization type'); // Verify reveal and selection were set - assert.isTrue((editor.revealRange as sinon.SinonStub).calledOnce, 'Should reveal the new cell range'); - const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + assert.isTrue(editor.revealRange.calledOnce, 'Should reveal the new cell range'); + const revealCall = editor.revealRange.firstCall; assert.equal(revealCall.args[0].start, 0, 'Should reveal correct range start'); assert.equal(revealCall.args[0].end, 1, 'Should reveal correct range end'); assert.equal(revealCall.args[1], 0, 'Should use NotebookEditorRevealType.Default (value 0)'); From 35cbe71a84badb108add7a565a591712853df82a Mon Sep 17 00:00:00 2001 From: tomas Date: Sun, 9 Aug 2026 19:40:15 +0000 Subject: [PATCH 64/80] Fix typo --- src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index ff95daad98..91e8a795e4 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -1197,7 +1197,7 @@ project: ] } }; - // Force casting wihtout metadata + // Force casting without metadata deepnoteFile.project.notebooks[0].blocks = [ { id: 'block-1', type: 'code', sortingKey: 'a0' } as DeepnoteBlock, { id: 'block-2', type: 'code', sortingKey: 'a1' } as DeepnoteBlock From 676c0976d1489cf3ff0bc6bbce0fea649841f31b Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 10 Aug 2026 11:26:04 +0000 Subject: [PATCH 65/80] fix(agent-block): close the race on the one-agent-block rule MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The existence check ran before the queued notebook update, so two invocations could both pass it while neither edit had applied — a double-click on the toolbar button was enough to insert two agent blocks. Repeat the check inside the serialized callback, where a concurrent edit has already landed, and compute insertIndex there too so it cannot go stale. The outer check stays as a fast path that keeps an obviously-redundant edit off the queue. The distinct-ids test drove a second insertion into one notebook, which production refuses; it only passed because the mocked update never mutates the notebook. Give it a second notebook instead. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- .../deepnoteNotebookCommandListener.ts | 24 ++++++++- ...epnoteNotebookCommandListener.unit.test.ts | 51 ++++++++++++++++++- 2 files changed, 71 insertions(+), 4 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts index 7d99d93f3b..f3726f89df 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts @@ -247,17 +247,30 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation } const document = editor.notebook; const selection = editor.selection; + const agentBlockExistsMessage = l10n.t('This notebook already contains an agent block.'); if (document.getCells().some(isAgentCell)) { - void window.showInformationMessage(l10n.t('This notebook already contains an agent block.')); + void window.showInformationMessage(agentBlockExistsMessage); return; } - const insertIndex = selection ? selection.end : document.cellCount; const blockId = generateBlockId(); + let alreadyHasAgentBlock = false; + let insertIndex = 0; + const result = await notebookUpdaterUtils.chainWithPendingUpdates(document, (edit) => { + // Repeated inside the serialized callback: a concurrent invocation passes the check above + // while its edit is still queued, and only here has that edit already applied. + if (document.getCells().some(isAgentCell)) { + alreadyHasAgentBlock = true; + + return; + } + + insertIndex = selection ? selection.end : document.cellCount; + const newCell = new NotebookCellData(NotebookCellKind.Code, '', 'plaintext'); newCell.metadata = { __deepnotePocket: { @@ -270,6 +283,13 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation const nbEdit = NotebookEdit.insertCells(insertIndex, [newCell]); edit.set(document.uri, [nbEdit]); }); + + if (alreadyHasAgentBlock) { + void window.showInformationMessage(agentBlockExistsMessage); + + return; + } + if (result !== true) { throw new Error(l10n.t('Failed to insert agent block')); } diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts index 78c5d29427..ac7127b0d2 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts @@ -914,15 +914,18 @@ suite('DeepnoteNotebookCommandListener', () => { assert.equal(metadata.__deepnoteBlockId, metadata.id, 'Backup id key must match id'); }); - test('should give consecutive agent blocks distinct ids', async () => { + test('should give each notebook its own agent block id', async () => { // Catches: a hoisted/constant id, which would make two agent blocks fight over the - // same generated cells. + // same generated cells. Two notebooks, because one notebook only ever gets one block. const { editor } = createMockEditor([], undefined); const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); await commandListener.addAgentBlock(); const first = insertedCell(getCapturedNotebookEdits).metadata.id; + const { editor: otherEditor } = createMockEditor([], undefined); + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(otherEditor); + await commandListener.addAgentBlock(); const second = insertedCell(getCapturedNotebookEdits).metadata.id; @@ -954,6 +957,50 @@ suite('DeepnoteNotebookCommandListener', () => { assert.isFalse(editor.revealRange.called, 'Must not reveal anything'); }); + test('should insert only one agent block when two invocations race', async () => { + // Catches: an existence check that runs before the queued update — both invocations + // pass it while neither edit has applied yet. + const { editor, document } = createMockEditor([], undefined); + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(editor); + when(mockedVSCodeNamespaces.commands.executeCommand(anything())).thenResolve(undefined as any); + + const cells = document.getCells(); + const insertedCells: NotebookCellData[] = []; + let pending: Promise = Promise.resolve(); + + // Mirrors chainWithPendingUpdates: a callback runs only once the previous edit applied. + sandbox + .stub(notebookUpdater.notebookUpdaterUtils, 'chainWithPendingUpdates') + .callsFake((_doc: NotebookDocument, callback: (edit: WorkspaceEdit) => void) => { + const applied = pending.then(() => { + const edit = new WorkspaceEdit(); + sandbox.stub(edit, 'set').callsFake((_uri, edits) => { + for (const newCell of (edits[0] as any).newCells as NotebookCellData[]) { + insertedCells.push(newCell); + cells.push( + createMockCell({ + metadata: newCell.metadata, + index: cells.length, + notebook: document + }) + ); + } + }); + callback(edit); + + return true; + }); + pending = applied; + + return applied; + }); + + await Promise.all([commandListener.addAgentBlock(), commandListener.addAgentBlock()]); + + assert.equal(insertedCells.length, 1, 'Should insert exactly one agent block'); + assert.equal(cells.length, 1, 'Notebook should end up with a single cell'); + }); + test('should still add the block when other cells carry no agent pocket', async () => { // Catches: a guard that trips on any cell, blocking the first agent block outright. const { editor } = createMockEditor([ From e531c5047733cf364c762ffbdfc7945bb24fdf86 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 12 Aug 2026 15:03:34 +0000 Subject: [PATCH 66/80] feat(analytics): count agent blocks and mark ephemeral cells addAgentBlock was the only add-block command that never called trackAddBlock, so agent-block adoption reported zero: the auto-tracker skips pocket-typed cells, and nothing else observed the insert. Agent scratch cells had the inverse problem. They were invisible to add_block for the same reason, while the agent running them through notebook.cell.execute re-enters the kernel path and emitted execute_cell as though a user had pressed Run. isEphemeral separates the two. Required rather than optional so no call site can omit it -- queries that mean "a human did this" need `isEphemeral != true`, since rows written before this change carry no such field. Co-Authored-By: Claude Opus 4.6 Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- .../deepnoteCellExecutionAnalytics.ts | 27 ++++++++-- ...eepnoteCellExecutionAnalytics.unit.test.ts | 51 ++++++++++++++++++- .../deepnoteNotebookCommandListener.ts | 6 ++- ...epnoteNotebookCommandListener.unit.test.ts | 16 +++++- src/platform/analytics/types.ts | 6 ++- 5 files changed, 96 insertions(+), 10 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.ts b/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.ts index 3ab4a7a2c8..0bb6386189 100644 --- a/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.ts +++ b/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.ts @@ -2,7 +2,7 @@ import { inject, injectable } from 'inversify'; import { Disposable, NotebookCellKind, workspace } from 'vscode'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; -import { ITelemetryService } from '../../platform/analytics/types'; +import { ITelemetryService, TelemetryEventProperties } from '../../platform/analytics/types'; import { IDisposableRegistry } from '../../platform/common/types'; import { isDeepnoteNotebook } from '../../platform/common/utils'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; @@ -11,10 +11,12 @@ import { toTelemetryIntegrationType } from '../../platform/notebooks/deepnote/integrationTypes'; import { IDeepnoteNotebookManager } from '../types'; +import { isEphemeralCell } from './dataConversionUtils'; /** * Tracks cell executions, plus the plain code/markdown insertions from VS Code's built-in - * "+ Code" / "+ Markdown" controls that never reach an extension command. + * "+ Code" / "+ Markdown" controls that never reach an extension command, plus the scratch + * cells the agent writes and runs on the user's behalf. */ @injectable() export class DeepnoteCellExecutionAnalytics implements IExtensionSyncActivationService { @@ -39,6 +41,18 @@ export class DeepnoteCellExecutionAnalytics implements IExtensionSyncActivationS } for (const cell of change.addedCells) { + const blockType = cell.kind === NotebookCellKind.Code ? 'code' : 'markdown'; + + // Agent scratch cells stamp a pocket like any typed block, but no command + // inserts them, so this is the only place they can be counted. + if (isEphemeralCell(cell)) { + this.analytics.trackEvent({ + eventName: 'add_block', + properties: { blockType, isEphemeral: true } + }); + continue; + } + // Typed Deepnote blocks stamp a pocket on insert and are already counted by // DeepnoteNotebookCommandListener. if (cell.metadata?.__deepnotePocket?.type) { @@ -47,7 +61,7 @@ export class DeepnoteCellExecutionAnalytics implements IExtensionSyncActivationS this.analytics.trackEvent({ eventName: 'add_block', - properties: { blockType: cell.kind === NotebookCellKind.Code ? 'code' : 'markdown' } + properties: { blockType, isEphemeral: false } }); } } @@ -67,7 +81,12 @@ export class DeepnoteCellExecutionAnalytics implements IExtensionSyncActivationS const languageId = e.cell.document.languageId; const cellType = languageId === 'sql' ? 'sql' : languageId === 'markdown' ? 'markdown' : 'code'; - const properties: { cellType: 'sql' | 'markdown' | 'code'; integrationType?: string } = { cellType }; + // The agent runs its generated cells through `notebook.cell.execute`, which re-enters the + // kernel path; unmarked they are indistinguishable here from a user pressing Run. + const properties: TelemetryEventProperties['execute_cell'] = { + cellType, + isEphemeral: isEphemeralCell(e.cell) + }; if (cellType === 'sql') { // The status-bar switch updates only this key, so the __deepnotePocket copy can go stale. diff --git a/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.unit.test.ts b/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.unit.test.ts index dc72115043..bdb08f633a 100644 --- a/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.unit.test.ts @@ -31,6 +31,17 @@ suite('DeepnoteCellExecutionAnalytics', () => { } as unknown as NotebookCell; } + /** Agent scratch cells go through the converter, so they carry a pocket like any typed block. */ + function ephemeralCell(kind: NotebookCellKind): NotebookCell { + return { + kind, + metadata: { + __deepnotePocket: { type: kind === NotebookCellKind.Code ? 'code' : 'markdown' }, + is_ephemeral: true + } + } as unknown as NotebookCell; + } + function executingCell(languageId: string, sqlIntegrationId?: string, notebookType = 'deepnote'): NotebookCell { return { document: { languageId }, @@ -106,7 +117,9 @@ suite('DeepnoteCellExecutionAnalytics', () => { row.expected.forEach((blockType) => verify( - telemetry.trackEvent(deepEqual({ eventName: 'add_block', properties: { blockType } })) + telemetry.trackEvent( + deepEqual({ eventName: 'add_block', properties: { blockType, isEphemeral: false } }) + ) ).once() ); verify(telemetry.trackEvent(anything())).times(row.expected.length); @@ -118,6 +131,22 @@ suite('DeepnoteCellExecutionAnalytics', () => { verify(telemetry.trackEvent(anything())).never(); }); + + test('counts agent scratch cells, which no command reports, as ephemeral', () => { + fireContentChange([ephemeralCell(NotebookCellKind.Code), ephemeralCell(NotebookCellKind.Markup)], []); + + verify( + telemetry.trackEvent( + deepEqual({ eventName: 'add_block', properties: { blockType: 'code', isEphemeral: true } }) + ) + ).once(); + verify( + telemetry.trackEvent( + deepEqual({ eventName: 'add_block', properties: { blockType: 'markdown', isEphemeral: true } }) + ) + ).once(); + verify(telemetry.trackEvent(anything())).twice(); + }); }); suite('execute_cell', () => { @@ -171,12 +200,30 @@ suite('DeepnoteCellExecutionAnalytics', () => { ); verify( - telemetry.trackEvent(deepEqual({ eventName: 'execute_cell', properties: { ...row.expected } })) + telemetry.trackEvent( + deepEqual({ + eventName: 'execute_cell', + properties: { ...row.expected, isEphemeral: false } + }) + ) ).once(); verify(telemetry.trackEvent(anything())).once(); }); }); + test('an agent-generated cell reports isEphemeral true', () => { + const cell = { ...executingCell('python'), metadata: { is_ephemeral: true } } as unknown as NotebookCell; + + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Executing); + + verify( + telemetry.trackEvent( + deepEqual({ eventName: 'execute_cell', properties: { cellType: 'code', isEphemeral: true } }) + ) + ).once(); + verify(telemetry.trackEvent(anything())).once(); + }); + test('ignores Pending and Idle transitions, and non-Deepnote notebooks', () => { const cell = executingCell('python'); diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts index f3726f89df..c01c73d0c5 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts @@ -294,6 +294,8 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation throw new Error(l10n.t('Failed to insert agent block')); } + this.trackAddBlock('agent'); + const notebookRange = new NotebookRange(insertIndex, insertIndex + 1); editor.revealRange(notebookRange, NotebookEditorRevealType.Default); editor.selection = notebookRange; @@ -661,6 +663,8 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation } private trackAddBlock(blockType: string): void { - this.analytics.trackEvent({ eventName: 'add_block', properties: { blockType } }); + // Commands only ever insert blocks the user asked for; agent scratch cells are counted + // in DeepnoteCellExecutionAnalytics, which is the only observer that sees them. + this.analytics.trackEvent({ eventName: 'add_block', properties: { blockType, isEphemeral: false } }); } } diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts index ac7127b0d2..f75e9ece90 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts @@ -1,6 +1,6 @@ import { assert } from 'chai'; import * as sinon from 'sinon'; -import { when, reset, anything, mock, instance, verify } from 'ts-mockito'; +import { when, reset, anything, deepEqual, mock, instance, verify } from 'ts-mockito'; import { NotebookCell, NotebookDocument, @@ -955,6 +955,20 @@ suite('DeepnoteNotebookCommandListener', () => { assert.isFalse(chainStub.called, 'Must not edit a notebook that already has an agent block'); verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); assert.isFalse(editor.revealRange.called, 'Must not reveal anything'); + verify(mockTelemetryService.trackEvent(anything())).never(); + }); + + test('should report the added agent block to analytics', async () => { + const { editor } = createMockEditor([], undefined); + mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + verify( + mockTelemetryService.trackEvent( + deepEqual({ eventName: 'add_block', properties: { blockType: 'agent', isEphemeral: false } }) + ) + ).once(); }); test('should insert only one agent block when two invocations race', async () => { diff --git a/src/platform/analytics/types.ts b/src/platform/analytics/types.ts index 4ce4f15a39..f63bcb7665 100644 --- a/src/platform/analytics/types.ts +++ b/src/platform/analytics/types.ts @@ -33,7 +33,8 @@ export type CommandOutcome = 'completed' | 'cancelled' | 'failed'; /** Caller-supplied properties per event; `undefined` means none beyond the common properties the service attaches. */ export interface TelemetryEventProperties { - add_block: { blockType: string }; + /** `isEphemeral` marks agent scratch blocks, which no add-block command reports. */ + add_block: { blockType: string; isEphemeral: boolean }; authenticate_integration: { integrationType: string; outcome: CommandOutcome }; configure_integration: { integrationType: string }; copy_notebook_details: undefined; @@ -44,7 +45,8 @@ export interface TelemetryEventProperties { delete_integration: { integrationType: string }; delete_notebook: { outcome: CommandOutcome }; duplicate_notebook: { outcome: CommandOutcome }; - execute_cell: { cellType: 'sql' | 'markdown' | 'code'; integrationType?: string }; + /** `isEphemeral` is true for agent-generated cells, which the agent runs itself via `notebook.cell.execute`. */ + execute_cell: { cellType: 'sql' | 'markdown' | 'code'; isEphemeral: boolean; integrationType?: string }; execute_notebook: undefined; export_notebook: { outcome: CommandOutcome; format?: string }; import_notebook: { outcome: CommandOutcome; source: 'deepnote' | 'jupyter' }; From a4ca6feedc89810f14b87739790a93a2fa606259 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 13 Aug 2026 06:19:40 +0000 Subject: [PATCH 67/80] fix(deepnote): keep every streamed output item when saving transformOutputsForDeepnote took the first stdout or stderr item of an output and dropped the rest. Agent runs append every streamed delta as a new item on one output, so saving kept only "[Agent] Planning next steps..." and lost the whole transcript -- 82% of it in the case that prompted this. Ordinary Jupyter cells whose stdout arrives in several chunks were truncated the same way; this is not agent-specific. Note the agent's context serializer runs the same converter, so a later run now sees an earlier agent cell's full output. That is correct, and it grows the prompt in a way the truncation was hiding. Co-Authored-By: Claude Opus 4.6 Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- .../deepnote/deepnoteDataConverter.ts | 27 +++--- .../deepnoteDataConverter.unit.test.ts | 96 +++++++++++++++++++ 2 files changed, 111 insertions(+), 12 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteDataConverter.ts b/src/notebooks/deepnote/deepnoteDataConverter.ts index 09b06cc2be..42e1eb3eaa 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.ts @@ -496,19 +496,22 @@ export class DeepnoteDataConverter { } } - // Check if this is a stream output - const stdoutItem = output.items.find((item) => item.mime === 'application/vnd.code.notebook.stdout'); - const stderrItem = output.items.find((item) => item.mime === 'application/vnd.code.notebook.stderr'); - - if (stdoutItem || stderrItem) { - const item = stdoutItem || stderrItem; - const text = new TextDecoder().decode(item!.data); - - return { - name: stderrItem ? 'stderr' : 'stdout', + // Streamed deltas are appended as new items on one output, so every item must be joined. + const streamItems = output.items.filter( + (item) => + item.mime === 'application/vnd.code.notebook.stdout' || + item.mime === 'application/vnd.code.notebook.stderr' + ); + + if (streamItems.length > 0) { + const decoder = new TextDecoder(); + const streamOutput: DeepnoteOutput = { + name: streamItems[0].mime === 'application/vnd.code.notebook.stderr' ? 'stderr' : 'stdout', output_type: 'stream', - text - } as DeepnoteOutput; + text: streamItems.map((item) => decoder.decode(item.data)).join('') + }; + + return streamOutput; } // Rich output (execute_result or display_data) diff --git a/src/notebooks/deepnote/deepnoteDataConverter.unit.test.ts b/src/notebooks/deepnote/deepnoteDataConverter.unit.test.ts index c1ac9b7b94..b7198127b0 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.unit.test.ts @@ -282,6 +282,102 @@ suite('DeepnoteDataConverter', () => { assert.strictEqual(new TextDecoder().decode(outputs[0].items[0].data), 'Hello world\n'); }); + test('joins every streamed stdout item into one stream output', () => { + const cells: NotebookCellData[] = [ + { + kind: NotebookCellKind.Code, + value: 'summarize the data', + languageId: 'plaintext', + metadata: { + __deepnotePocket: { + type: 'agent', + sortingKey: 'a0' + }, + id: 'agent-block-1' + }, + outputs: [ + new NotebookCellOutput([ + NotebookCellOutputItem.stdout('[Agent] Planning next steps...'), + NotebookCellOutputItem.stdout('\n\n[Agent] Tool called: add_code'), + NotebookCellOutputItem.stdout('\n\n[Agent] Text:\nDone.') + ]) + ] + } + ]; + + const blocks = converter.convertCellsToBlocks(cells); + + assert.deepStrictEqual((blocks[0] as ExecutableBlock).outputs, [ + { + name: 'stdout', + output_type: 'stream', + text: '[Agent] Planning next steps...\n\n[Agent] Tool called: add_code\n\n[Agent] Text:\nDone.' + } + ]); + }); + + test('converts a single stderr item to a stderr stream output', () => { + const cells: NotebookCellData[] = [ + { + kind: NotebookCellKind.Code, + value: 'raise ValueError()', + languageId: 'python', + metadata: { + __deepnotePocket: { + type: 'code', + sortingKey: 'a0' + }, + id: 'block-1' + }, + outputs: [new NotebookCellOutput([NotebookCellOutputItem.stderr('Agent execution failed: boom')])] + } + ]; + + const blocks = converter.convertCellsToBlocks(cells); + + assert.deepStrictEqual((blocks[0] as ExecutableBlock).outputs, [ + { + name: 'stderr', + output_type: 'stream', + text: 'Agent execution failed: boom' + } + ]); + }); + + test('keeps every item of a mixed stdout/stderr output and labels it by the first', () => { + const cells: NotebookCellData[] = [ + { + kind: NotebookCellKind.Code, + value: 'print("a")', + languageId: 'python', + metadata: { + __deepnotePocket: { + type: 'code', + sortingKey: 'a0' + }, + id: 'block-1' + }, + outputs: [ + new NotebookCellOutput([ + NotebookCellOutputItem.stdout('out-1\n'), + NotebookCellOutputItem.stderr('err-1\n'), + NotebookCellOutputItem.stdout('out-2\n') + ]) + ] + } + ]; + + const blocks = converter.convertCellsToBlocks(cells); + + assert.deepStrictEqual((blocks[0] as ExecutableBlock).outputs, [ + { + name: 'stdout', + output_type: 'stream', + text: 'out-1\nerr-1\nout-2\n' + } + ]); + }); + test('converts error output', () => { const deepnoteOutputs: DeepnoteOutput[] = [ { From 040e49852eadfcfac2cd778c3ab35b0dcfd74fc5 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 13 Aug 2026 06:19:49 +0000 Subject: [PATCH 68/80] fix(deepnote): only re-identify blocks that arrived without an id recoverBlockIdsFromOriginal matched on trimmed content alone -- not type, not cell kind -- and rewrote the id, sortingKey and blockGroup of any block whose id was absent from the stored project. Deleting an empty block and adding an empty agent block in the same save handed the agent the deleted block's identity. That matters now because addAgentBlock mints its id up front so each run can stamp its generated cells with a stable owner; the recovery silently voided it on the first save, leaving the main file and the snapshot disagreeing about which block the outputs belong to. Recovery still runs for cells VS Code stripped metadata from, which is what it was added for -- those have no id, so they stay candidates. Adding type to the match key would not work: a metadata-stripped SQL block arrives as 'code' and would stop matching its own original. Co-Authored-By: Claude Opus 4.6 Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- src/notebooks/deepnote/deepnoteSerializer.ts | 21 +- .../deepnote/deepnoteSerializer.unit.test.ts | 276 ++++++++++++++++++ 2 files changed, 294 insertions(+), 3 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 0365216810..8e57fa7596 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -7,7 +7,7 @@ import { workspace, type CancellationToken, type NotebookData, type NotebookSeri import { logger } from '../../platform/logging'; import { IDeepnoteNotebookManager } from '../types'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; -import { isEphemeralCell } from './dataConversionUtils'; +import { getBlockId, isEphemeralCell } from './dataConversionUtils'; import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; import { SnapshotService } from './snapshots/snapshotService'; import { computeHash } from '../../platform/common/crypto'; @@ -254,9 +254,13 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { logger.debug(`SerializeNotebook: Converted to ${blocks.length} blocks`); + // An id the cell still carries is its real identity, so only cells that arrived without one may be + // re-identified. convertCellsToBlocks maps 1:1 in order, so index i lines up with nonEphemeralCells[i]. + const recoverableBlocks = new Set(blocks.filter((_, index) => !getBlockId(nonEphemeralCells[index]))); + // Try to recover block IDs from original blocks when VS Code fails to preserve metadata // This uses content-based matching as a fallback when metadata.id is missing - this.recoverBlockIdsFromOriginal(blocks, notebook.blocks ?? []); + this.recoverBlockIdsFromOriginal(blocks, notebook.blocks ?? [], recoverableBlocks); // Log block IDs after conversion and recovery for (let i = 0; i < blocks.length; i++) { @@ -504,8 +508,13 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { * Uses content-based matching as a fallback strategy to recover id, sortingKey, and blockGroup. * @param blocks Blocks converted from cells (may have generated values if metadata was lost) * @param originalBlocks Original blocks from the stored project + * @param recoverableBlocks Blocks whose cell carried no id; only these may take an original's identity */ - private recoverBlockIdsFromOriginal(blocks: DeepnoteBlock[], originalBlocks: DeepnoteBlock[]): void { + private recoverBlockIdsFromOriginal( + blocks: DeepnoteBlock[], + originalBlocks: DeepnoteBlock[], + recoverableBlocks: Set + ): void { // Build a map of original blocks by content for quick lookup // Key: content (trimmed), Value: array of blocks with that content (in case of duplicates) const contentToOriginalBlocks = new Map(); @@ -533,6 +542,12 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { let recoveredCount = 0; for (const block of blocks) { + // A cell that carried an id already owns its identity - an agent block mints one when it is inserted - + // so it must never adopt the id of a block deleted in the same save. + if (!recoverableBlocks.has(block)) { + continue; + } + // Skip if this block already has an original ID if (claimedIds.has(block.id)) { continue; diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index 43e0be82cd..8a00f04a1a 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -728,6 +728,282 @@ project: 'Block ID should be newly generated when content differs' ); }); + + test('should keep a minted agent block ID when a deleted block had the same content', async () => { + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2023-01-01T00:00:00Z', + modifiedAt: '2023-01-02T00:00:00Z' + }, + project: { + id: 'project-agent-id', + name: 'Agent ID Test', + notebooks: [ + { + id: 'notebook-1', + name: 'Test Notebook', + blocks: [ + { + blockGroup: 'deleted-group', + id: 'deleted-block-id', + content: '', + sortingKey: 'a0', + metadata: {}, + type: 'code' + } + ], + executionMode: 'block', + isModule: false + } + ], + settings: {} + } + }; + + manager.storeOriginalProject('project-agent-id', 'notebook-1', projectData); + + // The empty code block was deleted and an empty agent block added in the same save + const notebookData = { + cells: [ + { + kind: 2, + value: '', + languageId: 'plaintext', + metadata: { + id: 'minted-agent-id', + __deepnoteBlockId: 'minted-agent-id', + __deepnotePocket: { type: 'agent' } + } + } + ], + metadata: { + deepnoteProjectId: 'project-agent-id', + deepnoteNotebookId: 'notebook-1' + } + }; + + const result = await serializer.serializeNotebook(notebookData 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[0].id, 'minted-agent-id', 'Agent block should keep its minted ID'); + assert.notStrictEqual( + notebook!.blocks[0].blockGroup, + 'deleted-group', + 'Agent block should not inherit the deleted block blockGroup' + ); + }); + + test('should keep an ID the cell carried even when the deleted block has the same type', async () => { + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2023-01-01T00:00:00Z', + modifiedAt: '2023-01-02T00:00:00Z' + }, + project: { + id: 'project-same-type', + name: 'Same Type Test', + notebooks: [ + { + id: 'notebook-1', + name: 'Test Notebook', + blocks: [ + { + blockGroup: 'deleted-group', + id: 'deleted-code-id', + content: '', + sortingKey: 'a0', + metadata: {}, + type: 'code' + } + ], + executionMode: 'block', + isModule: false + } + ], + settings: {} + } + }; + + manager.storeOriginalProject('project-same-type', 'notebook-1', projectData); + + const notebookData = { + cells: [ + { + kind: 2, + value: '', + languageId: 'python', + metadata: { + id: 'minted-code-id', + __deepnoteBlockId: 'minted-code-id' + } + } + ], + metadata: { + deepnoteProjectId: 'project-same-type', + deepnoteNotebookId: 'notebook-1' + } + }; + + const result = await serializer.serializeNotebook(notebookData 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[0].id, 'minted-code-id', 'Block should keep the ID its cell carried'); + }); + + test('should not recover an ID that another cell still carries', async () => { + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2023-01-01T00:00:00Z', + modifiedAt: '2023-01-02T00:00:00Z' + }, + project: { + id: 'project-claimed-id', + name: 'Claimed ID Test', + notebooks: [ + { + id: 'notebook-1', + name: 'Test Notebook', + blocks: [ + { + blockGroup: 'group-kept', + id: 'kept-id', + content: '', + sortingKey: 'a0', + metadata: {}, + type: 'code' + }, + { + blockGroup: 'group-stripped', + id: 'stripped-id', + content: '', + sortingKey: 'a1', + metadata: {}, + type: 'code' + } + ], + executionMode: 'block', + isModule: false + } + ], + settings: {} + } + }; + + manager.storeOriginalProject('project-claimed-id', 'notebook-1', projectData); + + // Both blocks are empty, so content matching alone cannot tell them apart + const notebookData = { + cells: [ + { + kind: 2, + value: '', + languageId: 'python', + metadata: { + id: 'kept-id', + __deepnoteBlockId: 'kept-id', + __deepnotePocket: { type: 'code', sortingKey: 'a0', blockGroup: 'group-kept' } + } + }, + { + kind: 2, + value: '', + languageId: 'python', + metadata: {} + } + ], + metadata: { + deepnoteProjectId: 'project-claimed-id', + deepnoteNotebookId: 'notebook-1' + } + }; + + const result = await serializer.serializeNotebook(notebookData 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[0].id, 'kept-id', 'Cell that carried an ID should keep it'); + assert.strictEqual(notebook!.blocks[1].id, 'stripped-id', 'Stripped cell should take the remaining ID'); + }); + + test('should recover IDs by cell position after ephemeral cells are dropped', 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-offset', + name: 'Ephemeral Offset Test', + notebooks: [ + { + id: 'notebook-1', + name: 'Test Notebook', + blocks: [ + { + blockGroup: 'group-1', + id: 'real-block-id', + content: 'print("kept")', + sortingKey: 'a0', + metadata: {}, + type: 'code' + } + ], + executionMode: 'block', + isModule: false + } + ], + settings: {} + } + }; + + manager.storeOriginalProject('project-ephemeral-offset', 'notebook-1', projectData); + + const notebookData = { + cells: [ + { + kind: 2, + value: 'print("scratch")', + languageId: 'python', + metadata: { + id: 'ephemeral-block-id', + is_ephemeral: true, + agent_source_block_id: 'agent-1' + } + }, + { + kind: 2, + value: 'print("kept")', + languageId: 'python', + metadata: {} + } + ], + metadata: { + deepnoteProjectId: 'project-ephemeral-offset', + deepnoteNotebookId: 'notebook-1' + } + }; + + const result = await serializer.serializeNotebook(notebookData 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].id, 'real-block-id', 'Stripped cell should recover its ID'); + }); }); suite('integration scenarios', () => { From 7eab35d738d41a7d34ce7217bb47aefd02401cb4 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 13 Aug 2026 06:19:58 +0000 Subject: [PATCH 69/80] fix(deepnote): reload when an external edit changes only block metadata contentActuallyChanged compared cell count, kind, languageId and source. An external edit that changed only deepnote_agent_model or a block id was read as "no change", the reload was skipped, and the next save wrote the stale in-memory value back over the file -- silently reverting the edit. Editing a .deepnote on disk while it is open is the case this watcher exists for. Comparing raw cell metadata would be worse than the bug: the save path rewrites contentHash and normalizes sortingKey every time, so every user save would reload, and reloading replaces all cells and destroys agent scratch cells. So compare what the file actually carries -- run both sides through convertCellToBlock, the same conversion the serializer saves through, and compare the resulting block. Anything the write path derives, normalizes or strips is excluded because it never reaches block.metadata, so there is no field list here to drift out of sync. Co-Authored-By: Claude Opus 4.6 Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- .../deepnote/deepnoteFileChangeWatcher.ts | 51 +++++- .../deepnoteFileChangeWatcher.unit.test.ts | 162 ++++++++++++++++++ 2 files changed, 204 insertions(+), 9 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index 5a304abc81..b6c3f47339 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -12,6 +12,7 @@ import { } from 'vscode'; import { inject, injectable, optional } from 'inversify'; import type { DeepnoteBlock } from '@deepnote/blocks'; +import fastDeepEqual from 'fast-deep-equal'; import { IControllerRegistration } from '../controllers/types'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; @@ -159,9 +160,9 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } /** - * Checks whether the source code content has actually changed between the - * live notebook and the new cells from disk. If only outputs differ (disk - * has fewer/no outputs), it's an auto-save of stripped content — skip reload. + * Checks whether anything the file carries has actually changed between the live notebook and + * the new cells from disk. Outputs and execution state are ignored: in snapshot mode the main + * file has them stripped, so our own auto-save would otherwise look like an external edit. */ private contentActuallyChanged(notebook: NotebookDocument, newCells: NotebookCellData[]): boolean { // Ephemeral cells aren't persisted; counting them looks like an external delete. @@ -169,12 +170,24 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic if (liveCells.length !== newCells.length) { return true; } - return liveCells.some( - (live, i) => - live.kind !== newCells[i].kind || - live.document.languageId !== newCells[i].languageId || - live.document.getText() !== newCells[i].value - ); + + return liveCells.some((live, i) => { + const fromDisk = newCells[i]; + + if ( + live.kind !== fromDisk.kind || + live.document.languageId !== fromDisk.languageId || + live.document.getText() !== fromDisk.value || + getBlockId(live) !== getBlockId(fromDisk) + ) { + return true; + } + + const liveBlock = this.persistedBlock(live, i); + const diskBlock = this.persistedBlock(fromDisk, i); + + return liveBlock.type !== diskBlock.type || !fastDeepEqual(liveBlock.metadata, diskBlock.metadata); + }); } /** @@ -649,6 +662,26 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic return true; } + /** + * The block a cell would be written as, produced by the very conversion the serializer runs on + * save. Comparing that instead of raw cell metadata keeps the check honest in both directions: + * whatever the write path derives, normalizes or strips ends up on the block rather than in + * `block.metadata`, while everything an external editor can put in the file survives. + * Outputs are left out — they are deliberately not compared and converting them is the + * expensive part. + */ + private persistedBlock(cell: NotebookCell | NotebookCellData, index: number): DeepnoteBlock { + const isLiveCell = 'document' in cell; + const projection = new NotebookCellData( + cell.kind, + isLiveCell ? cell.document.getText() : cell.value, + isLiveCell ? cell.document.languageId : cell.languageId + ); + projection.metadata = { ...cell.metadata }; + + return this.converter.convertCellToBlock(projection, index); + } + private selfWriteKey(uri: Uri): string { return uri.with({ query: '', fragment: '' }).toString(); } diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index 91e8a795e4..adbbca7dd2 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -158,6 +158,65 @@ project: content: print("hello") `; + // The same block after our own save: the serializer recomputes contentHash and the package + // normalizes sortingKey (generateSortingKey -> String(index).padStart(6, '0')). + const rewrittenBySaveYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00Z' +project: + id: e132b172-b114-410e-8331-011517db664f + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + type: code + sortingKey: '000000' + blockGroup: '1' + contentHash: 'sha256:beefbeef' + content: print("hello") +`; + + const agentYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00Z' +project: + id: e132b172-b114-410e-8331-011517db664f + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + type: agent + sortingKey: '000000' + blockGroup: '1' + content: summarise the dataframe + metadata: + deepnote_agent_model: gpt-5 +`; + + const renamedBlockYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00Z' +project: + id: e132b172-b114-410e-8331-011517db664f + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-9 + type: code + sortingKey: '000000' + blockGroup: '1' + content: print("hello") +`; + test('should skip reload when content matches notebook cells', async () => { const uri = Uri.file('/workspace/test.deepnote'); // Create a notebook whose cell content already matches validYaml @@ -441,6 +500,109 @@ project: assert.strictEqual(applyEditCount, 0, 'applyEdit should NOT be called for auto-save (same source)'); }); + test('should skip reload when only save-rewritten block fields differ', async function () { + this.timeout(8000); + const uri = Uri.file('/workspace/test.deepnote'); + // State right after our own save: contentHash recomputed, sortingKey normalized, outputs + // stripped from the main file. None of that is an edit, and reloading would drop the + // agent's ephemeral cells. + const notebook = createMockNotebook({ + uri, + cells: [ + { + metadata: { + id: 'block-1', + __deepnoteBlockId: 'block-1', + __hadOutputs: true, + __deepnotePocket: { + blockGroup: '1', + contentHash: 'sha256:deadbeef', + sortingKey: 'a0', + type: 'code' + } + }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(rewrittenBySaveYaml); + + onDidChangeFile.fire(uri); + + await waitFor(() => readFileCalls > 0); + await new Promise((resolve) => setTimeout(resolve, autoSaveGraceMs)); + + assert.strictEqual(applyEditCount, 0, 'a no-op save must not reload'); + assert.strictEqual(saveCount, 0, 'a no-op save must not trigger another save'); + }); + + test('should reload when only block metadata changed on disk', async function () { + this.timeout(8000); + const uri = Uri.file('/workspace/test.deepnote'); + // Live cell holds the model the file carried before an external editor rewrote it; + // source, kind and language are identical on both sides. + const notebook = createMockNotebook({ + uri, + cells: [ + { + metadata: { + deepnote_agent_model: 'gpt-4o', + id: 'block-1', + __deepnoteBlockId: 'block-1', + __hadOutputs: false, + __deepnotePocket: { blockGroup: '1', sortingKey: '000000', type: 'agent' } + }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'summarise the dataframe', languageId: 'plaintext' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(agentYaml); + + onDidChangeFile.fire(uri); + + await waitFor(() => applyEditCount > 0); + + assert.isAtLeast(applyEditCount, 1, 'a metadata-only external edit must reload'); + }); + + test('should reload when a block id changed on disk', async function () { + this.timeout(8000); + const uri = Uri.file('/workspace/test.deepnote'); + const notebook = createMockNotebook({ + uri, + cells: [ + { + metadata: { + id: 'block-1', + __deepnoteBlockId: 'block-1', + __hadOutputs: false, + __deepnotePocket: { blockGroup: '1', sortingKey: '000000', type: 'code' } + }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(renamedBlockYaml); + + onDidChangeFile.fire(uri); + + await waitFor(() => applyEditCount > 0); + + assert.isAtLeast(applyEditCount, 1, 'an external block-id change must reload'); + }); + suite('normalized one-shot self-write markers', () => { // YAML matching `print("hello")` lives in `validYaml` above. const helloYaml = validYaml; From a5b12979f8a91d3ee59b01c86b190034c8be3709 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 13 Aug 2026 06:20:14 +0000 Subject: [PATCH 70/80] fix(agent-block): make a run one batch that a stop actually ends Three faults in the same execution frame, all from splitting a Run All around agent cells and letting each generated cell re-enter it. Run All no longer continues past a failure. Before the split this function was one body where a failing segment's `return` ended the whole run; splitting it demoted those returns to ending one segment, so failing Python -> agent -> Python ran everything. They rethrow again, which is the pre-existing control flow rather than new bookkeeping. Cancellation needs more, because a cancelled execution resolves rather than rejects: the queue already latches that verdict, so expose it as INotebookKernelExecution.failed instead of tracking it again. Queue completion is now per gesture, not per queue. An agent run opens a fresh CellExecutionQueue per generated cell, each announcing completion, so SnapshotService saved and cleared execution state during ordinary LLM pauses. The controller owns the batch, so it announces completion once, when its re-entrancy depth unwinds to zero. Retiring a run's metadata moved off the save. Clearing it in performSnapshotSave's finally meant the save that follows a run -- and any file save after it -- serialized nothing, and it wiped the captured environment, so an agent run re-ran pip freeze per generated cell. It is now dropped when the next run starts, signalled by the same frame that announces completion so a run that opens no kernel queue still retires the previous one. Stopping an agent run does something. interruptHandler leaves NotebookCellExecution.token inert, so the agent never saw a stop: the kernel interrupt ended its in-flight cell, which the model read as a failure worth retrying, and a cell cancelled before it started left the agent waiting out a five minute timeout. The controller now owns a cancellation source per notebook, cancelled before the kernel interrupt so the agent sees the stop first. The model call itself still runs to the end of its turn -- that needs the AbortSignal support sitting unreleased in runtime-core, and executeAgentCell documents where it plugs in. Co-Authored-By: Claude Opus 4.6 Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- src/kernels/execution/cellExecutionQueue.ts | 5 - src/kernels/kernelExecution.ts | 4 + src/kernels/types.ts | 6 + .../controllers/vscodeNotebookController.ts | 86 ++++-- .../vscodeNotebookController.unit.test.ts | 281 ++++++++++++++++-- .../deepnote/agentCellExecutionHandler.ts | 42 ++- .../agentCellExecutionHandler.unit.test.ts | 105 ++++++- .../deepnote/snapshots/snapshotService.ts | 27 +- .../snapshots/snapshotService.unit.test.ts | 86 ++++++ .../notebooks/cellExecutionStateService.ts | 17 ++ 10 files changed, 592 insertions(+), 67 deletions(-) diff --git a/src/kernels/execution/cellExecutionQueue.ts b/src/kernels/execution/cellExecutionQueue.ts index 66d78822e4..6b1597f1b7 100644 --- a/src/kernels/execution/cellExecutionQueue.ts +++ b/src/kernels/execution/cellExecutionQueue.ts @@ -325,10 +325,5 @@ export class CellExecutionQueue implements Disposable { break; } } - - // Notify listeners that execution queue is complete - if (this.notebook) { - notebookCellExecutions.notifyQueueComplete(this.notebook.uri.toString()); - } } } diff --git a/src/kernels/kernelExecution.ts b/src/kernels/kernelExecution.ts index 99b36ce964..8d15d6599b 100644 --- a/src/kernels/kernelExecution.ts +++ b/src/kernels/kernelExecution.ts @@ -135,6 +135,10 @@ export class NotebookKernelExecution implements INotebookKernelExecution { } }); } + public get failed(): boolean { + return this.documentExecutions.get(this.notebook)?.failed === true; + } + public get pendingCells(): readonly NotebookCell[] { return this.documentExecutions.get(this.notebook)?.queue || []; } diff --git a/src/kernels/types.ts b/src/kernels/types.ts index fbc44b2ebc..642ca583ff 100644 --- a/src/kernels/types.ts +++ b/src/kernels/types.ts @@ -461,6 +461,12 @@ export interface INotebookKernelExecution { * Total execution count on this kernel */ readonly executionCount: number; + /** + * Whether the cell execution queue stopped early: a cell failed, or queued cells were cancelled + * (interrupt/restart). Cancelled cell executions resolve rather than reject, so this is the only + * way to tell an interrupted run from a clean one. + */ + readonly failed: boolean; readonly onDidReceiveDisplayUpdate: Event; /** * Cells that are still being executed (or pending). diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index 3238d968a3..c127dbfc18 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -611,6 +611,11 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont private handleInterrupt(notebook: NotebookDocument) { logger.debug(`VS Code interrupted kernel for ${getDisplayPath(notebook.uri)}`); notebook.getCells().forEach((cell) => traceCellMessage(cell, 'Cell cancellation requested')); + // Before the kernel interrupt: that interrupt ends the generated cell the agent is waiting on, + // and an agent that has not yet seen the stop reads the ended cell as a failure worth retrying. + // Setting an interruptHandler leaves NotebookCellExecution.token inert, so this is the agent's + // only stop signal. + this.agentCancellations.get(notebook)?.cancel(); commands .executeCommand(Commands.InterruptKernel, { notebookEditor: { notebookUri: notebook.uri } }) .then(noop, (ex) => logger.error('Failed to interrupt', ex)); @@ -634,7 +639,15 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont return currentExecution; } + /** Stop signal for the agent cell currently running in a notebook, read by `handleInterrupt`. */ + private readonly agentCancellations = new WeakMap(); private cellQueue = new WeakMap(); + /** + * Frames of `executeQueuedCells` in flight per notebook. An agent cell dispatches each cell it + * generates through `notebook.cell.execute`, which lands back here while the outer batch is still + * running; completion belongs to the gesture, so only the frame unwinding to zero announces it. + */ + private readonly executionDepth = new WeakMap(); private async executeQueuedCells(doc: NotebookDocument) { if (!this.cellQueue.has(doc)) { return; @@ -642,36 +655,63 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont const queuedCells = this.cellQueue.get(doc) || []; // Clear before await — agent runs can re-enter with an empty queue. this.cellQueue.delete(doc); + // Nothing may run between here and the `try`: a throw in between would strand the depth above + // zero and silence completion — and with it the snapshot save — for the rest of the session. + const depthOnEntry = this.executionDepth.get(doc) ?? 0; - const cellsToExecute = await removeEphemeralCellsForAgentBlocks(doc, queuedCells); + this.executionDepth.set(doc, depthOnEntry + 1); - let pendingKernelCells: NotebookCell[] = []; - let ranAgentCell = false; + if (depthOnEntry === 0) { + // Paired with the notify below so a run that opens no kernel queue — an agent cell that + // generates nothing — still marks a boundary the previous run's metadata is retired at. + notebookCellExecutions.notifyQueueStart(doc.uri.toString()); + } try { + const cellsToExecute = await removeEphemeralCellsForAgentBlocks(doc, queuedCells); + + let pendingKernelCells: NotebookCell[] = []; + for (const cell of cellsToExecute) { if (!isAgentCell(cell)) { pendingKernelCells.push(cell); continue; } - ranAgentCell = true; await this.executeKernelCells(doc, pendingKernelCells); pendingKernelCells = []; logger.trace(`Executing agent cell ${cell.index} for ${getDisplayPath(doc.uri)} without kernel`); - await executeAgentCell( - cell, - this.controller, - this.serviceContainer.get(IEncryptedStorage) - ).catch(noop); + + const agentCancellation = new CancellationTokenSource(); + + this.agentCancellations.set(doc, agentCancellation); + + try { + await executeAgentCell( + cell, + this.controller, + this.serviceContainer.get(IEncryptedStorage), + agentCancellation.token + ).catch(noop); + } finally { + this.agentCancellations.delete(doc); + agentCancellation.dispose(); + } } await this.executeKernelCells(doc, pendingKernelCells); + } catch (ex) { + // The failing cell already carries the error; unwinding here only stops the agent cell and the + // segments after it, the way CellExecutionQueue stops a run once a cell fails. + logger.debug(`Stopped the rest of the batch for ${getDisplayPath(doc.uri)}`, ex); } finally { - // Batches without an agent cell (including reentrant ephemeral-cell runs) already notify via - // CellExecutionQueue. Explicit notify covers agent-only and agent-then-kernel batches. - if (ranAgentCell) { + // Re-read rather than reuse a captured local: nested frames decrement the shared value. + const depth = (this.executionDepth.get(doc) ?? 1) - 1; + + this.executionDepth.set(doc, depth); + + if (depth === 0) { notebookCellExecutions.notifyQueueComplete(doc.uri.toString()); } } @@ -724,12 +764,12 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont } catch (ex) { if (ex instanceof KernelError) { // Kernel errors would have been handled and displayed - return; + throw ex; } ex = WrappedError.unwrap(ex); if (ex instanceof CellExecutionOutputError) { // CellExecution already wrote this message to the cell output. - return; + throw ex; } if (!isCancellationError(ex)) { logger.error(`Error in notebook cell execution`, ex); @@ -748,10 +788,8 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont await errorHandler.getErrorMessageForDisplayInCellOutput(ex, currentContext, doc.uri), isCancelled ); - } - if (!kernel) { - return; + throw ex; } const kernelExecution = this.kernelProvider.getKernelExecution(kernel); @@ -777,11 +815,11 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont } catch (ex) { if (ex instanceof KernelError) { // Kernel errors would have been handled and displayed - return; + throw ex; } ex = WrappedError.unwrap(ex); if (ex instanceof CellExecutionOutputError) { - return; + throw ex; } if (!isCancellationError(ex)) { logger.error(`Error in cell execution`, ex); @@ -795,9 +833,17 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont await errorHandler.getErrorMessageForDisplayInCellOutput(ex, currentContext, doc.uri), isCancelled ); + + throw ex; } }) - ).catch(noop); + ); + + if (kernelExecution.failed) { + // An interrupt resolves the cells it cancelled, so the awaits above stay clean; the queue's + // own verdict is the only thing that separates a stopped run from a successful one. + throw new CancellationError(); + } } public async startKernel(notebook: NotebookDocument) { diff --git a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts index 318b23aabd..c7cbdd2ffe 100644 --- a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts +++ b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts @@ -21,12 +21,15 @@ import { VSCodeNotebookController, warnWhenUsingOutdatedPython } from './vscodeN import { IKernel, IKernelProvider, + INotebookKernelExecution, KernelConnectionMetadata, LiveRemoteKernelConnectionMetadata, LocalKernelConnectionMetadata, LocalKernelSpecConnectionMetadata, RemoteKernelSpecConnectionMetadata } from '../../kernels/types'; +import { KernelError } from '../../kernels/errors/kernelError'; +import { LastCellExecutionTracker } from '../../kernels/execution/lastCellExecutionTracker'; import { anything, deepEqual, instance, mock, verify, when } from 'ts-mockito'; import { ITelemetryService } from '../../platform/analytics/types'; import { IEncryptedStorage } from '../../platform/common/application/types'; @@ -985,6 +988,130 @@ suite(`Notebook Controller`, function () { sinon.restore(); }); + // The connected kernel is a plain object, not `instance(mock())`: a ts-mockito proxy + // answers `then` with a function, so awaiting the connect promise would never settle. + // Anything left unstubbed throws inside the per-cell try and would abort the batch for the wrong reason. + function stubKernelForExecution(kernelExecution: Partial): void { + const neverFires = () => new Disposable(() => undefined); + const connectedKernel = { + controller: { + id: 'test-controller-id', + createNotebookCellExecution: (cell: NotebookCell) => + vscodeController.controller.createNotebookCellExecution(cell) + }, + disposing: false, + onDisposed: neverFires, + onStatusChanged: neverFires + } as unknown as IKernel; + + const oldConnectToNotebook = KernelConnector.connectToNotebookKernel; + KernelConnector.connectToNotebookKernel = async () => connectedKernel; + disposables.push(new Disposable(() => (KernelConnector.connectToNotebookKernel = oldConnectToNotebook))); + + when(serviceContainer.get(LastCellExecutionTracker)).thenReturn( + instance(mock()) + ); + when(kernelProvider.getKernelExecution(anything())).thenReturn(kernelExecution as INotebookKernelExecution); + } + + test('a failed kernel segment stops the agent cell and the trailing segment', async function () { + // Catches: executeKernelCells swallowing a KernelError, so Run All continues past a failed cell. + const { + notebook, + cells: [failingCell, agentCell, trailingCell] + } = createMockNotebookWithCells([ + { metadata: { id: 'code-1' }, text: 'raise ValueError()' }, + { metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, text: 'Test prompt' }, + { metadata: { id: 'code-2' }, text: 'print(2)' } + ]); + + const executedIndexes: number[] = []; + stubKernelForExecution({ + failed: false, + executeCell: async (cell: NotebookCell) => { + executedIndexes.push(cell.index); + throw new KernelError({ ename: 'ValueError', evalue: 'boom', traceback: [] }); + } + }); + + await vscodeController.controller.executeHandler( + [failingCell, agentCell, trailingCell], + notebook, + vscodeController.controller + ); + + assert.deepStrictEqual(executedIndexes, [0], 'the trailing segment must not run after a failure'); + assert.isFalse( + createNotebookCellExecutionStub.getCalls().some((call) => call.args[0] === agentCell), + 'the agent cell must not start after a failed segment' + ); + }); + + test('an interrupted kernel segment stops the agent cell even though its cells resolve', async function () { + // Catches: relying on a rejection alone - cancelled cell executions resolve + // (CellExecution.completedDueToCancellation), so only the queue's verdict shows the interrupt. + const { + notebook, + cells: [interruptedCell, agentCell, trailingCell] + } = createMockNotebookWithCells([ + { metadata: { id: 'code-1' }, text: 'time.sleep(30)' }, + { metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, text: 'Test prompt' }, + { metadata: { id: 'code-2' }, text: 'print(2)' } + ]); + + const executedIndexes: number[] = []; + stubKernelForExecution({ + failed: true, + executeCell: async (cell: NotebookCell) => { + executedIndexes.push(cell.index); + } + }); + + await vscodeController.controller.executeHandler( + [interruptedCell, agentCell, trailingCell], + notebook, + vscodeController.controller + ); + + assert.deepStrictEqual(executedIndexes, [0], 'the trailing segment must not run after an interrupt'); + assert.isFalse( + createNotebookCellExecutionStub.getCalls().some((call) => call.args[0] === agentCell), + 'the agent cell must not start after an interrupt' + ); + }); + + test('a clean kernel segment still runs the agent cell and the trailing segment', async function () { + // Catches: aborting the batch when nothing failed (e.g. consulting the queue verdict too early). + const { + notebook, + cells: [firstCell, agentCell, trailingCell] + } = createMockNotebookWithCells([ + { metadata: { id: 'code-1' }, text: 'x = 1' }, + { metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, text: 'Test prompt' }, + { metadata: { id: 'code-2' }, text: 'print(2)' } + ]); + + const executedIndexes: number[] = []; + stubKernelForExecution({ + failed: false, + executeCell: async (cell: NotebookCell) => { + executedIndexes.push(cell.index); + } + }); + + await vscodeController.controller.executeHandler( + [firstCell, agentCell, trailingCell], + notebook, + vscodeController.controller + ); + + assert.deepStrictEqual(executedIndexes, [0, 2], 'both kernel segments must run when nothing failed'); + assert.isTrue( + createNotebookCellExecutionStub.getCalls().some((call) => call.args[0] === agentCell), + 'the agent cell must run between the segments' + ); + }); + test('agent-only batch fires notifyQueueComplete (arms deferred snapshot save)', async function () { // Catches: agent-only runs never reach CellExecutionQueue, so snapshot save never arms. const { @@ -1016,8 +1143,9 @@ suite(`Notebook Controller`, function () { assert.isTrue(createNotebookCellExecutionStub.calledOnce, 'agent cell should run through executeAgentCell'); }); - test('kernel-only batch after agent batch does not fire a second explicit queue completion notify', async function () { - // Catches: explicit notify on kernel-only executeQueuedCells (e.g. reentrant ephemeral runs) when ranAgentCell is false. + test('an agent batch and a later kernel-only batch each fire exactly one completion', async function () { + // Catches: tying completion to an "this batch ran an agent cell" flag — CellExecutionQueue no + // longer notifies, so a plain Run would arm no deferred snapshot save. const { notebook: agentNotebook, cells: [agentCell, codeCell] @@ -1027,33 +1155,26 @@ suite(`Notebook Controller`, function () { text: 'Test prompt' }, { - metadata: { id: 'ephemeral-code-1' }, + metadata: { id: 'code-1' }, text: 'print(1)' } ]); - const executeHandler = vscodeController.controller.executeHandler; - assert.isDefined(executeHandler); - - notifyQueueCompleteSpy.resetHistory(); - - await executeHandler([agentCell], agentNotebook, vscodeController.controller); + stubKernelForExecution({ failed: false, executeCell: async () => undefined }); - try { - await executeHandler([codeCell], agentNotebook, vscodeController.controller); - } catch { - // Kernel harness may not fully mock cell execution startup. - } + await vscodeController.controller.executeHandler([agentCell], agentNotebook, vscodeController.controller); + await vscodeController.controller.executeHandler([codeCell], agentNotebook, vscodeController.controller); - assert.strictEqual( - notifyQueueCompleteSpy.callCount, - 1, - 'only the agent batch should fire explicit queue completion' + assert.deepStrictEqual( + notifyQueueCompleteSpy.getCalls().map((call) => call.args[0]), + [agentNotebook.uri.toString(), agentNotebook.uri.toString()], + 'each gesture fires one completion for its own notebook' ); }); - test('kernel-only batch does not fire explicit queue completion notify', async function () { - // Catches: explicit notify on pure kernel batches that already signal via CellExecutionQueue. + test('a kernel-only batch fires the completion itself', async function () { + // Catches: leaving completion to CellExecutionQueue, which no longer announces it — the + // deferred snapshot save would never arm for an ordinary run. const { notebook: codeNotebook, cells: [codeCell] @@ -1064,19 +1185,121 @@ suite(`Notebook Controller`, function () { } ]); - const executeHandler = vscodeController.controller.executeHandler; - assert.isDefined(executeHandler); + stubKernelForExecution({ failed: false, executeCell: async () => undefined }); + + await vscodeController.controller.executeHandler([codeCell], codeNotebook, vscodeController.controller); + + assert.isTrue(notifyQueueCompleteSpy.calledOnce, 'a kernel-only batch must fire one completion'); + assert.strictEqual(notifyQueueCompleteSpy.firstCall.args[0], codeNotebook.uri.toString()); + }); + + test('a run that re-enters once per generated cell fires exactly one completion', async function () { + // Catches the N+1 snapshot saves an agent run produced: each generated cell is dispatched + // through `notebook.cell.execute`, which lands back in this handler while the outer batch is + // still in flight. Completion belongs to the gesture, not to every nested run. + const { + notebook, + cells: [runCell, generatedFirst, generatedSecond] + } = createMockNotebookWithCells([ + { metadata: { id: 'code-1' }, text: 'x = 1' }, + { metadata: { id: 'generated-1' }, text: 'print(1)' }, + { metadata: { id: 'generated-2' }, text: 'print(2)' } + ]); + + const executedIds: string[] = []; + stubKernelForExecution({ + failed: false, + executeCell: async (cell: NotebookCell) => { + executedIds.push(cell.metadata.id as string); + + if (cell !== runCell) { + return; + } + + await vscodeController.controller.executeHandler( + [generatedFirst], + notebook, + vscodeController.controller + ); + await vscodeController.controller.executeHandler( + [generatedSecond], + notebook, + vscodeController.controller + ); + } + }); + + await vscodeController.controller.executeHandler([runCell], notebook, vscodeController.controller); + + assert.deepStrictEqual( + executedIds, + ['code-1', 'generated-1', 'generated-2'], + 'both re-entrant runs must have executed' + ); + assert.strictEqual( + notifyQueueCompleteSpy.callCount, + 1, + 'the gesture owns the completion; the runs nested inside it must not fire their own' + ); + }); + + test('a batch that fails still fires its own completion and does not silence the next one', async function () { + // Catches: skipping the completion when the batch unwinds — an interrupted run must still + // snapshot what it produced, and must not leave the re-entrancy depth above zero. + const { + notebook, + cells: [failingCell, laterCell] + } = createMockNotebookWithCells([ + { metadata: { id: 'code-1' }, text: 'raise ValueError()' }, + { metadata: { id: 'code-2' }, text: 'x = 1' } + ]); + + stubKernelForExecution({ + failed: false, + executeCell: async (cell: NotebookCell) => { + if (cell === failingCell) { + throw new KernelError({ ename: 'ValueError', evalue: 'boom', traceback: [] }); + } + } + }); + + await vscodeController.controller.executeHandler([failingCell], notebook, vscodeController.controller); + await vscodeController.controller.executeHandler([laterCell], notebook, vscodeController.controller); + + assert.strictEqual(notifyQueueCompleteSpy.callCount, 2, 'a failed batch must not silence later runs'); + }); + + test('a rejected scratch-cell cleanup neither escapes nor strands the completion', async function () { + // Catches: clearing prior scratch cells outside the frame that owns the re-entrancy depth — + // a rejected workspace edit would escape before the depth is handed back, silencing every + // completion for the rest of the session. + const { + notebook, + cells: [agentCell, , laterCell] + } = createMockNotebookWithCells([ + { metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, text: 'Test prompt' }, + { + metadata: { agent_source_block_id: 'agent-block-1', id: 'scratch-1', is_ephemeral: true }, + text: 'print(1)' + }, + { metadata: { id: 'code-1' }, text: 'x = 1' } + ]); + + stubKernelForExecution({ failed: false, executeCell: async () => undefined }); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReject(new Error('edit failed')); + + let escaped: unknown; try { - await executeHandler([codeCell], codeNotebook, vscodeController.controller); - } catch { - // Kernel harness may not fully mock cell execution startup. + await vscodeController.controller.executeHandler([agentCell], notebook, vscodeController.controller); + } catch (ex) { + escaped = ex; } - assert.isFalse( - notifyQueueCompleteSpy.called, - 'batches without agent cells must rely on CellExecutionQueue for completion notify' - ); + await vscodeController.controller.executeHandler([laterCell], notebook, vscodeController.controller); + + assert.isUndefined(escaped, 'a rejected cleanup edit must not escape the execute handler'); + assert.strictEqual(notifyQueueCompleteSpy.callCount, 2, 'a cleanup failure must not strand the depth'); }); }); }); diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index 748aecbc79..d98c35f331 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -30,6 +30,7 @@ import { dispose } from '../../platform/common/utils/lifecycle'; import { uuidUtils } from '../../platform/common/uuid'; import { ServiceContainer } from '../../platform/ioc/container'; import { logger } from '../../platform/logging'; +import { Cancellation } from '../../platform/common/cancellation'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; import { IDeepnoteNotebookManager } from '../types'; import { @@ -153,14 +154,28 @@ export interface ExecuteAgentCellOptions { executeAgentBlockFn?: typeof executeAgentBlock; } +/** + * True for both the host's own cancellation and the `AbortError` that runtime-core raises from + * `signal.throwIfAborted()`. `isCancellationError` covers only the former. + */ +function isStopped(error: unknown): boolean { + return error instanceof CancellationError || (error instanceof Error && error.name === 'AbortError'); +} + /** * Runs an agent block into the cell output and inserts generated cells below. * Call `removeEphemeralCellsForAgentBlocks` on the batch first. Never rejects — errors become stderr on the cell. + * + * `token` stops the run. It reaches the model only indirectly: the host refuses tool calls and throws + * from the event callback, so an in-flight model turn still finishes. Once `AgentBlockContext` carries + * an `AbortSignal` (present in runtime-core's `main`, unreleased), bridge the token to one and pass it + * as `signal` — runtime-core forwards it to `agent.stream`, which aborts the request itself. */ export async function executeAgentCell( cell: NotebookCell, controller: NotebookController, encryptedStorage: IEncryptedStorage, + token: CancellationToken, options?: ExecuteAgentCellOptions ): Promise { const executeAgentBlockFn = options?.executeAgentBlockFn ?? executeAgentBlock; @@ -208,7 +223,11 @@ export async function executeAgentCell( openAiToken, ...getProjectAgentContext(cell.notebook), notebookContext, + // The guards sit outside the `try`s: those turn every throw into a string the model reads + // as a retryable tool failure, which is how a stop used to make the agent do more work. addMarkdownBlock: async ({ content }: { content: string }) => { + Cancellation.throwIfCanceled(token); + try { await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'markdown', content); @@ -218,6 +237,8 @@ export async function executeAgentCell( } }, addAndExecuteCodeBlock: async ({ code }: { code: string }) => { + Cancellation.throwIfCanceled(token); + try { const insertedCell = await insertEphemeralCell( cell.notebook, @@ -227,15 +248,23 @@ export async function executeAgentCell( code ); - const { success, outputs, error } = await executeEphemeralCell(insertedCell, execution.token); + const { success, outputs, error } = await executeEphemeralCell(insertedCell, token); const outputText = error ?? describeExecutionOutputs(outputs); return success ? `Output:\n${outputText}` : `Execution failed:\n${outputText}`; } catch (error) { + if (isStopped(error)) { + throw error; + } + return `Execution error: ${toError(error).message}`; } }, onAgentEvent: async (event: AgentStreamEvent) => { + // Runs in runtime-core's own stream loop, which has no catch — the one place the host + // can end the run rather than merely refuse it. + Cancellation.throwIfCanceled(token); + logger.trace(`Agent event: ${event.type}`); let delta = lastAgentEventType != null && lastAgentEventType !== event.type ? `\n\n` : ''; @@ -277,6 +306,17 @@ export async function executeAgentCell( execution.end(true, Date.now()); } catch (error) { + if (isStopped(error)) { + logger.info('Agent cell execution stopped'); + + const stoppedOutput = new NotebookCellOutput([NotebookCellOutputItem.stderr('[Agent] Stopped')]); + + await execution.appendOutput([stoppedOutput]).then(undefined, () => undefined); + execution.end(false, Date.now()); + + return; + } + // logger.error does not print stacks unless isJupyterError — log stack explicitly. logger.error('Agent cell execution failed', error); if (error instanceof Error) { diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index 18c0b96ad7..1e8f72652e 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -3,6 +3,7 @@ import * as sinon from 'sinon'; import { anything, capture, instance, mock, reset, verify, when } from 'ts-mockito'; import { CancellationError, + CancellationToken, CancellationTokenSource, Disposable, NotebookCell, @@ -173,11 +174,17 @@ suite('AgentCellExecutionHandler', () => { let executeAgentBlockStub: sinon.SinonStub; let mockServiceContainer: ServiceContainer; let encryptedStorage: IEncryptedStorage; + let neverCancelled: CancellationToken; setup(() => { secretStorage.clear(); secretStorage.set('openAiApiKey', 'test-key'); encryptedStorage = createEncryptedStorageFake(secretStorage); + + const neverCancelledSource = new CancellationTokenSource(); + + neverCancelled = neverCancelledSource.token; + disposables.push(neverCancelledSource); mockServiceContainer = stubServiceContainerInstance(); disposables.push(new Disposable(() => sinon.restore())); @@ -237,7 +244,7 @@ suite('AgentCellExecutionHandler', () => { test('creates execution, clears output, sets planning output, and ends successfully', async () => { const cell = createAgentCell('Analyze data'); - await executeAgentCell(cell, mockController, encryptedStorage, { + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { executeAgentBlockFn: executeAgentBlockStub }); @@ -268,7 +275,7 @@ suite('AgentCellExecutionHandler', () => { const cell = createAgentCell(); - await executeAgentCell(cell, mockController, encryptedStorage, { + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { executeAgentBlockFn: executeAgentBlockStub }); @@ -290,7 +297,7 @@ suite('AgentCellExecutionHandler', () => { const cell = createAgentCell(); - await executeAgentCell(cell, mockController, encryptedStorage, { + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { executeAgentBlockFn: executeAgentBlockStub }); @@ -304,7 +311,7 @@ suite('AgentCellExecutionHandler', () => { const cell = createAgentCell(); - await executeAgentCell(cell, mockController, encryptedStorage, { + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { executeAgentBlockFn: executeAgentBlockStub }); @@ -325,7 +332,7 @@ suite('AgentCellExecutionHandler', () => { test('handles empty prompt', async () => { const cell = createAgentCell(''); - await executeAgentCell(cell, mockController, encryptedStorage, { + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { executeAgentBlockFn: executeAgentBlockStub }); @@ -343,7 +350,7 @@ suite('AgentCellExecutionHandler', () => { const cell = createAgentCell(); - await executeAgentCell(cell, mockController, encryptedStorage, { + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { executeAgentBlockFn: executeAgentBlockStub }); @@ -365,7 +372,7 @@ suite('AgentCellExecutionHandler', () => { }); const { agentCell, cells } = createAgentCellInMutableNotebook([previousResult]); - await executeAgentCell(agentCell, mockController, encryptedStorage, { + await executeAgentCell(agentCell, mockController, encryptedStorage, neverCancelled, { executeAgentBlockFn: executeAgentBlockStub }); @@ -388,7 +395,7 @@ suite('AgentCellExecutionHandler', () => { return { finalOutput: '' }; }); - await executeAgentCell(agentCell, mockController, encryptedStorage, { + await executeAgentCell(agentCell, mockController, encryptedStorage, neverCancelled, { executeAgentBlockFn: executeAgentBlockStub }); @@ -413,7 +420,7 @@ suite('AgentCellExecutionHandler', () => { return { finalOutput: '' }; }); - await executeAgentCell(agentCell, mockController, encryptedStorage, { + await executeAgentCell(agentCell, mockController, encryptedStorage, neverCancelled, { executeAgentBlockFn: executeAgentBlockStub }); @@ -440,7 +447,7 @@ suite('AgentCellExecutionHandler', () => { notebookMetadata: { deepnoteProjectId: 'project-1', deepnoteNotebookId: 'notebook-1' } }); - await executeAgentCell(cell, mockController, encryptedStorage, { + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { executeAgentBlockFn: executeAgentBlockStub }); @@ -448,6 +455,84 @@ suite('AgentCellExecutionHandler', () => { expect(context.mcpServers).to.deep.equal(mcpServers); expect(context.integrations).to.deep.equal(integrations); }); + + suite('cancellation', () => { + let runTokenSource: CancellationTokenSource; + + setup(() => { + runTokenSource = new CancellationTokenSource(); + disposables.push(runTokenSource); + }); + + test('refuses to add a generated cell once stopped', async () => { + const { agentCell, cells } = createAgentCellInMutableNotebook(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + runTokenSource.cancel(); + await context.addMarkdownBlock({ content: 'after the stop' }); + + return { finalOutput: '' }; + }); + + await executeAgentCell(agentCell, mockController, encryptedStorage, runTokenSource.token, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(cells.map((cell) => cell.document.getText())).to.deep.equal(['Test prompt']); + }); + + test('stops at the next stream event rather than running to completion', async () => { + const cell = createAgentCell(); + let eventsAfterStop = 0; + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + runTokenSource.cancel(); + await context.onAgentEvent?.({ type: 'text_delta', text: 'first' }); + eventsAfterStop += 1; + await context.onAgentEvent?.({ type: 'text_delta', text: 'second' }); + + return { finalOutput: '' }; + }); + + await executeAgentCell(cell, mockController, encryptedStorage, runTokenSource.token, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(eventsAfterStop).to.equal(0); + }); + + test('reports a stop as stopped rather than as a failed run', async () => { + const cell = createAgentCell(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + runTokenSource.cancel(); + await context.onAgentEvent?.({ type: 'text_delta', text: 'first' }); + + return { finalOutput: '' }; + }); + + await executeAgentCell(cell, mockController, encryptedStorage, runTokenSource.token, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(mockExecution.end.firstCall.args[0]).to.be.false; + + const [outputs] = mockExecution.appendOutput.firstCall.args as [NotebookCellOutput[]]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('Stopped'); + expect(text).to.not.include('Canceled'); + }); + + test('a run that is never stopped still completes', async () => { + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, encryptedStorage, runTokenSource.token, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(mockExecution.end.firstCall.args[0]).to.be.true; + }); + }); }); suite('removeEphemeralCellsForAgentBlocks', () => { diff --git a/src/notebooks/deepnote/snapshots/snapshotService.ts b/src/notebooks/deepnote/snapshots/snapshotService.ts index d39603d6b7..ce64336d8d 100644 --- a/src/notebooks/deepnote/snapshots/snapshotService.ts +++ b/src/notebooks/deepnote/snapshots/snapshotService.ts @@ -118,6 +118,7 @@ function generateTimestamp(): string { @injectable() export class SnapshotService implements ISnapshotMetadataService, IExtensionSyncActivationService { private readonly converter = new DeepnoteDataConverter(); + private readonly endedExecutionSessions = new Set(); private readonly environmentStates = new Map(); private readonly fileWrittenCallbacks: ((uri: Uri) => void)[] = []; private readonly pendingSnapshotSaves = new Map< @@ -193,11 +194,21 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync this, this.disposables ); + + notebookCellExecutions.onDidStartQueueExecution( + (e) => this.retireFinishedExecutionSession(e.notebookUri), + this, + this.disposables + ); } async captureEnvironmentBeforeExecution(notebookUri: string): Promise { logger.info(`[Snapshot] captureEnvironmentBeforeExecution called for ${notebookUri}`); + // Covers queues opened outside a controller run — the interactive window, execution resumed + // after a reload — which never signal a run start. + this.retireFinishedExecutionSession(notebookUri); + // Seed the session start at capture time so `startedAt` reflects capture, not the first cell. this.tracker.ensureExecutionState(notebookUri, Date.now()); @@ -216,6 +227,7 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync } clearExecutionState(notebookUri: string): void { + this.endedExecutionSessions.delete(notebookUri); this.tracker.clear(notebookUri); this.environmentStates.delete(notebookUri); @@ -701,6 +713,10 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync private async onExecutionComplete(notebookUri: string): Promise { logger.debug(`[Snapshot] onExecutionComplete called for ${notebookUri}`); + // The run is over, but its metadata stays readable — the deferred save and any file save that + // follows still serialize it. The next run's first queue resets it. + this.endedExecutionSessions.add(notebookUri); + // Wait for any pending cell state change events to be processed. // This is needed because the queue completion event can fire before the // last cell's Idle state change event has been processed (race condition). @@ -836,8 +852,15 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync } catch (error) { // Fire-and-forget save: swallow so a failure never becomes an unhandled rejection. logger.error(`[Snapshot] Failed to save deferred snapshot for ${notebookUri}`, error); - } finally { - // Clear execution state so the next run starts fresh, even if the save above failed. + } + } + + /** + * Drops the previous run's metadata once a new run begins. Deferred until then so the save that + * follows a run — and any file save after it — still serializes what that run did. + */ + private retireFinishedExecutionSession(notebookUri: string): void { + if (this.endedExecutionSessions.delete(notebookUri)) { this.clearExecutionState(notebookUri); } } diff --git a/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts b/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts index b817d8b229..c64128a94d 100644 --- a/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts +++ b/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts @@ -1148,6 +1148,7 @@ project: suite('deferred snapshot save timing', () => { const notebookUri = activatedServiceNotebookUri; let clock: fakeTimers.InstalledClock; + let activatedService: SnapshotService; let changeEmitter: EventEmitter; let closeEmitter: EventEmitter; let flush: sinon.SinonStub; @@ -1158,6 +1159,7 @@ project: clock = fakeTimers.install(); const built = buildActivatedSnapshotService(); + activatedService = built.service; changeEmitter = built.changeEmitter; closeEmitter = built.closeEmitter; @@ -1278,6 +1280,90 @@ project: await clock.tickAsync(3000); assert.isFalse(flush.called, 'an output change with no pending save must not arm a deferred save'); }); + + test('keeps the run metadata readable after the deferred save flushes (catches wiping the state the next file save serializes)', async () => { + await arm(); + await clock.tickAsync(150); + + assert.isTrue(flush.calledOnce, 'the deferred save must have flushed'); + // deepnoteSerializer reads this on every save; wiping it at flush time meant a Ctrl+S a + // moment later wrote the .deepnote file with no execution metadata at all. + assert.isDefined(activatedService.getExecutionMetadata(notebookUri)); + }); + + test('a run that opens no kernel queue still clears the finished run (an agent run generating no cells)', async () => { + await arm(); + await clock.tickAsync(150); + + // Nothing calls captureEnvironmentBeforeExecution here: an agent cell runs off the kernel, + // so without the run-start signal the finished run's counters would be what gets saved. + notebookCellExecutions.notifyQueueStart(notebookUri); + + assert.isUndefined( + activatedService.getExecutionMetadata(notebookUri), + "a run that executes nothing must not report the previous run's counters" + ); + }); + + test('a run start with no completion before it keeps the run alive', async () => { + notebookCellExecutions.notifyQueueStart(notebookUri); + + assert.strictEqual( + activatedService.getExecutionMetadata(notebookUri)?.summary?.blocksExecuted, + 1, + 'only a finished run may be retired' + ); + }); + + test("the next run's first queue clears the finished run (catches counters accumulating across runs)", async () => { + await arm(); + await clock.tickAsync(150); + + // A queue opening after the completion belongs to the next run. + await activatedService.captureEnvironmentBeforeExecution(notebookUri); + + assert.isUndefined( + activatedService.getExecutionMetadata(notebookUri), + "a new run must not inherit the previous run's counters" + ); + }); + + test('a second queue inside the same run keeps the run alive (catches clearing per queue, which an agent run opens one of per generated cell)', async () => { + // No completion since the fixture recorded its executed cell: still the same run. + await activatedService.captureEnvironmentBeforeExecution(notebookUri); + + assert.strictEqual( + activatedService.getExecutionMetadata(notebookUri)?.summary?.blocksExecuted, + 1, + 'a mid-run queue must not reset the session' + ); + }); + + test("the run's own captured environment survives its first executing cell (catches clearing on Executing, which lands after capture)", async () => { + const capturedEnvironment: Environment = { + hash: 'sha256:abc', + packages: {}, + platform: 'linux-x64', + python: { environment: 'venv', version: '3.12.0' } + }; + when(mockEnvironmentCapture.captureEnvironment(anything())).thenResolve(capturedEnvironment); + + await arm(); + await clock.tickAsync(150); + + await activatedService.captureEnvironmentBeforeExecution(notebookUri); + + const cellNotebook = mock(); + when(cellNotebook.uri).thenReturn(Uri.parse(notebookUri)); + + const cell = mock(); + when(cell.notebook).thenReturn(instance(cellNotebook)); + when(cell.metadata).thenReturn({ id: 'cell-1' }); + + notebookCellExecutions.changeCellState(instance(cell), NotebookCellExecutionState.Executing); + + assert.deepStrictEqual(await activatedService.getEnvironmentMetadata(notebookUri), capturedEnvironment); + }); }); suite('createSnapshot', () => { diff --git a/src/platform/notebooks/cellExecutionStateService.ts b/src/platform/notebooks/cellExecutionStateService.ts index 25b25dd376..c0800b72e2 100644 --- a/src/platform/notebooks/cellExecutionStateService.ts +++ b/src/platform/notebooks/cellExecutionStateService.ts @@ -57,6 +57,7 @@ const STATE_NAMES: Record = { export namespace notebookCellExecutions { const eventEmitter = trackDisposable(new EventEmitter()); const queueCompletionEmitter = trackDisposable(new EventEmitter()); + const queueStartEmitter = trackDisposable(new EventEmitter()); /** * An {@link Event} which fires when the execution state of a cell has changed. @@ -71,6 +72,13 @@ export namespace notebookCellExecutions { */ export const onDidCompleteQueueExecution = queueCompletionEmitter.event; + /** + * An {@link Event} which fires when a user-initiated run begins, before any cell executes. + * A run that executes no cells at all still fires it, which is what separates it from the + * first cell going Executing. + */ + export const onDidStartQueueExecution = queueStartEmitter.event; + /** * Notify listeners that a notebook's cell execution queue has completed. * @param notebookUri The URI of the notebook whose queue completed @@ -80,6 +88,15 @@ export namespace notebookCellExecutions { queueCompletionEmitter.fire({ notebookUri }); } + /** + * Notify listeners that a user-initiated run is starting. + * @param notebookUri The URI of the notebook whose run is starting + */ + export function notifyQueueStart(notebookUri: string) { + logger.debug(`[CellExecState] Queue execution starting for ${notebookUri}`); + queueStartEmitter.fire({ notebookUri }); + } + export function changeCellState(cell: NotebookCell, state: NotebookCellExecutionState, executionOrder?: number) { const cellId = cell.metadata?.id as string | undefined; const stateName = STATE_NAMES[state] || String(state); From bd7e29c04b5e80b4e4c067df2e4b7b69b4fb9ce4 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 13 Aug 2026 07:58:26 +0000 Subject: [PATCH 71/80] fix(agent-block): stop the batch when the agent itself is interrupted executeAgentCell reports a stop by ending its cell and returning, not by throwing, so a run interrupted during the agent cell reached the loop looking like one that finished and the cells after it still executed. The batch already aborts when a kernel segment is interrupted; this is the one branch that did not, because it was the one that does not throw. Reported by CodeRabbit on #358. Co-Authored-By: Claude Opus 4.6 Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- .../controllers/vscodeNotebookController.ts | 6 +++ .../vscodeNotebookController.unit.test.ts | 39 +++++++++++++++++++ 2 files changed, 45 insertions(+) diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index c127dbfc18..6abf173d05 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -694,6 +694,12 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont this.serviceContainer.get(IEncryptedStorage), agentCancellation.token ).catch(noop); + + // A stopped agent ends its own cell and returns, so it arrives here looking like a + // run that finished. Without this the cells after it would still execute. + if (agentCancellation.token.isCancellationRequested) { + throw new CancellationError(); + } } finally { this.agentCancellations.delete(doc); agentCancellation.dispose(); diff --git a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts index c7cbdd2ffe..9a4019f78f 100644 --- a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts +++ b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts @@ -1080,6 +1080,45 @@ suite(`Notebook Controller`, function () { ); }); + test('an interrupt during the agent cell stops the trailing segment', async function () { + // Catches: reading the stop from a rejection - executeAgentCell reports one by ending its + // cell and returning, so the batch sees a run that finished normally. + const { + notebook, + cells: [agentCell, trailingCell] + } = createMockNotebookWithCells([ + { metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, text: 'Test prompt' }, + { metadata: { id: 'code-1' }, text: 'print(1)' } + ]); + + const executedIndexes: number[] = []; + stubKernelForExecution({ + failed: false, + executeCell: async (cell: NotebookCell) => { + executedIndexes.push(cell.index); + } + }); + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(undefined); + + // The agent's execution object is created as its run starts, which is where Stop lands. + createNotebookCellExecutionStub.callsFake((cell: NotebookCell) => { + if (cell === agentCell) { + vscodeController.controller.interruptHandler!(notebook); + } + mockExecution.end = sinon.stub(); + + return mockExecution; + }); + + await vscodeController.controller.executeHandler( + [agentCell, trailingCell], + notebook, + vscodeController.controller + ); + + assert.deepStrictEqual(executedIndexes, [], 'the trailing segment must not run after an interrupt'); + }); + test('a clean kernel segment still runs the agent cell and the trailing segment', async function () { // Catches: aborting the batch when nothing failed (e.g. consulting the queue verdict too early). const { From 8bd754751fafaa49ad5586d2338294f5566a2f8c Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 13 Aug 2026 11:47:59 +0000 Subject: [PATCH 72/80] fix(tests): enhance notebook controller tests and update deepnote file watcher - Added logging for error handling in the notebook controller's interrupt handler to ensure proper error reporting when interrupting notebook execution. - Updated comments in the agent cell execution handler for clarity on tool failure handling. - Corrected content hash and spelling in deepnote file change watcher tests to maintain consistency and accuracy. These changes improve the robustness of the tests and clarify the code's intent. --- .../controllers/vscodeNotebookController.unit.test.ts | 11 ++++++++++- src/notebooks/deepnote/agentCellExecutionHandler.ts | 2 +- .../deepnote/deepnoteFileChangeWatcher.unit.test.ts | 6 +++--- 3 files changed, 14 insertions(+), 5 deletions(-) diff --git a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts index 9a4019f78f..412df10542 100644 --- a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts +++ b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts @@ -59,6 +59,7 @@ import { Environment, PythonExtension } from '@vscode/python-extension'; import { crateMockedPythonApi, whenResolveEnvironment } from '../../kernels/helpers.unit.test'; import { IJupyterVariablesProvider } from '../../kernels/variables/types'; import { notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; +import { logger } from '../../platform/logging'; import { createMockNotebookWithCells } from '../deepnote/deepnoteTestHelpers'; // executeAgentCell takes IEncryptedStorage from the controller's container; getProjectAgentContext @@ -1100,10 +1101,18 @@ suite(`Notebook Controller`, function () { }); when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(undefined); + const { interruptHandler } = vscodeController.controller; + + if (!interruptHandler) { + assert.fail('the controller must install an interrupt handler for Stop to reach the agent'); + } + // The agent's execution object is created as its run starts, which is where Stop lands. createNotebookCellExecutionStub.callsFake((cell: NotebookCell) => { if (cell === agentCell) { - vscodeController.controller.interruptHandler!(notebook); + Promise.resolve(interruptHandler(notebook)).catch((ex) => + logger.error('Failed to interrupt the notebook', ex) + ); } mockExecution.end = sinon.stub(); diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index d98c35f331..2914cefcfc 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -224,7 +224,7 @@ export async function executeAgentCell( ...getProjectAgentContext(cell.notebook), notebookContext, // The guards sit outside the `try`s: those turn every throw into a string the model reads - // as a retryable tool failure, which is how a stop used to make the agent do more work. + // as a tool failure worth retrying, which is how a stop used to make the agent do more work. addMarkdownBlock: async ({ content }: { content: string }) => { Cancellation.throwIfCanceled(token); diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index adbbca7dd2..18f92991ff 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -175,7 +175,7 @@ project: type: code sortingKey: '000000' blockGroup: '1' - contentHash: 'sha256:beefbeef' + contentHash: 'sha256:0badc0de1' content: print("hello") `; @@ -194,7 +194,7 @@ project: type: agent sortingKey: '000000' blockGroup: '1' - content: summarise the dataframe + content: summarize the dataframe metadata: deepnote_agent_model: gpt-5 `; @@ -558,7 +558,7 @@ project: }, outputs: [], kind: NotebookCellKind.Code, - document: { getText: () => 'summarise the dataframe', languageId: 'plaintext' } + document: { getText: () => 'summarize the dataframe', languageId: 'plaintext' } } ] }); From ecf533eb752b6200f78eeb1b2500dd6b5c83298c Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 13 Aug 2026 13:41:04 +0000 Subject: [PATCH 73/80] Remove unnecessary code comment --- src/notebooks/deepnote/deepnoteNotebookCommandListener.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts index c01c73d0c5..e65717a4e0 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts @@ -663,8 +663,6 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation } private trackAddBlock(blockType: string): void { - // Commands only ever insert blocks the user asked for; agent scratch cells are counted - // in DeepnoteCellExecutionAnalytics, which is the only observer that sees them. this.analytics.trackEvent({ eventName: 'add_block', properties: { blockType, isEphemeral: false } }); } } From 689e1e3243a2e1c0d059e67b216edc23e9df68bb Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 13 Aug 2026 14:10:32 +0000 Subject: [PATCH 74/80] fix(build): keep runtime-core out of the web module graph Marking @deepnote/runtime-core external for the web target left a bare top-level import in extension.web.bundle.js: agentCellExecutionHandler imports it statically, and the web-registered VSCodeNotebookController pulls that handler in through controllerRegistration. runtime-core needs Node built-ins (net, child_process) and .vscodeignore excludes node_modules from the VSIX, so the specifier can never resolve at runtime -- and dropping the external turns it into a build failure (tcp-port-used and @ai-sdk/mcp reach for net/child_process), which is what the external was actually silencing rather than fixing. Alias it to a stub instead, the same way @nteract/presentational-components is already aliased in this file. Agent blocks are desktop-only; the web build now throws a clear error if either export is ever called instead of shipping an unresolvable import. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- build/esbuild/build.ts | 14 ++++++++++++-- src/notebooks/deepnote/runtimeCore.web.ts | 7 +++++++ 2 files changed, 19 insertions(+), 2 deletions(-) create mode 100644 src/notebooks/deepnote/runtimeCore.web.ts diff --git a/build/esbuild/build.ts b/build/esbuild/build.ts index 54ae574977..9816dc8cd5 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' // Node built-ins (net, child_process); agent blocks run on desktop only + 'mathjax-electron' // Uses Node.js path module, MathJax rendering handled differently in browser ]; const desktopExternals = [...commonExternals, ...deskTopNodeModulesToExternalize]; const bundleConfig = getBundleConfiguration(); @@ -275,6 +274,17 @@ function createConfig( if (target === 'desktop') { alias['jsonc-parser'] = path.join(extensionFolder, 'node_modules', 'jsonc-parser', 'lib', 'esm', 'main.js'); } + // @deepnote/runtime-core needs Node built-ins (net, child_process) and is excluded from the VSIX; + // externalizing it (like desktop) would leave an unresolvable bare import in the web bundle. + if (target === 'web') { + alias['@deepnote/runtime-core'] = path.join( + extensionFolder, + 'src', + 'notebooks', + 'deepnote', + 'runtimeCore.web.ts' + ); + } // Desktop builds use CommonJS for VS Code/Cursor compatibility // Web builds use ESM for browser compatibility const config: SameShape = { diff --git a/src/notebooks/deepnote/runtimeCore.web.ts b/src/notebooks/deepnote/runtimeCore.web.ts new file mode 100644 index 0000000000..01cc943a8f --- /dev/null +++ b/src/notebooks/deepnote/runtimeCore.web.ts @@ -0,0 +1,7 @@ +/** Web stands in for @deepnote/runtime-core, which needs Node built-ins. Agent blocks run on desktop only. */ +const unsupported = (): never => { + throw new Error('Deepnote agent blocks are not supported in the web extension host.'); +}; + +export const executeAgentBlock = unsupported; +export const serializeNotebookContextFromBlocks = unsupported; From 5c9fc6377dae0877a15e669ccfbd9130d4e092c7 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 13 Aug 2026 14:10:41 +0000 Subject: [PATCH 75/80] fix(agent-block): feed agent cell execution into snapshot accounting executeAgentCell called controller.createNotebookCellExecution directly and never touched the internal execution-state shim that SnapshotService and execute_cell analytics actually listen to -- start()/end() on a raw NotebookCellExecution fires no event either one sees. The agent cell still counts toward totalCodeCells since it's Code-kind, so a Run All containing an agent block could never make executedBlockCount equal totalCodeCells and always fell back to updating the latest snapshot only, silently losing timestamped history for every run of the PR's headline feature. The agent block also never got execution timing on save, and never showed up in execute_cell analytics. Route start/end through notebookCellExecutions.changeCellState so the run is visible on the same shim every kernel execution reports to. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- .../deepnote/agentCellExecutionHandler.ts | 15 ++++- .../agentCellExecutionHandler.unit.test.ts | 42 ++++++++++++++ .../snapshots/snapshotService.unit.test.ts | 58 +++++++++++++++++++ 3 files changed, 112 insertions(+), 3 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index 2914cefcfc..3a4c842ead 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -180,7 +180,16 @@ export async function executeAgentCell( ): Promise { const executeAgentBlockFn = options?.executeAgentBlockFn ?? executeAgentBlock; const execution = controller.createNotebookCellExecution(cell); + + // The agent runs off the kernel, so nothing announces it on the internal shim — the only source + // SnapshotService and the execute_cell analytics read. Without this the run is invisible to both. + const endExecution = (success: boolean) => { + execution.end(success, Date.now()); + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Idle); + }; + execution.start(Date.now()); + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Executing); try { await execution.clearOutput(); @@ -304,7 +313,7 @@ export async function executeAgentCell( const result = await executeAgentBlockFn(agentBlock, context); logger.info(`Agent cell: executeAgentBlock completed, finalOutput length=${result.finalOutput.length}`); - execution.end(true, Date.now()); + endExecution(true); } catch (error) { if (isStopped(error)) { logger.info('Agent cell execution stopped'); @@ -312,7 +321,7 @@ export async function executeAgentCell( const stoppedOutput = new NotebookCellOutput([NotebookCellOutputItem.stderr('[Agent] Stopped')]); await execution.appendOutput([stoppedOutput]).then(undefined, () => undefined); - execution.end(false, Date.now()); + endExecution(false); return; } @@ -331,7 +340,7 @@ export async function executeAgentCell( const message = error instanceof Error ? error.message : String(error); const stderrOutput = new NotebookCellOutput([NotebookCellOutputItem.stderr(message)]); await execution.appendOutput([stderrOutput]).then(undefined, () => undefined); - execution.end(false, Date.now()); + endExecution(false); } } diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index 1e8f72652e..fe18fe187f 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -264,6 +264,48 @@ suite('AgentCellExecutionHandler', () => { expect(mockExecution.end.firstCall.args[0]).to.be.true; }); + // SnapshotService and execute_cell analytics read this shim, not the raw NotebookCellExecution — + // without these events a Run All containing an agent block never matches its own code-cell count. + test('reports the run on the execution-state shim so SnapshotService can see it', async () => { + const cell = createAgentCell('Analyze data'); + const seenStates: NotebookCellExecutionState[] = []; + + disposables.push( + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell === cell) { + seenStates.push(e.state); + } + }) + ); + + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(seenStates).to.deep.equal([NotebookCellExecutionState.Executing, NotebookCellExecutionState.Idle]); + }); + + test('reports Idle on the shim even when the run fails', async () => { + mockExecution.clearOutput.rejects(new Error('Something went wrong')); + + const cell = createAgentCell(); + const seenStates: NotebookCellExecutionState[] = []; + + disposables.push( + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell === cell) { + seenStates.push(e.state); + } + }) + ); + + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(seenStates).to.deep.equal([NotebookCellExecutionState.Executing, NotebookCellExecutionState.Idle]); + }); + // Incremental deltas only — full transcript per event is O(n²) over the EH boundary. test('streams text_delta events via appendOutputItems with incremental stdout chunks', async () => { executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { diff --git a/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts b/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts index c64128a94d..0a042032bc 100644 --- a/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts +++ b/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts @@ -1934,6 +1934,64 @@ project: ); }); + // An agent block is a Code-kind cell (agentBlockConverter) that runs off the kernel, so the + // tracker never sees it — this is the shape executeAgentCell left behind before it started + // reporting to the execution-state shim. + test('an untracked Code cell (an agent block) blocks Run-All even though every kernel cell executed', async () => { + const mockConfig = mock(); + when(mockConfig.get('snapshots.enabled', true)).thenReturn(true); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + + const projectId = 'test-project-id'; + const notebookId = 'test-notebook-id'; + + const mockNotebook = mockNotebookDoc({ + uri: Uri.parse(notebookUri), + projectId, + notebookId, + cells: [mockCell({ id: 'cell-1', source: 'print(1)' }), mockCell({ id: 'agent-cell' })] + }); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([mockNotebook]); + + const originalProject: DeepnoteFile = { + metadata: { createdAt: '2025-01-01T00:00:00Z' }, + version: '1.0.0', + project: { + id: projectId, + name: 'Test Project', + notebooks: [{ id: notebookId, name: 'Test Notebook', blocks: [] }] + } + }; + const mockNotebookManager = mock(); + when(mockNotebookManager.getProjectForNotebook(projectId, notebookId)).thenReturn(originalProject); + + const testService = new SnapshotService( + instance(mockEnvironmentCapture), + mockDisposables, + instance(mockNotebookManager), + tracker + ); + + // Only cell-1 is tracked as executed — agent-cell never reaches recordCellExecutionStart/End. + const startTime = Date.now(); + tracker.recordCellExecutionStart(notebookUri, 'cell-1', startTime); + tracker.recordCellExecutionEnd(notebookUri, 'cell-1', startTime + 100, true); + + const writtenUris = captureSnapshotWrites(); + + testService.activate(); + await flushDeferredSave(notebookUri); + + assert.isFalse( + wroteTimestampedSnapshot(writtenUris), + 'a Code-kind cell the tracker never saw must keep the run off the Run-All branch' + ); + assert.isTrue( + wroteLatestSnapshot(writtenUris), + 'the run instead falls back to the partial-run (latest-only) path' + ); + }); + test('writes the snapshot next to the saved notebook, not a sibling that shares the project id', async () => { const mockConfig = mock(); when(mockConfig.get('snapshots.enabled', true)).thenReturn(true); From 1ef480008b7c5881f7798587099655ba249d6d51 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 13 Aug 2026 14:10:51 +0000 Subject: [PATCH 76/80] fix(kernels): notify queue completion after a resumed cell execution notifyQueueComplete now has a single production caller, the controller's executeQueuedCells, after the per-queue notification moved out of CellExecutionQueue to stop an agent batch's own per-segment queues from retiring the run mid-batch. NotebookKernelExecution.resumeCellExecution opens a queue through the same path but never goes through the controller's batch, so SnapshotService starts tracking a resumed execution and never sees it finish -- its counters and startedAt survive into whatever runs next on that document. restoreConnection is reachable only for Jupyter/interactive documents (a .deepnote file cannot take that path), so this doesn't affect Deepnote snapshots today, but a resumed queue should still announce its own completion. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- src/kernels/kernelExecution.ts | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/kernels/kernelExecution.ts b/src/kernels/kernelExecution.ts index 8d15d6599b..aeca91137d 100644 --- a/src/kernels/kernelExecution.ts +++ b/src/kernels/kernelExecution.ts @@ -161,6 +161,8 @@ export class NotebookKernelExecution implements INotebookKernelExecution { .then(() => true) .catch(() => false); + notebookCellExecutions.notifyQueueComplete(cell.notebook.uri.toString()); + traceCellMessage( cell, `NotebookKernelExecution.resumeCellExecution (completed), ${getDisplayPath(cell.notebook.uri)}` From b3947c716f0970b3d90fa38dc7d5544a26ba2aac Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 13 Aug 2026 14:26:26 +0000 Subject: [PATCH 77/80] Remove unnecessary comments --- src/platform/analytics/types.ts | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/platform/analytics/types.ts b/src/platform/analytics/types.ts index f63bcb7665..8428043247 100644 --- a/src/platform/analytics/types.ts +++ b/src/platform/analytics/types.ts @@ -33,7 +33,6 @@ export type CommandOutcome = 'completed' | 'cancelled' | 'failed'; /** Caller-supplied properties per event; `undefined` means none beyond the common properties the service attaches. */ export interface TelemetryEventProperties { - /** `isEphemeral` marks agent scratch blocks, which no add-block command reports. */ add_block: { blockType: string; isEphemeral: boolean }; authenticate_integration: { integrationType: string; outcome: CommandOutcome }; configure_integration: { integrationType: string }; @@ -45,7 +44,6 @@ export interface TelemetryEventProperties { delete_integration: { integrationType: string }; delete_notebook: { outcome: CommandOutcome }; duplicate_notebook: { outcome: CommandOutcome }; - /** `isEphemeral` is true for agent-generated cells, which the agent runs itself via `notebook.cell.execute`. */ execute_cell: { cellType: 'sql' | 'markdown' | 'code'; isEphemeral: boolean; integrationType?: string }; execute_notebook: undefined; export_notebook: { outcome: CommandOutcome; format?: string }; From c83c0877b9cd38dc80e5a7aab1bb252b329a05ac Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 17 Aug 2026 07:58:17 +0000 Subject: [PATCH 78/80] test(agent-block): cover the post-review fixes in E2E The agent block E2E last changed four days before the review, so none of the fixes it prompted had integration coverage: a mixed Run All that must stop, a Stop that must reach the agent, and a transcript that must survive the save. Four tests, in two groups that each bind the notebook they run: - a failing cell before the agent ends the batch, so neither the agent nor the trailing cell runs. The failure comes from the fixture rather than the agent, which is the reported repro; the mock is still scripted so a batch that carried on has markers to render, and those are asserted absent. - Interrupt during the agent run stops it and the cell after it. The generated cell prints and then sleeps, giving a bounded window in which the notebook is demonstrably running. It clicks "Interrupt" (notebook.interruptExecution) specifically -- VS Code shows that only while notebookInterruptibleKernel is set, and it is the one toolbar action reaching the controller's interruptHandler. "Stop Execution" cancels the cells without telling the agent, so it is not a fallback. - a generated cell that raises comes back to the agent as "Execution failed:" and the run carries on. The mock runs --strict and the second leg matches on that prefix, so a swallowed failure leaves the request unmatched and the later markers never render. - the streamed transcript is read back out of the snapshot sidecar, where it lands once outputs are stripped from the main file. Asserted through the parsed block, not the raw YAML: serializeDeepnoteFile folds at 120 columns, and a fold inside a marker makes a raw substring match fail on transcript length alone. The two groups share one workspace and one environment -- a second one costs about 90s of CI and every test wants the same kernel -- but nothing else. Reopening the notebook drops the block's generated cells, so that happens once per group rather than between tests, and the pre-existing serial pair keeps working. Written and typechecked; not yet run against a workbench. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- test/e2e/fixtures/agent-block-batch.deepnote | 36 + test/e2e/fixtures/agent-block-stop.deepnote | 29 + test/e2e/helpers/mockOpenAiServer.ts | 28 + test/e2e/helpers/notebook.ts | 96 ++- test/e2e/helpers/yaml.ts | 30 +- test/e2e/suite/agentBlock.e2e.test.ts | 667 +++++++++++++------ 6 files changed, 687 insertions(+), 199 deletions(-) create mode 100644 test/e2e/fixtures/agent-block-batch.deepnote create mode 100644 test/e2e/fixtures/agent-block-stop.deepnote diff --git a/test/e2e/fixtures/agent-block-batch.deepnote b/test/e2e/fixtures/agent-block-batch.deepnote new file mode 100644 index 0000000000..5d426fdc42 --- /dev/null +++ b/test/e2e/fixtures/agent-block-batch.deepnote @@ -0,0 +1,36 @@ +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00.000Z' + modifiedAt: '2025-01-01T00:00:00.000Z' +project: + id: e2e-agent-batch-project + name: E2E Agent Batch + notebooks: + - id: e2e-agent-batch-notebook + name: Agent Batch + blocks: + - id: e2e-agent-batch-failing-block + blockGroup: e2e-agent-batch-group + type: code + content: |- + raise ValueError("e2e-batch-boom") + sortingKey: a0 + metadata: {} + - id: e2e-agent-batch-agent-block + blockGroup: e2e-agent-batch-group + type: agent + content: |- + Run some Python, then add a markdown block summarizing this notebook. + sortingKey: a1 + metadata: + deepnote_agent_model: 'gpt-5.6-sol' + - id: e2e-agent-batch-trailing-block + blockGroup: e2e-agent-batch-group + type: code + content: |- + print("e2e-batch-trailing-ran") + sortingKey: a2 + metadata: {} + executionMode: block + isModule: false + settings: {} diff --git a/test/e2e/fixtures/agent-block-stop.deepnote b/test/e2e/fixtures/agent-block-stop.deepnote new file mode 100644 index 0000000000..a8a2177794 --- /dev/null +++ b/test/e2e/fixtures/agent-block-stop.deepnote @@ -0,0 +1,29 @@ +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00.000Z' + modifiedAt: '2025-01-01T00:00:00.000Z' +project: + id: e2e-agent-stop-project + name: E2E Agent Stop + notebooks: + - id: e2e-agent-stop-notebook + name: Agent Stop + blocks: + - id: e2e-agent-stop-agent-block + blockGroup: e2e-agent-stop-group + type: agent + content: |- + Run some Python, then add a markdown block summarizing this notebook. + sortingKey: a0 + metadata: + deepnote_agent_model: 'gpt-5.6-sol' + - id: e2e-agent-stop-trailing-block + blockGroup: e2e-agent-stop-group + type: code + content: |- + print("e2e-stop-trailing-ran") + sortingKey: a1 + metadata: {} + executionMode: block + isModule: false + settings: {} diff --git a/test/e2e/helpers/mockOpenAiServer.ts b/test/e2e/helpers/mockOpenAiServer.ts index bb718dc284..f9fb50cf6e 100644 --- a/test/e2e/helpers/mockOpenAiServer.ts +++ b/test/e2e/helpers/mockOpenAiServer.ts @@ -4,6 +4,10 @@ import { connect } from 'net'; import * as os from 'os'; import * as path from 'path'; import { setTimeout as delay } from 'timers/promises'; +import { InputBox, Workbench } from 'vscode-extension-tester'; + +import { QUICK_PICK_TIMEOUT } from './constants'; +import { waitForNotification } from './notifications'; // npx aimock — keep jest/vitest peers out of the lockfile. const AIMOCK_VERSION = '1.37.4'; @@ -17,6 +21,30 @@ export function pointExtensionHostAtMockServer(): void { process.env.OPENAI_BASE_URL = `http://127.0.0.1:${MOCK_OPENAI_PORT}/v1`; } +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'; +const API_KEY_SAVED_NOTIFICATION = /OpenAI API key has been saved/; + +/** + * Stores a throwaway key so the agent has credentials to send. The request goes to the mock, which + * never checks it — `pointExtensionHostAtMockServer` is what keeps it off the real API. + */ +export 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(); + + await waitForNotification(API_KEY_SAVED_NOTIFICATION, QUICK_PICK_TIMEOUT, true); +} + +/** Removes the stored key. The key outlives a suite, so every suite that stores one clears it. */ +export async function clearStoredOpenAiApiKey(): Promise { + await new Workbench().executeCommand(CLEAR_API_KEY_COMMAND); +} + const START_TIMEOUT = 90_000; const POLL_INTERVAL = 200; const STOP_TIMEOUT = 2_000; diff --git a/test/e2e/helpers/notebook.ts b/test/e2e/helpers/notebook.ts index 1945608472..36784deb5c 100644 --- a/test/e2e/helpers/notebook.ts +++ b/test/e2e/helpers/notebook.ts @@ -4,12 +4,12 @@ import { OUTPUT_FRAME_SWITCH_TIMEOUT, OUTPUT_POLL_INTERVAL, OUTPUT_SELECTOR, WOR import { dismissAllNotifications } from './notifications'; /** - * Focuses the given notebook editor and clicks its toolbar "Run All" button. The command-palette - * entry for `deepnote.runallcells` ("Jupyter: Run All Cells") is gated behind context keys - * (`deepnote.ispythonornativeactive`, …) that are not reliably set under automation, so driving it + * Focuses the given notebook editor and clicks the toolbar button carrying `ariaLabel`. The + * command-palette entries for these actions are gated behind context keys + * (`deepnote.ispythonornativeactive`, …) that are not reliably set under automation, so driving them * through `Workbench.executeCommand` can silently miss and trigger the wrong command. */ -export async function clickRunAll(notebookFileName: string): Promise { +async function clickNotebookToolbarButton(notebookFileName: string, ariaLabel: string): Promise { const driver = VSBrowser.instance.driver; await new EditorView().openEditor(notebookFileName); @@ -17,12 +17,11 @@ export async function clickRunAll(notebookFileName: string): Promise { // Locate AND click inside the same wait loop. The notebook toolbar can re-render between finding // the button and clicking it (the editor re-focuses, kernel status / notifications change), which // would otherwise surface as a StaleElementReferenceError. Re-finding and clicking on the next - // tick is still a SINGLE "Run All" — the run is only issued once the click actually lands, so - // this does not re-run a notebook whose first execution was accepted. + // tick still issues the action only ONCE — it is only issued when the click actually lands. await driver.wait( async () => { try { - const [button] = await driver.findElements(By.css('a.action-label[aria-label="Run All"]')); + const [button] = await driver.findElements(By.css(`a.action-label[aria-label="${ariaLabel}"]`)); if (!button) { return false; } @@ -31,16 +30,33 @@ export async function clickRunAll(notebookFileName: string): Promise { return true; } catch (error) { - console.warn('[deepnote-e2e] locate/click notebook Run All (retrying):', error); + console.warn(`[deepnote-e2e] locate/click notebook "${ariaLabel}" (retrying):`, error); return false; } }, WORKBENCH_TIMEOUT, - 'notebook "Run All" button did not appear or could not be clicked' + `notebook "${ariaLabel}" button did not appear or could not be clicked` ); } +export async function clickRunAll(notebookFileName: string): Promise { + return clickNotebookToolbarButton(notebookFileName, 'Run All'); +} + +/** + * Clicks the toolbar's "Interrupt" button — VS Code's `notebook.interruptExecution`, shown while + * `notebookHasSomethingRunning && notebookInterruptibleKernel`. + * + * It is the only toolbar action that reaches the controller's `interruptHandler`, and therefore the + * only one that signals a running agent to stop. "Stop Execution" (`notebook.cancelExecution`, + * which VS Code shows in its place for a kernel that declares no interrupt handler) cancels the + * cells without ever telling the agent, so it must not stand in as a fallback here. + */ +export async function clickInterrupt(notebookFileName: string): Promise { + return clickNotebookToolbarButton(notebookFileName, 'Interrupt'); +} + /** * Clicks the notebook cell status bar item whose text contains `label`. Cell chrome lives in the * main window DOM (not the output iframe), so this switches out of the webview first and matches on @@ -138,6 +154,68 @@ export async function readRenderedOutput(): Promise { }); } +/** + * Polls the notebook webview until every marker in `markers` is rendered and none of `absentMarkers` + * is, then returns the text it settled on. `context` names the state being waited for; it is only + * used to make the timeout message say what did not happen. + */ +export async function awaitWebviewMarkers( + markers: string[], + timeout: number, + context: string, + absentMarkers: string[] = [] +): Promise { + const driver = VSBrowser.instance.driver; + const deadline = Date.now() + timeout; + let text = ''; + + while (Date.now() < deadline) { + text = await readNotebookWebviewText(); + const missing = markers.filter((marker) => !text.includes(marker)); + const lingering = absentMarkers.filter((marker) => text.includes(marker)); + if (missing.length === 0 && lingering.length === 0) { + return text; + } + + await driver.sleep(OUTPUT_POLL_INTERVAL); + } + + const missing = markers.filter((marker) => !text.includes(marker)); + const lingering = absentMarkers.filter((marker) => text.includes(marker)); + throw new Error( + `Timed out after ${timeout}ms waiting for notebook webview (${context}). Missing: ${JSON.stringify( + missing + )}. ` + `Lingering: ${JSON.stringify(lingering)}. Last text: ${JSON.stringify(text)}` + ); +} + +/** + * Fails if any of `markers` renders in the notebook webview during the next `windowMs`. + * + * Non-occurrence needs a window rather than one read: the regressions this guards render the + * forbidden text a beat *after* the state the test waited for — a batch that should have stopped + * carries on into the agent's round trip to the local mock and then the trailing cell. Size the + * window well above that round trip, since the whole window is spent on every passing run. + */ +export async function assertMarkersStayAbsent(markers: string[], windowMs: number, context: string): Promise { + const driver = VSBrowser.instance.driver; + const deadline = Date.now() + windowMs; + + while (Date.now() < deadline) { + const text = await readNotebookWebviewText(); + const rendered = markers.filter((marker) => text.includes(marker)); + + if (rendered.length > 0) { + throw new Error( + `Notebook webview rendered ${JSON.stringify(rendered)}, which must not appear (${context}). ` + + `Full text: ${JSON.stringify(text)}` + ); + } + + await driver.sleep(OUTPUT_POLL_INTERVAL); + } +} + /** * Issues a SINGLE "Run All" after the kernel has been selected and polls the notebook output webview * until the expected text renders. It deliberately does NOT re-issue "Run All" when output is diff --git a/test/e2e/helpers/yaml.ts b/test/e2e/helpers/yaml.ts index d179bd8b1f..e09f89038e 100644 --- a/test/e2e/helpers/yaml.ts +++ b/test/e2e/helpers/yaml.ts @@ -1,6 +1,34 @@ -import { deserializeDeepnoteFile } from '@deepnote/blocks'; +import { deserializeDeepnoteFile, isExecutableBlock } from '@deepnote/blocks'; /** Counts notebooks in a serialized `.deepnote` file by parsing it with the canonical schema. */ export function notebookCount(yaml: string): number { return deserializeDeepnoteFile(yaml).project.notebooks.length; } + +/** + * The text a block's stream outputs carry in a serialized `.deepnote`, concatenated in order. + * + * Parse rather than search the raw YAML: `serializeDeepnoteFile` folds at 120 columns, so a marker + * that is one unbroken string in the block can sit across two lines in the file. + */ +export function blockStreamOutputText(yaml: string, blockId: string): string { + const block = deserializeDeepnoteFile(yaml) + .project.notebooks.flatMap((notebook) => notebook.blocks ?? []) + .find((candidate) => candidate.id === blockId); + + if (!block) { + throw new Error(`No block ${JSON.stringify(blockId)} in the serialized project.`); + } + + if (!isExecutableBlock(block)) { + throw new Error(`Block ${JSON.stringify(blockId)} is a ${block.type} block, which carries no outputs.`); + } + + return (block.outputs ?? []) + .map((output: { text?: string | string[] }) => { + const text = output.text ?? ''; + + return Array.isArray(text) ? text.join('') : text; + }) + .join(''); +} diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts index 54a3e6d9e4..1c28181e55 100644 --- a/test/e2e/suite/agentBlock.e2e.test.ts +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -1,15 +1,27 @@ -/** Agent block E2E vs local aimock; legs 2–3 advance on real tool results (no live OpenAI). */ - -import { EditorView, InputBox, VSBrowser, WebView, Workbench } from 'vscode-extension-tester'; +/** + * Agent block E2E vs local aimock; legs 2–3 advance on real tool results (no live OpenAI). + * + * Three fixtures share one workspace and one environment — provisioning a second environment costs + * ~90s of CI and every test here wants the same kernel. Only that setup is shared: each group below + * opens and binds the notebook it runs, so the groups are order-independent and either can run on + * its own. + */ + +import * as fs from 'fs'; +import * as path from 'path'; +import { EditorView, VSBrowser, WebView, Workbench } from 'vscode-extension-tester'; import { FIRST_RUN_OUTPUT_TIMEOUT, MockOpenAiServer, - OUTPUT_POLL_INTERVAL, - QUICK_PICK_TIMEOUT, SUITE_TIMEOUT, WORKBENCH_TIMEOUT, + assertMarkersStayAbsent, + awaitWebviewMarkers, + blockStreamOutputText, + clearStoredOpenAiApiKey, clickCellStatusBarItem, + clickInterrupt, clickRunAll, confirmModalDialog, copyFixtureToTempDir, @@ -19,15 +31,18 @@ import { openFolderViaDialog, openWorkspaceFile, pointExtensionHostAtMockServer, - readNotebookWebviewText, selectEnvironmentForNotebook, startMockOpenAiServer, - waitForNotification + storeMockOpenAiApiKey } from '../helpers'; pointExtensionHostAtMockServer(); const AGENT_FILE = 'agent-block.deepnote'; +const AGENT_BLOCK_ID = 'e2e-agent-block'; +// Mixed-batch fixtures: a failing cell before the agent, and a trailing cell after it. +const BATCH_FILE = 'agent-block-batch.deepnote'; +const STOP_FILE = 'agent-block-stop.deepnote'; const CODE_TOOL_NAME = 'add_code_block'; const MARKDOWN_TOOL_NAME = 'add_markdown_block'; // Leg 3 match: agentCellExecutionHandler add_markdown_block tool result. @@ -56,39 +71,66 @@ const CLEAR_EPHEMERAL_BUTTON = 'Clear ephemeral blocks'; const CLEAR_EPHEMERAL_CONFIRM = 'Clear'; const CLEAR_EPHEMERAL_CONFIRM_TEXT = 'ephemeral block'; const CLEAR_EPHEMERAL_TIMEOUT = 30_000; -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'; const REVERT_FILE_COMMAND = 'File: Revert File'; const DISCARD_CHANGES_BUTTON = "Don't Save"; -const API_KEY_SAVED_NOTIFICATION = /OpenAI API key has been saved/; - -async function awaitWebviewMarkers( - markers: string[], - timeout: number, - context: string, - absentMarkers: string[] = [] -): Promise { - const driver = VSBrowser.instance.driver; - const deadline = Date.now() + timeout; - let text = ''; - - while (Date.now() < deadline) { - text = await readNotebookWebviewText(); - const missing = markers.filter((marker) => !text.includes(marker)); - const lingering = absentMarkers.filter((marker) => text.includes(marker)); - if (missing.length === 0 && lingering.length === 0) { - return text; - } - - await driver.sleep(OUTPUT_POLL_INTERVAL); - } +// Fourth run: the generated cell raises, so the agent must be handed the failure and carry on. +const FAILING_PYTHON_MARKER = 'e2e-agent-code-boom'; +const FAILING_GENERATED_PYTHON = `raise ValueError("${FAILING_PYTHON_MARKER}")`; +// addAndExecuteCodeBlock's prefix for a cell that ran and failed (as opposed to one that never ran). +const EXECUTION_FAILED_TEXT = 'Execution failed:'; +const FAILURE_RECOVERY_MARKDOWN = 'Recovered from the failed cell and carried on'; +const FAILURE_FINAL_AGENT_TEXT = 'Reported the failure as a markdown block.'; +// Fifth run: read back from the snapshot sidecar rather than the webview. +const PERSISTED_PYTHON_OUTPUT_MARKER = 'persisted-python-ran'; +const PERSISTED_GENERATED_PYTHON = `print("${PERSISTED_PYTHON_OUTPUT_MARKER}")`; +const PERSISTED_MARKDOWN_TEXT = 'Fifth-run markdown from the E2E agent'; +const PERSISTED_FINAL_AGENT_TEXT = 'Transcript that must survive the save in full.'; +// executeAgentCell's first output item — all that survived the save before the streamed-item fix. +const AGENT_FIRST_OUTPUT_ITEM = '[Agent] Planning next steps...'; +const SNAPSHOT_WRITE_TIMEOUT = 60_000; +const SNAPSHOT_POLL_INTERVAL = 1_500; + +// Mixed batch, failing first cell: the two markers after it must never render. +const BATCH_FAILURE_MARKER = 'e2e-batch-boom'; +const BATCH_TRAILING_MARKER = 'e2e-batch-trailing-ran'; +const BATCH_AGENT_PYTHON_MARKER = 'e2e-batch-agent-python-ran'; +const BATCH_AGENT_MARKDOWN_TEXT = 'Markdown the agent must never get to write'; +const BATCH_AGENT_FINAL_TEXT = 'Summary the agent must never get to write.'; +/** + * How long the forbidden markers must stay away. A batch that carried on renders them within one + * agent round trip to the local mock (~1–2s) plus one cell execution (~1–3s); this keeps a wide + * margin over that without being open-ended, since a passing run spends the whole window. + */ +const BATCH_SETTLE_WINDOW = 10_000; + +// The generated cell prints, then sleeps: a bounded window in which the notebook is demonstrably +// running and Stop has something to interrupt. Long enough that a slow click still lands inside it, +// and never actually waited out on a passing run. +const STOP_SLEEP_MARKER = 'e2e-stop-sleeping'; +const STOP_GENERATED_PYTHON = `print("${STOP_SLEEP_MARKER}", flush=True)\nimport time\ntime.sleep(60)`; +const STOP_TRAILING_MARKER = 'e2e-stop-trailing-ran'; +const STOP_MARKDOWN_TEXT = 'Markdown the stopped agent must never write'; +const STOP_FINAL_TEXT = 'Summary the stopped agent must never write.'; +const AGENT_STOPPED_TEXT = '[Agent] Stopped'; +const STOP_ACKNOWLEDGED_TIMEOUT = 30_000; +// Shorter than the batch window: `[Agent] Stopped` already proves the run ended, so this only has +// to outlast the trailing cell that a batch which ignored the stop would dispatch next. +const STOP_SETTLE_WINDOW = 8_000; + +/** Keeps exactly one editor open: `clickRunAll` takes the first toolbar in DOM order. */ +async function openOnly(fileName: string): Promise { + await new WebView().switchBack().catch((error) => { + console.warn('[agent-block] switch back from webview before opening an editor:', error); + }); + await new EditorView().closeAllEditors().catch((error) => { + console.warn('[agent-block] close editors before opening the next notebook:', error); + }); - const missing = markers.filter((marker) => !text.includes(marker)); - const lingering = absentMarkers.filter((marker) => text.includes(marker)); - throw new Error( - `Timed out after ${timeout}ms waiting for notebook webview (${context}). Missing: ${JSON.stringify(missing)}. ` + - `Lingering: ${JSON.stringify(lingering)}. Last text: ${JSON.stringify(text)}` + await openWorkspaceFile(fileName); + await VSBrowser.instance.driver.wait( + async () => (await new EditorView().getOpenEditorTitles()).some((title) => title.includes(fileName)), + WORKBENCH_TIMEOUT, + `${fileName} did not open` ); } @@ -116,46 +158,40 @@ function assertOccurrences(rendered: string, needle: string, expected: number): ); } -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(); - - 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; + let workspaceDir = ''; before(async function () { screenshot = createScreenshotter(this); const copy = copyFixtureToTempDir(AGENT_FILE); cleanupTempDir = copy.cleanup; + workspaceDir = copy.tempDir; + + for (const fixture of [BATCH_FILE, STOP_FILE]) { + fs.copyFileSync( + path.resolve(process.cwd(), 'test', 'e2e', 'fixtures', fixture), + path.join(copy.tempDir, fixture) + ); + } 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` - ); - + // createEnvironment needs an active deepnote notebook; which one does not matter, and each + // group below binds the kernel for the notebook it actually runs. + await openOnly(AGENT_FILE); await createEnvironment(ENVIRONMENT_NAME); - await selectEnvironmentForNotebook(ENVIRONMENT_NAME, AGENT_FILE); await dismissAllNotifications(); await storeMockOpenAiApiKey(); - await screenshot('kernel-connected'); + await screenshot('environment-created'); }); async function releaseMockServer(): Promise { @@ -191,162 +227,415 @@ describe('Deepnote — running an agent block against a stand-in OpenAI API', fu console.warn('[agent-block] discard unsaved changes during cleanup:', error); }); } - await new Workbench().executeCommand(CLEAR_API_KEY_COMMAND).catch((error) => { + await clearStoredOpenAiApiKey().catch((error) => { console.warn('[agent-block] clear the stored OpenAI API key during cleanup:', error); }); }); - it('executes the generated code block, inserts its markdown block, and streams one transcript', async function () { - mockServer = await startMockOpenAiServer([ - { - match: { hasToolResult: false }, - response: { - toolCall: { - arguments: JSON.stringify({ code: GENERATED_PYTHON }), - id: 'call_e2e_code', - name: CODE_TOOL_NAME + describe('one agent block on its own', function () { + // Opens and binds the notebook this group runs, so the group does not care what ran before + // it. Closing and reopening drops the block's generated cells, which is why it happens here + // once and never between the tests below. + before(async function () { + await openOnly(AGENT_FILE); + await selectEnvironmentForNotebook(ENVIRONMENT_NAME, AGENT_FILE); + await dismissAllNotifications(); + }); + + it('executes the generated code block, inserts its markdown block, and streams one transcript', 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 } } - }, - { - match: { toolResultContains: PYTHON_OUTPUT_MARKER }, - response: { - toolCall: { - arguments: JSON.stringify({ content: EPHEMERAL_MARKDOWN_TEXT }), - id: 'call_e2e_markdown', - name: MARKDOWN_TOOL_NAME + ]); + + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); + + await awaitWebviewMarkers([PYTHON_OUTPUT_MARKER], FIRST_RUN_OUTPUT_TIMEOUT, 'generated code cell stdout'); + + const transcript = 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'); + + assertRenderedContiguously( + transcript, + `[Agent] Tool called: ${MARKDOWN_TOOL_NAME}\n\n[Agent] Tool output: ${MARKDOWN_TOOL_NAME}` + ); + assertRenderedContiguously(transcript, `[Agent] Text:\n${FINAL_AGENT_TEXT}`); + }); + + // Serial with prior it — block still owns first-run cells. + it('clears the cells its previous run generated instead of stacking a second copy', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: RERUN_GENERATED_PYTHON }), + id: 'call_e2e_rerun_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: RERUN_PYTHON_OUTPUT_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: RERUN_MARKDOWN_TEXT }), + id: 'call_e2e_rerun_markdown', + name: MARKDOWN_TOOL_NAME + } } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: RERUN_FINAL_AGENT_TEXT } } - }, - { - match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, - response: { content: FINAL_AGENT_TEXT } - } - ]); + ]); - await dismissAllNotifications(); - await clickRunAll(AGENT_FILE); + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); - await awaitWebviewMarkers([PYTHON_OUTPUT_MARKER], FIRST_RUN_OUTPUT_TIMEOUT, 'generated code cell stdout'); + const rendered = await awaitWebviewMarkers( + [RERUN_PYTHON_OUTPUT_MARKER, RERUN_MARKDOWN_TEXT, RERUN_FINAL_AGENT_TEXT], + AGENT_RUN_TIMEOUT, + 'second agent run' + ); - const transcript = 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'); - - assertRenderedContiguously( - transcript, - `[Agent] Tool called: ${MARKDOWN_TOOL_NAME}\n\n[Agent] Tool output: ${MARKDOWN_TOOL_NAME}` - ); - assertRenderedContiguously(transcript, `[Agent] Text:\n${FINAL_AGENT_TEXT}`); - }); + await screenshot('agent-rerun'); - // Serial with prior it — block still owns first-run cells. - it('clears the cells its previous run generated instead of stacking a second copy', async function () { - mockServer = await startMockOpenAiServer([ - { - match: { hasToolResult: false }, - response: { - toolCall: { - arguments: JSON.stringify({ code: RERUN_GENERATED_PYTHON }), - id: 'call_e2e_rerun_code', - name: CODE_TOOL_NAME + // assertOccurrences — retries would duplicate markers. + assertOccurrences(rendered, PYTHON_OUTPUT_MARKER, 0); + assertOccurrences(rendered, EPHEMERAL_MARKDOWN_TEXT, 0); + assertOccurrences(rendered, RERUN_PYTHON_OUTPUT_MARKER, 1); + assertOccurrences(rendered, RERUN_MARKDOWN_TEXT, 1); + assertOccurrences(rendered, STALE_CELLS_ERROR_TEXT, 0); + }); + + // Self-contained: generates the run it clears, so it survives --grep and a Mocha retry (Run All + // drops any stale generated cells first). + it('clears the whole generated run from the agent block status bar button', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: CLEAR_RUN_GENERATED_PYTHON }), + id: 'call_e2e_clear_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: CLEAR_RUN_PYTHON_OUTPUT_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: CLEAR_RUN_MARKDOWN_TEXT }), + id: 'call_e2e_clear_markdown', + name: MARKDOWN_TOOL_NAME + } } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: CLEAR_RUN_FINAL_AGENT_TEXT } } - }, - { - match: { toolResultContains: RERUN_PYTHON_OUTPUT_MARKER }, - response: { - toolCall: { - arguments: JSON.stringify({ content: RERUN_MARKDOWN_TEXT }), - id: 'call_e2e_rerun_markdown', - name: MARKDOWN_TOOL_NAME + ]); + + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); + + await awaitWebviewMarkers( + [CLEAR_RUN_PYTHON_OUTPUT_MARKER, CLEAR_RUN_MARKDOWN_TEXT, CLEAR_RUN_FINAL_AGENT_TEXT], + AGENT_RUN_TIMEOUT, + 'agent run whose cells the button clears' + ); + + await clickCellStatusBarItem(CLEAR_EPHEMERAL_BUTTON); + await confirmModalDialog(CLEAR_EPHEMERAL_CONFIRM, { messageIncludes: CLEAR_EPHEMERAL_CONFIRM_TEXT }); + + // The button lives on the agent block and takes both cells its run generated. Requiring the + // agent's own transcript to survive keeps an unreadable webview (which reads as '') from + // passing this as "the generated cells are gone". + await awaitWebviewMarkers( + [CLEAR_RUN_FINAL_AGENT_TEXT], + CLEAR_EPHEMERAL_TIMEOUT, + 'ephemeral cells cleared', + [CLEAR_RUN_PYTHON_OUTPUT_MARKER, CLEAR_RUN_MARKDOWN_TEXT] + ); + + await screenshot('agent-ephemeral-cleared'); + }); + + it('hands the agent a failed generated cell and lets the run carry on', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: FAILING_GENERATED_PYTHON }), + id: 'call_e2e_failing_code', + name: CODE_TOOL_NAME + } + } + }, + // Only a tool result carrying "Execution failed:" reaches this leg. The mock runs with + // --strict, so if the handler swallowed the failure and reported success instead, this + // request goes unmatched, the agent errors out, and the markers below never render — + // which is what makes the assertion about the tool result and not just about the cell. + { + match: { toolResultContains: EXECUTION_FAILED_TEXT }, + response: { + toolCall: { + arguments: JSON.stringify({ content: FAILURE_RECOVERY_MARKDOWN }), + id: 'call_e2e_failing_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: FAILURE_FINAL_AGENT_TEXT } + } + ]); + + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); + + await awaitWebviewMarkers( + [FAILING_PYTHON_MARKER, FAILURE_RECOVERY_MARKDOWN, FAILURE_FINAL_AGENT_TEXT], + AGENT_RUN_TIMEOUT, + 'agent run whose generated cell raises' + ); + + await screenshot('agent-generated-cell-failed'); + }); + + it('persists every streamed output item to the snapshot, not just the first', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: PERSISTED_GENERATED_PYTHON }), + id: 'call_e2e_persisted_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: PERSISTED_PYTHON_OUTPUT_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: PERSISTED_MARKDOWN_TEXT }), + id: 'call_e2e_persisted_markdown', + name: MARKDOWN_TOOL_NAME + } } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: PERSISTED_FINAL_AGENT_TEXT } } - }, - { - match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, - response: { content: RERUN_FINAL_AGENT_TEXT } + ]); + + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); + + await awaitWebviewMarkers( + [PERSISTED_PYTHON_OUTPUT_MARKER, PERSISTED_FINAL_AGENT_TEXT], + AGENT_RUN_TIMEOUT, + 'agent run whose transcript must survive the save' + ); + + // With snapshots on (the default) the main .deepnote has outputs stripped, so the agent's + // transcript only exists in the sidecar. The save is deferred off queue completion, hence + // the poll rather than a single read. + const snapshotsDir = path.join(workspaceDir, 'snapshots'); + const driver = VSBrowser.instance.driver; + const deadline = Date.now() + SNAPSHOT_WRITE_TIMEOUT; + let transcript = ''; + + while (Date.now() < deadline) { + const files = fs.existsSync(snapshotsDir) + ? fs.readdirSync(snapshotsDir).filter((file) => file.endsWith('_latest.snapshot.deepnote')) + : []; + transcript = + files.length > 0 + ? blockStreamOutputText( + fs.readFileSync(path.join(snapshotsDir, files[0]), 'utf8'), + AGENT_BLOCK_ID + ) + : ''; + + if (transcript.includes(PERSISTED_FINAL_AGENT_TEXT)) { + break; + } + + await driver.sleep(SNAPSHOT_POLL_INTERVAL); } - ]); - await dismissAllNotifications(); - await clickRunAll(AGENT_FILE); - - const rendered = await awaitWebviewMarkers( - [RERUN_PYTHON_OUTPUT_MARKER, RERUN_MARKDOWN_TEXT, RERUN_FINAL_AGENT_TEXT], - AGENT_RUN_TIMEOUT, - 'second agent run' - ); - - await screenshot('agent-rerun'); - - // assertOccurrences — retries would duplicate markers. - assertOccurrences(rendered, PYTHON_OUTPUT_MARKER, 0); - assertOccurrences(rendered, EPHEMERAL_MARKDOWN_TEXT, 0); - assertOccurrences(rendered, RERUN_PYTHON_OUTPUT_MARKER, 1); - assertOccurrences(rendered, RERUN_MARKDOWN_TEXT, 1); - assertOccurrences(rendered, STALE_CELLS_ERROR_TEXT, 0); + // The first item alone is what a truncating converter leaves behind, so requiring it AND the + // later ones is the whole assertion: the run is on disk and it is not just its opening line. + for (const expected of [ + AGENT_FIRST_OUTPUT_ITEM, + `[Agent] Tool called: ${CODE_TOOL_NAME}`, + `[Agent] Tool called: ${MARKDOWN_TOOL_NAME}`, + PERSISTED_FINAL_AGENT_TEXT + ]) { + if (!transcript.includes(expected)) { + throw new Error( + `Saved snapshot is missing ${JSON.stringify(expected)} — the agent's streamed output items ` + + `did not all survive the save. Saved transcript: ${JSON.stringify(transcript)}` + ); + } + } + }); }); - // Self-contained: generates the run it clears, so it survives --grep and a Mocha retry (Run All - // drops any stale generated cells first). - it('clears the whole generated run from the agent block status bar button', async function () { - mockServer = await startMockOpenAiServer([ - { - match: { hasToolResult: false }, - response: { - toolCall: { - arguments: JSON.stringify({ code: CLEAR_RUN_GENERATED_PYTHON }), - id: 'call_e2e_clear_code', - name: CODE_TOOL_NAME + describe('a mixed batch of kernel cells and an agent block', function () { + // Each test here opens and binds the notebook it needs, so both stand alone: run either by + // itself, in either order, or after the group above, and it sets up the same state. + it('stops at a failing cell instead of running the agent block and the cell after it', async function () { + // Scripted so a batch that carried on has something to render. Nothing should reach the mock: + // if these legs are never requested the agent never started, which is the point. + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: `print("${BATCH_AGENT_PYTHON_MARKER}")` }), + id: 'call_e2e_batch_code', + name: CODE_TOOL_NAME + } } - } - }, - { - match: { toolResultContains: CLEAR_RUN_PYTHON_OUTPUT_MARKER }, - response: { - toolCall: { - arguments: JSON.stringify({ content: CLEAR_RUN_MARKDOWN_TEXT }), - id: 'call_e2e_clear_markdown', - name: MARKDOWN_TOOL_NAME + }, + { + match: { toolResultContains: BATCH_AGENT_PYTHON_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: BATCH_AGENT_MARKDOWN_TEXT }), + id: 'call_e2e_batch_markdown', + name: MARKDOWN_TOOL_NAME + } } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: BATCH_AGENT_FINAL_TEXT } } - }, - { - match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, - response: { content: CLEAR_RUN_FINAL_AGENT_TEXT } - } - ]); + ]); + + await openOnly(BATCH_FILE); + await selectEnvironmentForNotebook(ENVIRONMENT_NAME, BATCH_FILE); + await dismissAllNotifications(); + await clickRunAll(BATCH_FILE); + + await awaitWebviewMarkers( + [BATCH_FAILURE_MARKER], + FIRST_RUN_OUTPUT_TIMEOUT, + 'traceback from the failing cell' + ); + + await assertMarkersStayAbsent( + [BATCH_AGENT_PYTHON_MARKER, BATCH_AGENT_MARKDOWN_TEXT, BATCH_AGENT_FINAL_TEXT, BATCH_TRAILING_MARKER], + BATCH_SETTLE_WINDOW, + 'a failing cell must end the batch, so neither the agent block nor the cell after it runs' + ); + + await screenshot('batch-stopped-at-failure'); + }); - await dismissAllNotifications(); - await clickRunAll(AGENT_FILE); - - await awaitWebviewMarkers( - [CLEAR_RUN_PYTHON_OUTPUT_MARKER, CLEAR_RUN_MARKDOWN_TEXT, CLEAR_RUN_FINAL_AGENT_TEXT], - AGENT_RUN_TIMEOUT, - 'agent run whose cells the button clears' - ); - - await clickCellStatusBarItem(CLEAR_EPHEMERAL_BUTTON); - await confirmModalDialog(CLEAR_EPHEMERAL_CONFIRM, { messageIncludes: CLEAR_EPHEMERAL_CONFIRM_TEXT }); - - // The button lives on the agent block and takes both cells its run generated. Requiring the - // agent's own transcript to survive keeps an unreadable webview (which reads as '') from - // passing this as "the generated cells are gone". - await awaitWebviewMarkers([CLEAR_RUN_FINAL_AGENT_TEXT], CLEAR_EPHEMERAL_TIMEOUT, 'ephemeral cells cleared', [ - CLEAR_RUN_PYTHON_OUTPUT_MARKER, - CLEAR_RUN_MARKDOWN_TEXT - ]); - - await screenshot('agent-ephemeral-cleared'); + it('stops the agent and the cell after it when Interrupt is clicked mid-run', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: STOP_GENERATED_PYTHON }), + id: 'call_e2e_stop_code', + name: CODE_TOOL_NAME + } + } + }, + // Reached only if the agent survived the stop and asked for its next turn. + { + match: { toolResultContains: STOP_SLEEP_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: STOP_MARKDOWN_TEXT }), + id: 'call_e2e_stop_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: STOP_FINAL_TEXT } + } + ]); + + await openOnly(STOP_FILE); + await selectEnvironmentForNotebook(ENVIRONMENT_NAME, STOP_FILE); + await dismissAllNotifications(); + await clickRunAll(STOP_FILE); + + // Both markers, not just the sleep one: a Mocha retry starts with the previous attempt's + // generated cell still on screen, and the agent's own transcript is cleared at run start, so + // the tool-call line is what proves we are looking at this attempt. + await awaitWebviewMarkers( + [`[Agent] Tool called: ${CODE_TOOL_NAME}`, STOP_SLEEP_MARKER], + AGENT_RUN_TIMEOUT, + 'the generated cell reached its sleep, so the notebook is running and Stop has something to interrupt' + ); + + await clickInterrupt(STOP_FILE); + + await awaitWebviewMarkers( + [AGENT_STOPPED_TEXT], + STOP_ACKNOWLEDGED_TIMEOUT, + 'the agent reports the stop rather than treating the interrupted cell as a retryable failure', + [STOP_MARKDOWN_TEXT, STOP_FINAL_TEXT] + ); + + await assertMarkersStayAbsent( + [STOP_MARKDOWN_TEXT, STOP_FINAL_TEXT, STOP_TRAILING_MARKER], + STOP_SETTLE_WINDOW, + 'a stopped agent ends the batch, so nothing after it runs' + ); + + await screenshot('batch-stopped-by-interrupt'); + }); }); }); From 4745d5bab3b3218fdfd3d2399b10405eee564cd1 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 17 Aug 2026 08:24:40 +0000 Subject: [PATCH 79/80] fix(tests): answer the save prompt when the E2E switches notebooks Both mixed-batch tests failed on the first CI run, ~5s into openOnly, with "TimeoutError: Waiting until element is visible" -- nowhere near what they assert. The cause is one line earlier in the log: closing the editor raised the save prompt for a notebook an agent run had dirtied, the intercepted click was swallowed by the surrounding catch, and the modal then dimmed the workbench so every later click landed on the overlay. Revert before closing so the prompt does not appear, and answer it if it does anyway. The check for surviving editors has to gate that answer: confirmModalDialog waits out its full timeout before throwing when no dialog is up, so calling it unconditionally would trade a 5s failure for a 60s one. The suite's own `after` already pairs revert, close and discard this way; this brings the mid-suite switch in line with it. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- test/e2e/suite/agentBlock.e2e.test.ts | 29 ++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts index 1c28181e55..87fa34be6e 100644 --- a/test/e2e/suite/agentBlock.e2e.test.ts +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -117,15 +117,42 @@ const STOP_ACKNOWLEDGED_TIMEOUT = 30_000; // to outlast the trailing cell that a batch which ignored the stop would dispatch next. const STOP_SETTLE_WINDOW = 8_000; -/** Keeps exactly one editor open: `clickRunAll` takes the first toolbar in DOM order. */ +/** + * Keeps exactly one editor open: `clickRunAll` takes the first toolbar in DOM order. + * + * A run leaves the notebook dirty, so closing it raises the save prompt. Revert first so the close + * is clean, and answer the prompt if one appears anyway: an unanswered modal dims the workbench and + * intercepts every later click, which surfaces as an unrelated "element is visible" timeout in + * whatever runs next rather than here. + */ async function openOnly(fileName: string): Promise { await new WebView().switchBack().catch((error) => { console.warn('[agent-block] switch back from webview before opening an editor:', error); }); + + const alreadyOpen = await new EditorView().getOpenEditorTitles().catch(() => [] as string[]); + if (alreadyOpen.length > 0) { + await new Workbench().executeCommand(REVERT_FILE_COMMAND).catch((error) => { + console.warn('[agent-block] revert notebook before closing it:', error); + }); + } + await new EditorView().closeAllEditors().catch((error) => { console.warn('[agent-block] close editors before opening the next notebook:', error); }); + // Only when editors survived the close, since `confirmModalDialog` waits out its full timeout + // and then throws when no dialog is up. + const stillOpen = await new EditorView().getOpenEditorTitles().catch(() => [] as string[]); + if (stillOpen.length > 0) { + await confirmModalDialog(DISCARD_CHANGES_BUTTON).catch((error) => { + console.warn('[agent-block] discard unsaved changes before opening the next notebook:', error); + }); + await new EditorView().closeAllEditors().catch((error) => { + console.warn('[agent-block] close editors after discarding unsaved changes:', error); + }); + } + await openWorkspaceFile(fileName); await VSBrowser.instance.driver.wait( async () => (await new EditorView().getOpenEditorTitles()).some((title) => title.includes(fileName)), From 3f73c76b5c5eea51a1881c295a2d3d459992849b Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 17 Aug 2026 14:56:44 +0000 Subject: [PATCH 80/80] fix(agent-block): abort the model request when a run is stopped The E2E caught what the unit tests could not: pressing Stop mid-run left the agent working and then reported the run as successful. The extension log shows it plainly -- the interrupt lands at 09:53:58, and six seconds later the run finishes down the success path with an empty result: 09:53:55.712 Agent cell: starting executeAgentBlock 09:53:58.255 [error] No kernel associated with the notebook (handleInterrupt) 09:54:04.289 Agent cell: executeAgentBlock completed, finalOutput length=0 The cancellation did fire; it just could not end the run. Throwing from a tool callback never could, because runtime-core wraps those callbacks: } catch (error) { ... return `Execution error: ${executionError.message}`; } The throw becomes a string the model reads as a retryable tool failure, so a stop made the agent do more work, and the loop only wound down once it ran out of turns -- landing on executeAgentCell's success branch, which called endExecution(true) for a run the user had stopped. 0.5.0 adds the AbortSignal the previous comment was waiting on. It calls signal.throwIfAborted() inside runtime-core, outside that catch, and forwards the signal to agent.stream as abortSignal, so the in-flight request is aborted rather than left to finish. Bridging the cancellation token to it is the whole fix; isStopped already recognised AbortError. Verified against the built extension, not mocks: the agent now reports "Agent cell execution stopped" 4ms after the interrupt, and the full agent E2E suite passes locally, 7/7. Note runtime-core 0.5.0 pins @deepnote/blocks 4.7.0, so the lock now carries a nested copy alongside the root ^4.6.0. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01WAA6y9KqdkR8p5H26VD2ee --- package-lock.json | 43 +++++++++++++++---- package.json | 2 +- .../deepnote/agentCellExecutionHandler.ts | 13 +++--- 3 files changed, 42 insertions(+), 16 deletions(-) diff --git a/package-lock.json b/package-lock.json index 20463d02e4..8f2fcbf25a 100644 --- a/package-lock.json +++ b/package-lock.json @@ -14,7 +14,7 @@ "@deepnote/blocks": "^4.6.0", "@deepnote/convert": "^4.0.0", "@deepnote/database-integrations": "^1.5.0", - "@deepnote/runtime-core": "^0.4.0", + "@deepnote/runtime-core": "^0.5.0", "@deepnote/sql-language-server": "^3.0.0", "@enonic/fnv-plus": "^1.3.0", "@jupyter-widgets/base": "^6.0.8", @@ -2165,14 +2165,14 @@ } }, "node_modules/@deepnote/runtime-core": { - "version": "0.4.0", - "resolved": "https://registry.npmjs.org/@deepnote/runtime-core/-/runtime-core-0.4.0.tgz", - "integrity": "sha512-iS5E2FUxAT83cRDt9cGvhO+CR9tuLG6+PmLT4Vm8ffQUzPSi5b0v0AuwnSOrMDPlt4SXM9DXxtB3iHhWUk/0kQ==", + "version": "0.5.0", + "resolved": "https://registry.npmjs.org/@deepnote/runtime-core/-/runtime-core-0.5.0.tgz", + "integrity": "sha512-+0Dbs5IhSsRFLg5XZqmr1gKo0GPeMwKNxezhAP6MgYmVECvxb+kHRoI5i4aazRsni/hk8i5AqipyYKYu4TT3Og==", "license": "Apache-2.0", "dependencies": { "@ai-sdk/mcp": "^1.0.25", "@ai-sdk/openai": "^3.0.0", - "@deepnote/blocks": "4.6.0", + "@deepnote/blocks": "4.7.0", "@jupyterlab/nbformat": "^4.3.2", "@jupyterlab/services": "^7.3.2", "ai": "^6.0.0", @@ -2181,6 +2181,21 @@ "zod": "3.25.76" } }, + "node_modules/@deepnote/runtime-core/node_modules/@deepnote/blocks": { + "version": "4.7.0", + "resolved": "https://registry.npmjs.org/@deepnote/blocks/-/blocks-4.7.0.tgz", + "integrity": "sha512-GW5jIpO2Sr7R2LaokuF8js60FHw4DLk1e4bq6BnMs2QkxGr2QF6JbRONl82V7zEV3iJQesXX+mFuTMWp9jCxPA==", + "license": "Apache-2.0", + "dependencies": { + "ts-dedent": "^2.2.0", + "yaml": "^2.8.3", + "zod": "3.25.76" + }, + "engines": { + "node": ">=22.14.0", + "pnpm": ">=10.17.1" + } + }, "node_modules/@deepnote/runtime-core/node_modules/ws": { "version": "8.21.0", "resolved": "https://registry.npmjs.org/ws/-/ws-8.21.0.tgz", @@ -36964,13 +36979,13 @@ } }, "@deepnote/runtime-core": { - "version": "0.4.0", - "resolved": "https://registry.npmjs.org/@deepnote/runtime-core/-/runtime-core-0.4.0.tgz", - "integrity": "sha512-iS5E2FUxAT83cRDt9cGvhO+CR9tuLG6+PmLT4Vm8ffQUzPSi5b0v0AuwnSOrMDPlt4SXM9DXxtB3iHhWUk/0kQ==", + "version": "0.5.0", + "resolved": "https://registry.npmjs.org/@deepnote/runtime-core/-/runtime-core-0.5.0.tgz", + "integrity": "sha512-+0Dbs5IhSsRFLg5XZqmr1gKo0GPeMwKNxezhAP6MgYmVECvxb+kHRoI5i4aazRsni/hk8i5AqipyYKYu4TT3Og==", "requires": { "@ai-sdk/mcp": "^1.0.25", "@ai-sdk/openai": "^3.0.0", - "@deepnote/blocks": "4.6.0", + "@deepnote/blocks": "4.7.0", "@jupyterlab/nbformat": "^4.3.2", "@jupyterlab/services": "^7.3.2", "ai": "^6.0.0", @@ -36979,6 +36994,16 @@ "zod": "3.25.76" }, "dependencies": { + "@deepnote/blocks": { + "version": "4.7.0", + "resolved": "https://registry.npmjs.org/@deepnote/blocks/-/blocks-4.7.0.tgz", + "integrity": "sha512-GW5jIpO2Sr7R2LaokuF8js60FHw4DLk1e4bq6BnMs2QkxGr2QF6JbRONl82V7zEV3iJQesXX+mFuTMWp9jCxPA==", + "requires": { + "ts-dedent": "^2.2.0", + "yaml": "2.8.3", + "zod": "3.25.76" + } + }, "ws": { "version": "8.21.0", "resolved": "https://registry.npmjs.org/ws/-/ws-8.21.0.tgz", diff --git a/package.json b/package.json index abdfce4b76..3596fb8007 100644 --- a/package.json +++ b/package.json @@ -2704,7 +2704,7 @@ "@deepnote/blocks": "^4.6.0", "@deepnote/convert": "^4.0.0", "@deepnote/database-integrations": "^1.5.0", - "@deepnote/runtime-core": "^0.4.0", + "@deepnote/runtime-core": "^0.5.0", "@deepnote/sql-language-server": "^3.0.0", "@enonic/fnv-plus": "^1.3.0", "@jupyter-widgets/base": "^6.0.8", diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index 3a4c842ead..d4f9632d25 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -166,10 +166,9 @@ function isStopped(error: unknown): boolean { * Runs an agent block into the cell output and inserts generated cells below. * Call `removeEphemeralCellsForAgentBlocks` on the batch first. Never rejects — errors become stderr on the cell. * - * `token` stops the run. It reaches the model only indirectly: the host refuses tool calls and throws - * from the event callback, so an in-flight model turn still finishes. Once `AgentBlockContext` carries - * an `AbortSignal` (present in runtime-core's `main`, unreleased), bridge the token to one and pass it - * as `signal` — runtime-core forwards it to `agent.stream`, which aborts the request itself. + * `token` stops the run, bridged to the `AbortSignal` runtime-core forwards to `agent.stream`, so the + * in-flight model request is aborted rather than left to finish. Throwing from a tool callback cannot + * stop it: runtime-core catches that and hands the model an `Execution error: …` string to retry. */ export async function executeAgentCell( cell: NotebookCell, @@ -180,10 +179,13 @@ export async function executeAgentCell( ): Promise { const executeAgentBlockFn = options?.executeAgentBlockFn ?? executeAgentBlock; const execution = controller.createNotebookCellExecution(cell); + const stopController = new AbortController(); + const stopSubscription = token.onCancellationRequested(() => stopController.abort()); // The agent runs off the kernel, so nothing announces it on the internal shim — the only source // SnapshotService and the execute_cell analytics read. Without this the run is invisible to both. const endExecution = (success: boolean) => { + stopSubscription.dispose(); execution.end(success, Date.now()); notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Idle); }; @@ -232,8 +234,7 @@ export async function executeAgentCell( openAiToken, ...getProjectAgentContext(cell.notebook), notebookContext, - // The guards sit outside the `try`s: those turn every throw into a string the model reads - // as a tool failure worth retrying, which is how a stop used to make the agent do more work. + signal: stopController.signal, addMarkdownBlock: async ({ content }: { content: string }) => { Cancellation.throwIfCanceled(token);