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/package-lock.json b/package-lock.json index 53cf33ee0e..8f7903ce4f 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 d8a7c43e06..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%", @@ -2676,6 +2686,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/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/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts b/src/kernels/deepnote/deepnoteLspClientManager.node.vscode.test.ts index 5195d0a28b..31a6c85845 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.node'; 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 fe41270e0f..112f7f0e09 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: + * - 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 } 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 MAX_OUTPUT_TRACKING_LENGTH = 5000; +const SERVER_STARTUP_TIMEOUT_MS = 120_000; +const GRACEFUL_SHUTDOWN_TIMEOUT_MS = 3000; + interface ServerLockFile { sessionId: string; pid: number; @@ -42,65 +49,52 @@ type PendingOperation = interface ProjectContext { environmentId: string; - serverProcess: ObservableExecutionResult | 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 projectContexts: Map = new Map(); + private readonly serverOutputByFile: Map = new Map(); 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, @@ -112,12 +106,10 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension token?: CancellationToken ): Promise { const fileKey = deepnoteFileUri.fsPath; - const serverKey = `${fileKey}-${environmentId}`; - // Wait for any pending operations on this environment to complete - 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 { @@ -125,44 +117,40 @@ 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; } - // Start the operation if not already pending - pendingOp = this.pendingOperations.get(serverKey); + pendingOp = this.pendingOperations.get(fileKey); 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}...` + `Stopping existing server for ${fileKey} with environmentId ${existingEnvironmentId} to start new one with environmentId ${environmentId}...` ); await this.stopServerForEnvironment(existingContext, deepnoteFileUri, token); - // TODO - Clear controllers for the notebook ? + existingContext.environmentId = environmentId; } } else { - const newContext = { + const newContext: ProjectContext = { environmentId, - serverProcess: null, serverInfo: null }; - this.projectContexts.set(serverKey, newContext); - + this.projectContexts.set(fileKey, newContext); existingContext = newContext; } - // Start the operation and track it const operation = { type: 'start' as const, promise: this.startServerForEnvironment( @@ -176,34 +164,34 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension token ) }; - this.pendingOperations.set(serverKey, operation); + this.pendingOperations.set(fileKey, operation); 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); + if (this.pendingOperations.get(fileKey) === operation) { + this.pendingOperations.delete(fileKey); } } } /** - * 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 + 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...`); @@ -216,7 +204,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 +213,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 +220,13 @@ 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) */ private async startServerForEnvironment( projectContext: ProjectContext, @@ -247,11 +239,9 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension token?: CancellationToken ): Promise { const fileKey = deepnoteFileUri.fsPath; - const serverKey = `${fileKey}-${environmentId}`; 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,173 +258,60 @@ 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); + 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 jupyter port ${jupyterPort} and lsp port ${lspPort} 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 || ''}`; - - // Also set VIRTUAL_ENV to indicate we're in a venv - env.VIRTUAL_ENV = venvPath.fsPath; + 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'; + // Initialize output tracking for error reporting + this.serverOutputByFile.set(fileKey, { stdout: '', stderr: '' }); - // 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'; + let serverInfo: DeepnoteServerInfo | undefined; + try { + serverInfo = await startServer({ + pythonEnv: venvPath.fsPath, + workingDirectory: path.dirname(deepnoteFileUri.fsPath), + startupTimeoutMs: SERVER_STARTUP_TIMEOUT_MS, + env: extraEnv + }); + } catch (error) { + const capturedOutput = this.serverOutputByFile.get(fileKey); + this.serverOutputByFile.delete(fileKey); - // Inject SQL integration environment variables - if (this.sqlIntegrationEnvVars) { - logger.debug( - `DeepnoteServerStarter: Injecting SQL integration env vars for ${fileKey} with environmentId ${environmentId}` + throw new DeepnoteServerStartupError( + interpreter.uri.fsPath, + serverInfo?.jupyterPort ?? 0, + 'unknown', + capturedOutput?.stdout || '', + capturedOutput?.stderr || '', + error instanceof Error ? error : new Error(`${error}`) ); - 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; + projectContext.serverInfo = serverInfo; - this.serverProcesses.set(serverKey, serverProcess); + // Set up output channel logging from the server process + this.monitorServerOutput(fileKey, serverInfo); - // Track disposables for this environment - const disposables: IDisposable[] = []; - this.disposablesByFile.set(serverKey, disposables); - - // 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 - ); - - // Wait for server to be ready - const url = `http://localhost:${jupyterPort}`; - const serverInfo = { url, jupyterPort, lspPort }; - this.serverInfos.set(serverKey, serverInfo); - - // Write lock file for the server process - const serverPid = serverProcess.proc?.pid; + // 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}`); - } - - 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.warn(`Could not get PID for server process for ${fileKey}`); } - 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 ${serverInfo.url} for ${fileKey}`); + this.outputChannel.appendLine(l10n.t('✓ Deepnote server running at {0}', serverInfo.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, + projectContext: ProjectContext, deepnoteFileUri: Uri, token?: CancellationToken ): Promise { @@ -442,23 +319,20 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension Cancellation.throwIfCanceled(token); - // const serverProcess = this.serverProcesses.get(fileKey); - const serverProcess = projectContext?.serverProcess; + const { serverInfo } = projectContext; - if (serverProcess) { - const serverPid = serverProcess.proc?.pid; + if (serverInfo) { + const serverPid = serverInfo.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(serverInfo); 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 + projectContext.serverInfo = null; + if (serverPid) { await this.deleteLockFile(serverPid); } @@ -467,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()); @@ -474,281 +350,141 @@ 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`, { signal: AbortSignal.timeout(5000) }); + 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. - * - * @param key The environment ID to reserve ports for - * @returns Object with jupyterPort and lspPort + * Gather SQL integration environment variables for the deepnote-toolkit server. */ - 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 - const previousLock = this.portAllocationLock; - let releaseLock: () => void; - const currentLock = new Promise((resolve) => { - releaseLock = resolve; - }); - 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); - } - } - - // 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' - })` - ); + private async gatherSqlIntegrationEnvVars( + deepnoteFileUri: Uri, + environmentId: string, + token?: CancellationToken + ): Promise> { + const extraEnv: Record = {}; - return { jupyterPort, lspPort }; - } finally { - // Release the lock to allow next allocation in the chain to proceed - releaseLock!(); + if (!this.sqlIntegrationEnvVars) { + logger.debug('DeepnoteServerStarter: SqlIntegrationEnvironmentVariablesProvider not available'); + return extraEnv; } - } - - /** - * 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 - */ - 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 - ); - - const nextPort = candidatePort + 1; - - // 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}` - ); - - if (isNextPortAvailable) { - // Found a consecutive pair! - return { jupyterPort: candidatePort, lspPort: nextPort }; - } - // Consecutive port not available - mark both as unavailable and try next - portsInUse.add(candidatePort); - portsInUse.add(nextPort); - } + const fileKey = deepnoteFileUri.fsPath; - // 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 - ) + logger.debug( + `DeepnoteServerStarter: Injecting SQL integration env vars for ${fileKey} with environmentId ${environmentId}` ); - } - - /** - * 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. - */ - private async isPortAvailable(port: number): Promise { try { - const inUse = await tcpPortUsed.check(port, '127.0.0.1'); - if (inUse) { - return false; - } - - // 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; - } - logger.warn(`Failed to check IPv6 port availability for ${port}:`, error); - return false; + 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'); } } catch (error) { - logger.warn(`Failed to check port availability for ${port}:`, error); - return false; + logger.error('DeepnoteServerStarter: Failed to get SQL integration env vars', error); } + + return extraEnv; } /** - * Find an available port starting from the given port number. - * Checks both our internal portsInUse set and system availability by actually binding to test. + * Stream stdout/stderr from the server process to the VSCode output channel. */ - 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; + private monitorServerOutput(fileKey: string, serverInfo: DeepnoteServerInfo): void { + const proc = serverInfo.process; + const disposables: IDisposable[] = []; + this.disposablesByFile.set(fileKey, disposables); + + if (proc.stdout) { + const stdout = proc.stdout; + const onData = (data: Buffer) => { + 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); } - } - - // Try next port - port++; - attempts++; + }; + stdout.on('data', onData); + disposables.push({ + dispose: () => { + stdout.off('data', onData); + } + }); } - 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(', ') - ) - ); + if (proc.stderr) { + const stderr = proc.stderr; + const onData = (data: Buffer) => { + 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({ + dispose: () => { + stderr.off('data', onData); + } + }); + } } 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.promise, 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.serverInfo) { + const pid = ctx.serverInfo.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.serverInfo).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 +493,16 @@ export class DeepnoteServerStarter implements IDeepnoteServerStarter, IExtension } } - // Clear all maps - this.serverProcesses.clear(); - this.serverInfos.clear(); this.disposablesByFile.clear(); this.pendingOperations.clear(); + this.projectContexts.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 +512,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 +531,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 +543,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 +555,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 +575,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 +603,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 +618,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 +644,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 +673,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 +691,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 +721,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 +736,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..f90e3a8ec1 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 * as fakeTimers from '@sinonjs/fake-timers'; 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 delegated to + * @deepnote/runtime-core's startServer/stopServer. These tests focus on the + * extension-specific layers: 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,115 @@ 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); + 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) + ); - const debugStub = sinon.stub(logger, 'debug'); + const gatherEnvVars = getPrivateMethod(starterWithoutSql, 'gatherSqlIntegrationEnvVars'); + const { Uri } = await import('vscode'); + const result = await gatherEnvVars(Uri.file('/test/file.deepnote'), 'env1'); - try { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const isPortAvailable = getPrivateMethod(serverStarter as any, 'isPortAvailable'); - const result = await isPortAvailable(port); + assert.deepStrictEqual(result, {}); - 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(); - } + await starterWithoutSql.dispose(); }); - test('should return false when IPv6 check throws non-EAFNOSUPPORT error', async () => { - const port = 54325; - const ipv6Error = new Error('Some other IPv6 error'); + test('should return empty object when provider rejects with cancellation error', async () => { + const { CancellationError, Uri } = await import('vscode'); - // IPv4 check succeeds (port is available) - checkStub.onFirstCall().resolves(false); + const cancelledProvider = mock(); + when(cancelledProvider.getEnvironmentVariables(anything(), anything())).thenReject(new CancellationError()); - // IPv6 check throws a different error - checkStub.onSecondCall().rejects(ipv6Error); + const starterWithCancelledSql = new DeepnoteServerStarter( + instance(mockProcessServiceFactory), + instance(mockToolkitInstaller), + instance(mockAgentSkillsManager), + instance(mockOutputChannel), + instance(mockAsyncRegistry), + instance(cancelledProvider) + ); - const warnStub = sinon.stub(logger, 'warn'); + const gatherEnvVars = getPrivateMethod(starterWithCancelledSql, 'gatherSqlIntegrationEnvVars'); + const result = await gatherEnvVars(Uri.file('/test/file.deepnote'), 'env1'); - try { - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const isPortAvailable = getPrivateMethod(serverStarter as any, 'isPortAvailable'); - const result = await isPortAvailable(port); + assert.deepStrictEqual(result, {}); - 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(); - } + await starterWithCancelledSql.dispose(); }); }); - suite('findAvailablePort', () => { - test('should find an available port starting from given port', async () => { - const portsInUse = new Set(); - const startPort = 54400; + suite('dispose', () => { + let clock: fakeTimers.InstalledClock; - // 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); - }); - - test('should skip ports in portsInUse set', async () => { - const portsInUse = new Set([54500, 54501, 54502]); - const startPort = 54500; - - // 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); + setup(() => { + clock = fakeTimers.install(); }); - 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); + teardown(() => { + clock.uninstall(); }); - }); - suite('allocatePorts - Consecutive Port Allocation (Critical Bug Fix)', () => { - test('should allocate consecutive ports (LSP = Jupyter + 1)', async () => { - const key = 'test-consecutive-1'; + 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 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)' - ); + const starter = serverStarter as any; + 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 allocate different consecutive port pairs for multiple servers', async () => { - const key1 = 'test-server-1'; - const key2 = 'test-server-2'; - + test('should wait for in-flight pending operations before completing', async () => { // eslint-disable-next-line @typescript-eslint/no-explicit-any - const allocatePorts = getPrivateMethod(serverStarter as any, 'allocatePorts'); + const starter = serverStarter as any; - 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(); - }); + let resolveDeferred!: () => void; + const deferred = new Promise((resolve) => { + resolveDeferred = resolve; }); - try { - const key = 'test-blocked-lsp-port'; - - // 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()); - }); - } - }); - - test('should handle rapid sequential allocations', async () => { - const keys = ['seq-1', 'seq-2', 'seq-3', 'seq-4', 'seq-5']; - - // 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 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'); - }); - }); - - 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); - }); - - test('should return valid port numbers', async () => { - const key = 'valid-ports'; - - // eslint-disable-next-line @typescript-eslint/no-explicit-any - const allocatePorts = getPrivateMethod(serverStarter as any, 'allocatePorts'); - const result = await allocatePorts(key); - - // 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); - }); - }); - - 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` - ]; - - // 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.` - ); - } + starter.pendingOperations.set('/test/inflight.deepnote', { + type: 'stop', + promise: deferred + }); - // 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' - ); + let disposeResolved = false; + const disposePromise = serverStarter.dispose().then(() => { + disposeResolved = true; + }); - // Verify uniqueness of individual ports - const allPorts = results.flatMap((r) => [r.jupyterPort, r.lspPort]); - const uniquePorts = new Set(allPorts); + await clock.tickAsync(0); assert.strictEqual( - uniquePorts.size, - allPorts.length, - 'All allocated ports (both Jupyter and LSP) must be unique across concurrent calls' + disposeResolved, + false, + 'dispose() should not resolve while a pending operation is in flight' ); - }); - }); - - 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}` - ); - } + resolveDeferred(); + await clock.tickAsync(0); + await disposePromise; - 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()); - }) - ) - ); - } + assert.strictEqual(disposeResolved, true, 'dispose() should resolve after pending operation completes'); + assert.strictEqual(starter.pendingOperations.size, 0); }); }); }); diff --git a/src/kernels/deepnote/deepnoteTestHelpers.node.ts b/src/kernels/deepnote/deepnoteTestHelpers.node.ts new file mode 100644 index 0000000000..524c6e0e47 --- /dev/null +++ b/src/kernels/deepnote/deepnoteTestHelpers.node.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/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/environments/deepnoteEnvironmentTreeDataProvider.unit.test.ts b/src/kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.unit.test.ts index 4cc6d5df5d..05b690f3b7 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.node'; 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 ef64ae04e0..a0c17a31ba 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 { ServerInfo as RuntimeCoreServerInfo } from '@deepnote/runtime-core'; import * as vscode from 'vscode'; import { serializePythonEnvironment } from '../../platform/api/pythonApi'; @@ -185,10 +186,7 @@ export interface IDeepnoteServerStarter { dispose(): Promise; } -export interface DeepnoteServerInfo { - url: string; - jupyterPort: number; - lspPort: number; +export interface DeepnoteServerInfo extends RuntimeCoreServerInfo { token?: string; } 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..3ed36cf0e2 --- /dev/null +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -0,0 +1,327 @@ +import { + CancellationError, + CancellationToken, + 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 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'; +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; + + 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 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()); + + try { + await execution.clearOutput(); + + const prompt = cell.document.getText(); + + let accumulated = `[Agent] Planning next steps...`; + const output = new NotebookCellOutput([NotebookCellOutputItem.text(accumulated)]); + 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 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'); + } + + await removeEphemeralCellsForAgent(cell.notebook, agentBlock.id); + + let lastAgentEventType: AgentStreamEvent['type'] | undefined; + + const notebookContext = serializeNotebookContext({ + cells: cell.notebook.getCells().filter((c) => c.index !== cell.index) + }); + + const openAiToken = await getOpenAiApiKey(); + + 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 insertedCell = cell.notebook.cellAt(cellIndex); + + const { success } = await executeEphemeralCell(insertedCell, execution.token); + return success ? { success } : { success: false, error: new Error('Ephemeral cell execution failed') }; + }, + 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}\n`; + 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 executeAgentBlockFn(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.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 (isEphemeralCell(cell) && 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; + +export async function executeEphemeralCell( + 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(); + } + }) + ); + + 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')); + }, EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS); + + try { + const cellIndex = cell.index; + + await commands.executeCommand('notebook.cell.execute', { + ranges: [{ start: cellIndex, end: cellIndex + 1 }], + document: cell.notebook.uri + }); + + await completionDeferred.promise; + + return { + success: cell.executionSummary?.success === true, + outputs: cell.outputs.map(translateCellDisplayOutput), + executionCount: cell.executionSummary?.executionOrder ?? null + }; + } catch (error) { + return { + success: false, + outputs: [], + executionCount: null + }; + } finally { + dispose(disposables); + clearTimeout(timeout); + } +} + +async function removeEphemeralCellsForAgent(notebook: NotebookDocument, agentBlockId: string): Promise { + 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) { + 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}`); + } +} diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts new file mode 100644 index 0000000000..b124742244 --- /dev/null +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -0,0 +1,386 @@ +import { expect } from 'chai'; +import * as sinon from 'sinon'; +import { anything, capture, instance, mock, reset, when } from 'ts-mockito'; +import { + CancellationTokenSource, + Disposable, + EventEmitter, + ExtensionMode, + NotebookCellOutput, + NotebookCellOutputItem, + NotebookController, + 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' } } }); + + 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('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())); + }); + + teardown(() => { + disposables = dispose(disposables); + }); + + test('returns key when configured', async () => { + secretStorage.set('openAiApiKey', 'test-key'); + + const key = await getOpenAiApiKey(); + + 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'); + } + }); + }); + + suite('executeAgentCell', () => { + let mockExecution: { + appendOutput: sinon.SinonStub; + clearOutput: sinon.SinonStub; + end: sinon.SinonStub; + replaceOutput: sinon.SinonStub; + replaceOutputItems: sinon.SinonStub; + start: sinon.SinonStub; + }; + let mockController: NotebookController; + let executeAgentBlockStub: sinon.SinonStub; + + 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(); + }); + disposables.push(new Disposable(() => sinon.restore())); + + mockExecution = { + appendOutput: sinon.stub().resolves(), + 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; + + executeAgentBlockStub = sinon.stub().resolves({ finalOutput: 'done' } as AgentBlockResult); + }); + + teardown(() => { + disposables = dispose(disposables); + }); + + function createAgentCell(text: string = 'Test prompt') { + return createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text + }); + } + + test('creates execution and starts it', 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[]; + 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] Planning next steps...'); + }); + + 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 cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.replaceOutputItems.callCount).to.equal(2); + }); + + test('streaming chunks accumulate text progressively', 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 }); + + 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); + + expect(chunk1).to.include('[Agent] Text:'); + expect(chunk1).to.include('first'); + expect(chunk2).to.include('first second'); + }); + + 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' }); + + return { finalOutput: '' } as AgentBlockResult; + }); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + const getChunkText = (callIndex: number): string => { + const item = mockExecution.replaceOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; + + return Buffer.from(item.data).toString('utf-8'); + }; + + 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 = 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')); + + 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; + }); + + 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[]; + 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 = createAgentCell(''); + + 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('[Agent] Planning next steps...'); + }); + + test('ends with failure and writes error when API key is not set', async () => { + secretStorage.clear(); + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + 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('OpenAI API key is not set'); + }); + }); + + 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 + }); + }); + + 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(); + + 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 new file mode 100644 index 0000000000..0bbc73f3b7 --- /dev/null +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -0,0 +1,279 @@ +import { + CancellationToken, + Disposable, + EventEmitter, + NotebookCell, + NotebookCellStatusBarItem, + NotebookCellStatusBarItemProvider, + NotebookEdit, + WorkspaceEdit, + commands, + l10n, + notebooks, + window, + 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 { 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); + +/** + * 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( + 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); + } + + public dispose(): void { + for (const disposable of this.disposables) { + disposable.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; + 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; + } + + 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..62adc562cc --- /dev/null +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -0,0 +1,311 @@ +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 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', () => { + 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..6f9ebbd31c --- /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 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. + */ +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 || '', 'plaintext'); + + 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..a3ce26acf8 --- /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 plaintext 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, 'plaintext'); + }); + + 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, 'plaintext'); + }); + + 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, 'plaintext'); + }); + + 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, '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', () => { + 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', + 'plaintext' + ); + + 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, '', 'plaintext'); + + 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', '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.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/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/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)}`); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts index 8141e32093..824635343c 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.node'; 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 }; 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/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/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 { diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts new file mode 100644 index 0000000000..19ca8b76fc --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts @@ -0,0 +1,126 @@ +import { + Disposable, + NotebookCell, + NotebookDocument, + OverviewRulerLane, + Range, + TextEditor, + TextEditorDecorationType, + ThemeColor, + window, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { isEphemeralCell } from './dataConversionUtils'; +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 { + 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) { + 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); + } + } +} diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts new file mode 100644 index 0000000000..60c0a67fba --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts @@ -0,0 +1,80 @@ +import { + CancellationToken, + Disposable, + EventEmitter, + NotebookCell, + NotebookCellStatusBarItem, + NotebookCellStatusBarItemProvider, + l10n, + notebooks, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { isEphemeralCell } from './dataConversionUtils'; +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 (!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') + }; + } +} 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); + }); }