diff --git a/.github/workflows/e2e.yml b/.github/workflows/e2e.yml index a28c014d37..81f46448d7 100644 --- a/.github/workflows/e2e.yml +++ b/.github/workflows/e2e.yml @@ -71,6 +71,10 @@ jobs: - name: Install the Python extension into the test instance run: npm run setup:e2e:deps + - name: Pre-download the mock LLM server + # aimock is npx-only (not in lockfile), so setup-node never restores ~/.npm/_npx. + run: npm run setup:e2e:mock + - name: Cache pip wheel downloads # Provisioning the Deepnote environment pip-installs the toolkit dependency tree into a # fresh venv on first kernel connect — the bulk of the E2E runtime. Caching pip's wheel diff --git a/build/esbuild/build.ts b/build/esbuild/build.ts index 0f198363d9..9816dc8cd5 100644 --- a/build/esbuild/build.ts +++ b/build/esbuild/build.ts @@ -274,6 +274,17 @@ function createConfig( if (target === 'desktop') { alias['jsonc-parser'] = path.join(extensionFolder, 'node_modules', 'jsonc-parser', 'lib', 'esm', 'main.js'); } + // @deepnote/runtime-core needs Node built-ins (net, child_process) and is excluded from the VSIX; + // externalizing it (like desktop) would leave an unresolvable bare import in the web bundle. + if (target === 'web') { + alias['@deepnote/runtime-core'] = path.join( + extensionFolder, + 'src', + 'notebooks', + 'deepnote', + 'runtimeCore.web.ts' + ); + } // Desktop builds use CommonJS for VS Code/Cursor compatibility // Web builds use ESM for browser compatibility const config: SameShape = { diff --git a/cspell.json b/cspell.json index ad6ec1a49b..c7053314ae 100644 --- a/cspell.json +++ b/cspell.json @@ -50,6 +50,7 @@ "evalue", "findstr", "getsitepackages", + "hubot", "IMAGENAME", "ipykernel", "ipynb", diff --git a/package-lock.json b/package-lock.json index 20463d02e4..8f2fcbf25a 100644 --- a/package-lock.json +++ b/package-lock.json @@ -14,7 +14,7 @@ "@deepnote/blocks": "^4.6.0", "@deepnote/convert": "^4.0.0", "@deepnote/database-integrations": "^1.5.0", - "@deepnote/runtime-core": "^0.4.0", + "@deepnote/runtime-core": "^0.5.0", "@deepnote/sql-language-server": "^3.0.0", "@enonic/fnv-plus": "^1.3.0", "@jupyter-widgets/base": "^6.0.8", @@ -2165,14 +2165,14 @@ } }, "node_modules/@deepnote/runtime-core": { - "version": "0.4.0", - "resolved": "https://registry.npmjs.org/@deepnote/runtime-core/-/runtime-core-0.4.0.tgz", - "integrity": "sha512-iS5E2FUxAT83cRDt9cGvhO+CR9tuLG6+PmLT4Vm8ffQUzPSi5b0v0AuwnSOrMDPlt4SXM9DXxtB3iHhWUk/0kQ==", + "version": "0.5.0", + "resolved": "https://registry.npmjs.org/@deepnote/runtime-core/-/runtime-core-0.5.0.tgz", + "integrity": "sha512-+0Dbs5IhSsRFLg5XZqmr1gKo0GPeMwKNxezhAP6MgYmVECvxb+kHRoI5i4aazRsni/hk8i5AqipyYKYu4TT3Og==", "license": "Apache-2.0", "dependencies": { "@ai-sdk/mcp": "^1.0.25", "@ai-sdk/openai": "^3.0.0", - "@deepnote/blocks": "4.6.0", + "@deepnote/blocks": "4.7.0", "@jupyterlab/nbformat": "^4.3.2", "@jupyterlab/services": "^7.3.2", "ai": "^6.0.0", @@ -2181,6 +2181,21 @@ "zod": "3.25.76" } }, + "node_modules/@deepnote/runtime-core/node_modules/@deepnote/blocks": { + "version": "4.7.0", + "resolved": "https://registry.npmjs.org/@deepnote/blocks/-/blocks-4.7.0.tgz", + "integrity": "sha512-GW5jIpO2Sr7R2LaokuF8js60FHw4DLk1e4bq6BnMs2QkxGr2QF6JbRONl82V7zEV3iJQesXX+mFuTMWp9jCxPA==", + "license": "Apache-2.0", + "dependencies": { + "ts-dedent": "^2.2.0", + "yaml": "^2.8.3", + "zod": "3.25.76" + }, + "engines": { + "node": ">=22.14.0", + "pnpm": ">=10.17.1" + } + }, "node_modules/@deepnote/runtime-core/node_modules/ws": { "version": "8.21.0", "resolved": "https://registry.npmjs.org/ws/-/ws-8.21.0.tgz", @@ -36964,13 +36979,13 @@ } }, "@deepnote/runtime-core": { - "version": "0.4.0", - "resolved": "https://registry.npmjs.org/@deepnote/runtime-core/-/runtime-core-0.4.0.tgz", - "integrity": "sha512-iS5E2FUxAT83cRDt9cGvhO+CR9tuLG6+PmLT4Vm8ffQUzPSi5b0v0AuwnSOrMDPlt4SXM9DXxtB3iHhWUk/0kQ==", + "version": "0.5.0", + "resolved": "https://registry.npmjs.org/@deepnote/runtime-core/-/runtime-core-0.5.0.tgz", + "integrity": "sha512-+0Dbs5IhSsRFLg5XZqmr1gKo0GPeMwKNxezhAP6MgYmVECvxb+kHRoI5i4aazRsni/hk8i5AqipyYKYu4TT3Og==", "requires": { "@ai-sdk/mcp": "^1.0.25", "@ai-sdk/openai": "^3.0.0", - "@deepnote/blocks": "4.6.0", + "@deepnote/blocks": "4.7.0", "@jupyterlab/nbformat": "^4.3.2", "@jupyterlab/services": "^7.3.2", "ai": "^6.0.0", @@ -36979,6 +36994,16 @@ "zod": "3.25.76" }, "dependencies": { + "@deepnote/blocks": { + "version": "4.7.0", + "resolved": "https://registry.npmjs.org/@deepnote/blocks/-/blocks-4.7.0.tgz", + "integrity": "sha512-GW5jIpO2Sr7R2LaokuF8js60FHw4DLk1e4bq6BnMs2QkxGr2QF6JbRONl82V7zEV3iJQesXX+mFuTMWp9jCxPA==", + "requires": { + "ts-dedent": "^2.2.0", + "yaml": "2.8.3", + "zod": "3.25.76" + } + }, "ws": { "version": "8.21.0", "resolved": "https://registry.npmjs.org/ws/-/ws-8.21.0.tgz", diff --git a/package.json b/package.json index f75c7aebfd..3596fb8007 100644 --- a/package.json +++ b/package.json @@ -173,6 +173,12 @@ "category": "Deepnote", "icon": "$(notebook)" }, + { + "command": "deepnote.addAgentBlock", + "title": "%deepnote.commands.addAgentBlock.title%", + "category": "Deepnote", + "icon": "$(hubot)" + }, { "command": "deepnote.addSqlBlock", "title": "%deepnote.commands.addSqlBlock.title%", @@ -340,6 +346,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%", @@ -1022,30 +1038,35 @@ "when": "notebookType == 'deepnote'" }, { - "command": "deepnote.addSqlBlock", + "command": "deepnote.addAgentBlock", "group": "navigation@3", "when": "notebookType == 'deepnote'" }, { - "command": "deepnote.addChartBlock", + "command": "deepnote.addSqlBlock", "group": "navigation@4", "when": "notebookType == 'deepnote'" }, { - "command": "deepnote.addBigNumberChartBlock", + "command": "deepnote.addChartBlock", "group": "navigation@5", "when": "notebookType == 'deepnote'" }, { - "command": "deepnote.addInputBlock", + "command": "deepnote.addBigNumberChartBlock", "group": "navigation@6", "when": "notebookType == 'deepnote'" }, { - "command": "deepnote.addTextBlock", + "command": "deepnote.addInputBlock", "group": "navigation@7", "when": "notebookType == 'deepnote'" }, + { + "command": "deepnote.addTextBlock", + "group": "navigation@8", + "when": "notebookType == 'deepnote'" + }, { "command": "deepnote.restartkernel", "group": "navigation/execute@5", @@ -2666,7 +2687,8 @@ "compile-e2e-watch": "tsc -p ./test/e2e/tsconfig.json --watch", "setup:e2e:vscode": "extest get-vscode -c max && extest get-chromedriver -c max", "setup:e2e:deps": "extest install-from-marketplace ms-python.python -e .test-extensions", - "setup:e2e": "npm run setup:e2e:vscode && npm run setup:e2e:deps", + "setup:e2e:mock": "npx -y -p @copilotkit/aimock@1.37.4 llmock --help", + "setup:e2e": "npm run setup:e2e:vscode && npm run setup:e2e:deps && npm run setup:e2e:mock", "test:e2e": "extest setup-and-run \"./out/e2e/suite/*.e2e.test.js\" -c max -o ./test/e2e/settings.json -e .test-extensions -m ./test/e2e/.mocharc.js -i", "test:e2e:prebuilt": "extest run-tests \"./out/e2e/suite/*.e2e.test.js\" -c max -o ./test/e2e/settings.json -e .test-extensions -m ./test/e2e/.mocharc.js", "test:unittests": "mocha --config ./build/.mocha.unittests.js.json ./out/**/*.unit.test.js", @@ -2682,7 +2704,7 @@ "@deepnote/blocks": "^4.6.0", "@deepnote/convert": "^4.0.0", "@deepnote/database-integrations": "^1.5.0", - "@deepnote/runtime-core": "^0.4.0", + "@deepnote/runtime-core": "^0.5.0", "@deepnote/sql-language-server": "^3.0.0", "@enonic/fnv-plus": "^1.3.0", "@jupyter-widgets/base": "^6.0.8", diff --git a/package.nls.json b/package.nls.json index 58ea84148d..2cf9d50a05 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", @@ -258,6 +260,7 @@ "deepnote.commands.newProject.title": "New Project", "deepnote.commands.importNotebook.title": "Import Notebook", "deepnote.commands.importJupyterNotebook.title": "Import Jupyter Notebook", + "deepnote.commands.addAgentBlock.title": "Add Agent Block", "deepnote.commands.addSqlBlock.title": "Add SQL Block", "deepnote.commands.addBigNumberChartBlock.title": "Add Big Number Block", "deepnote.commands.addChartBlock.title": "Add Chart Block", diff --git a/src/kernels/execution/cellExecutionQueue.ts b/src/kernels/execution/cellExecutionQueue.ts index 66d78822e4..6b1597f1b7 100644 --- a/src/kernels/execution/cellExecutionQueue.ts +++ b/src/kernels/execution/cellExecutionQueue.ts @@ -325,10 +325,5 @@ export class CellExecutionQueue implements Disposable { break; } } - - // Notify listeners that execution queue is complete - if (this.notebook) { - notebookCellExecutions.notifyQueueComplete(this.notebook.uri.toString()); - } } } diff --git a/src/kernels/kernelExecution.ts b/src/kernels/kernelExecution.ts index 99b36ce964..aeca91137d 100644 --- a/src/kernels/kernelExecution.ts +++ b/src/kernels/kernelExecution.ts @@ -135,6 +135,10 @@ export class NotebookKernelExecution implements INotebookKernelExecution { } }); } + public get failed(): boolean { + return this.documentExecutions.get(this.notebook)?.failed === true; + } + public get pendingCells(): readonly NotebookCell[] { return this.documentExecutions.get(this.notebook)?.queue || []; } @@ -157,6 +161,8 @@ export class NotebookKernelExecution implements INotebookKernelExecution { .then(() => true) .catch(() => false); + notebookCellExecutions.notifyQueueComplete(cell.notebook.uri.toString()); + traceCellMessage( cell, `NotebookKernelExecution.resumeCellExecution (completed), ${getDisplayPath(cell.notebook.uri)}` diff --git a/src/kernels/types.ts b/src/kernels/types.ts index fbc44b2ebc..642ca583ff 100644 --- a/src/kernels/types.ts +++ b/src/kernels/types.ts @@ -461,6 +461,12 @@ export interface INotebookKernelExecution { * Total execution count on this kernel */ readonly executionCount: number; + /** + * Whether the cell execution queue stopped early: a cell failed, or queued cells were cancelled + * (interrupt/restart). Cancelled cell executions resolve rather than reject, so this is the only + * way to tell an interrupted run from a clean one. + */ + readonly failed: boolean; readonly onDidReceiveDisplayUpdate: Event; /** * Cells that are still being executed (or pending). diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index 38ad19e63b..6abf173d05 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -58,6 +58,7 @@ import { IJupyterVariablesProvider } from '../../kernels/variables/types'; import { IPyWidgetMessages } from '../../messageTypes'; import { ITelemetryService } from '../../platform/analytics/types'; import { IPythonExtensionChecker } from '../../platform/api/types'; +import { IEncryptedStorage } from '../../platform/common/application/types'; import { isCancellationError } from '../../platform/common/cancellation'; import { Commands, @@ -76,6 +77,7 @@ import { IExtensionContext } from '../../platform/common/types'; import { getNotebookMetadata, isJupyterNotebook, updateNotebookMetadata } from '../../platform/common/utils'; +import { notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; import { createDeferred } from '../../platform/common/utils/async'; import { DisposableStore, dispose } from '../../platform/common/utils/lifecycle'; import { Common, DataScience } from '../../platform/common/utils/localize'; @@ -92,6 +94,8 @@ import { RemoteKernelReconnectBusyIndicator } from './remoteKernelReconnectBusyI import { IConnectionDisplayData, IConnectionDisplayDataProvider, IVSCodeNotebookController } from './types'; import { notebookPathToDeepnoteProjectFilePath } from '../../platform/deepnote/deepnoteProjectUtils'; import { DEEPNOTE_NOTEBOOK_TYPE, IDeepnoteKernelAutoSelector } from '../../kernels/deepnote/types'; +import { executeAgentCell, removeEphemeralCellsForAgentBlocks } from '../deepnote/agentCellExecutionHandler'; +import { isAgentCell } from '../deepnote/dataConversionUtils'; /** * Our implementation of the VSCode Notebook Controller. Called by VS code to execute cells in a notebook. Also displayed @@ -607,6 +611,11 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont private handleInterrupt(notebook: NotebookDocument) { logger.debug(`VS Code interrupted kernel for ${getDisplayPath(notebook.uri)}`); notebook.getCells().forEach((cell) => traceCellMessage(cell, 'Cell cancellation requested')); + // Before the kernel interrupt: that interrupt ends the generated cell the agent is waiting on, + // and an agent that has not yet seen the stop reads the ended cell as a failure worth retrying. + // Setting an interruptHandler leaves NotebookCellExecution.token inert, so this is the agent's + // only stop signal. + this.agentCancellations.get(notebook)?.cancel(); commands .executeCommand(Commands.InterruptKernel, { notebookEditor: { notebookUri: notebook.uri } }) .then(noop, (ex) => logger.error('Failed to interrupt', ex)); @@ -630,23 +639,108 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont return currentExecution; } + /** Stop signal for the agent cell currently running in a notebook, read by `handleInterrupt`. */ + private readonly agentCancellations = new WeakMap(); private cellQueue = new WeakMap(); + /** + * Frames of `executeQueuedCells` in flight per notebook. An agent cell dispatches each cell it + * generates through `notebook.cell.execute`, which lands back here while the outer batch is still + * running; completion belongs to the gesture, so only the frame unwinding to zero announces it. + */ + private readonly executionDepth = new WeakMap(); private async executeQueuedCells(doc: NotebookDocument) { if (!this.cellQueue.has(doc)) { return; } + const queuedCells = this.cellQueue.get(doc) || []; + // Clear before await — agent runs can re-enter with an empty queue. + this.cellQueue.delete(doc); + // Nothing may run between here and the `try`: a throw in between would strand the depth above + // zero and silence completion — and with it the snapshot save — for the rest of the session. + const depthOnEntry = this.executionDepth.get(doc) ?? 0; + + this.executionDepth.set(doc, depthOnEntry + 1); + + if (depthOnEntry === 0) { + // Paired with the notify below so a run that opens no kernel queue — an agent cell that + // generates nothing — still marks a boundary the previous run's metadata is retired at. + notebookCellExecutions.notifyQueueStart(doc.uri.toString()); + } + + try { + const cellsToExecute = await removeEphemeralCellsForAgentBlocks(doc, queuedCells); + + let pendingKernelCells: NotebookCell[] = []; + + for (const cell of cellsToExecute) { + if (!isAgentCell(cell)) { + pendingKernelCells.push(cell); + continue; + } + + await this.executeKernelCells(doc, pendingKernelCells); + pendingKernelCells = []; + + logger.trace(`Executing agent cell ${cell.index} for ${getDisplayPath(doc.uri)} without kernel`); + + const agentCancellation = new CancellationTokenSource(); + + this.agentCancellations.set(doc, agentCancellation); + + try { + await executeAgentCell( + cell, + this.controller, + this.serviceContainer.get(IEncryptedStorage), + agentCancellation.token + ).catch(noop); + + // A stopped agent ends its own cell and returns, so it arrives here looking like a + // run that finished. Without this the cells after it would still execute. + if (agentCancellation.token.isCancellationRequested) { + throw new CancellationError(); + } + } finally { + this.agentCancellations.delete(doc); + agentCancellation.dispose(); + } + } + + await this.executeKernelCells(doc, pendingKernelCells); + } catch (ex) { + // The failing cell already carries the error; unwinding here only stops the agent cell and the + // segments after it, the way CellExecutionQueue stops a run once a cell fails. + logger.debug(`Stopped the rest of the batch for ${getDisplayPath(doc.uri)}`, ex); + } finally { + // Re-read rather than reuse a captured local: nested frames decrement the shared value. + const depth = (this.executionDepth.get(doc) ?? 1) - 1; + + this.executionDepth.set(doc, depth); + + if (depth === 0) { + notebookCellExecutions.notifyQueueComplete(doc.uri.toString()); + } + } + } + + private async executeKernelCells(doc: NotebookDocument, cells: NotebookCell[]) { // Start execution now (from the user's point of view) // Creating these execution objects marks the cell as queued for execution (vscode will update cell UI). type CellExec = { cell: NotebookCell; exec: NotebookCellExecution }; - const cellExecs: CellExec[] = (this.cellQueue.get(doc) || []).map((cell) => { + + // Stale handles use index -1 and abort the whole batch in createNotebookCellExecution. + const kernelCells = cells.filter((cell) => cell.index >= 0); + + if (kernelCells.length === 0) { + return; + } + + const cellExecs: CellExec[] = kernelCells.map((cell) => { const exec = this.createCellExecutionIfNecessary(cell, new KernelController(this.controller)); return { cell, exec }; }); - this.cellQueue.delete(doc); - const firstCell = cellExecs.length ? cellExecs[0].cell : undefined; - if (!firstCell) { - return; - } + + const firstCell = cellExecs[0].cell; logger.trace(`Execute Notebook ${getDisplayPath(doc.uri)}. Step 1`); @@ -676,12 +770,12 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont } catch (ex) { if (ex instanceof KernelError) { // Kernel errors would have been handled and displayed - return; + throw ex; } ex = WrappedError.unwrap(ex); if (ex instanceof CellExecutionOutputError) { // CellExecution already wrote this message to the cell output. - return; + throw ex; } if (!isCancellationError(ex)) { logger.error(`Error in notebook cell execution`, ex); @@ -700,10 +794,8 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont await errorHandler.getErrorMessageForDisplayInCellOutput(ex, currentContext, doc.uri), isCancelled ); - } - if (!kernel) { - return; + throw ex; } const kernelExecution = this.kernelProvider.getKernelExecution(kernel); @@ -729,11 +821,11 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont } catch (ex) { if (ex instanceof KernelError) { // Kernel errors would have been handled and displayed - return; + throw ex; } ex = WrappedError.unwrap(ex); if (ex instanceof CellExecutionOutputError) { - return; + throw ex; } if (!isCancellationError(ex)) { logger.error(`Error in cell execution`, ex); @@ -747,9 +839,17 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont await errorHandler.getErrorMessageForDisplayInCellOutput(ex, currentContext, doc.uri), isCancelled ); + + throw ex; } }) - ).catch(noop); + ); + + if (kernelExecution.failed) { + // An interrupt resolves the cells it cancelled, so the awaits above stay clean; the queue's + // own verdict is the only thing that separates a stopped run from a successful one. + throw new CancellationError(); + } } public async startKernel(notebook: NotebookDocument) { diff --git a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts index 0f55f5d9d1..412df10542 100644 --- a/src/notebooks/controllers/vscodeNotebookController.unit.test.ts +++ b/src/notebooks/controllers/vscodeNotebookController.unit.test.ts @@ -7,27 +7,32 @@ /* eslint-disable @typescript-eslint/no-require-imports, @typescript-eslint/no-var-requires */ import { assert } from 'chai'; import * as fakeTimers from '@sinonjs/fake-timers'; +import * as sinon from 'sinon'; import { - NotebookDocument, + Disposable, EventEmitter, NotebookCell, NotebookCellKind, NotebookController, - Uri, - Disposable + NotebookDocument, + Uri } from 'vscode'; import { VSCodeNotebookController, warnWhenUsingOutdatedPython } from './vscodeNotebookController'; import { IKernel, IKernelProvider, + INotebookKernelExecution, KernelConnectionMetadata, LiveRemoteKernelConnectionMetadata, LocalKernelConnectionMetadata, LocalKernelSpecConnectionMetadata, RemoteKernelSpecConnectionMetadata } from '../../kernels/types'; +import { KernelError } from '../../kernels/errors/kernelError'; +import { LastCellExecutionTracker } from '../../kernels/execution/lastCellExecutionTracker'; import { anything, deepEqual, instance, mock, verify, when } from 'ts-mockito'; import { ITelemetryService } from '../../platform/analytics/types'; +import { IEncryptedStorage } from '../../platform/common/application/types'; import { IConfigurationService, IDisposable, @@ -35,6 +40,7 @@ import { IWatchableJupyterSettings } from '../../platform/common/types'; import { dispose } from '../../platform/common/utils/lifecycle'; +import { ServiceContainer } from '../../platform/ioc/container'; import { NotebookCellLanguageService } from '../languages/cellLanguageService'; import { IServiceContainer } from '../../platform/ioc/types'; import { IJupyterServerProviderRegistry } from '../../kernels/jupyter/types'; @@ -52,6 +58,57 @@ import { mockedVSCode, mockedVSCodeNamespaces, resetVSCodeMocks } from '../../te import { Environment, PythonExtension } from '@vscode/python-extension'; import { crateMockedPythonApi, whenResolveEnvironment } from '../../kernels/helpers.unit.test'; import { IJupyterVariablesProvider } from '../../kernels/variables/types'; +import { notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; +import { logger } from '../../platform/logging'; +import { createMockNotebookWithCells } from '../deepnote/deepnoteTestHelpers'; + +// executeAgentCell takes IEncryptedStorage from the controller's container; getProjectAgentContext +// still resolves the notebook manager off the static one. +function stubAgentDependencies(serviceContainer: IServiceContainer, openAiApiKey: string): void { + const encryptedStorage = mock(); + const staticServiceContainer = instance(mock()); + + when(encryptedStorage.retrieve(anything(), anything())).thenResolve(openAiApiKey); + when(serviceContainer.get(IEncryptedStorage)).thenReturn(instance(encryptedStorage)); + sinon.stub(ServiceContainer, 'instance').get(() => staticServiceContainer); +} + +function installMockedCreateNotebookController( + onDidChangeSelectedNotebooksEvent: EventEmitter<{ + readonly notebook: NotebookDocument; + readonly selected: boolean; + }>['event'], + createNotebookCellExecution: NotebookController['createNotebookCellExecution'] = () => ({}) as any +): void { + (mockedVSCode as any).notebooks.createNotebookController = ( + _id: string, + _view: string, + _label: string, + executeHandler: any, + _rendererScripts: any + ) => { + return { + id: _id, + label: _label, + description: '', + detail: '', + supportedLanguages: [], + supportsExecutionOrder: false, + interruptHandler: undefined, + executeHandler, + onDidChangeSelectedNotebooks: onDidChangeSelectedNotebooksEvent, + onDidReceiveMessage: new EventEmitter().event, + dispose: () => {}, + asWebviewUri: (uri: Uri) => uri, + postMessage: () => Promise.resolve(true), + updateNotebookAffinity: () => {}, + createNotebookCellExecution, + createNotebookExecution: () => ({}) as any, + notebookType: _view, + rendererScripts: _rendererScripts || [] + } as NotebookController; + }; +} suite(`Notebook Controller`, function () { let controller: NotebookController; @@ -112,42 +169,7 @@ suite(`Notebook Controller`, function () { when(controller.label).thenReturn('Test Controller'); when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([]); when(mockedVSCodeNamespaces.workspace.onDidCloseNotebookDocument).thenReturn(onDidCloseNotebookDocument.event); - // Override just the createNotebookController method on the existing notebooks object - (mockedVSCode as any).notebooks.createNotebookController = ( - _id: string, - _view: string, - _label: string, - _handler: any, - _rendererScripts: any - ) => { - console.log('MOCK createNotebookController CALLED with id:', _id); - const mockControllerObject: any = { - id: _id, - label: _label, - description: '', - detail: '', - supportedLanguages: [], - supportsExecutionOrder: false, - interruptHandler: undefined, - executeHandler: _handler, - onDidChangeSelectedNotebooks: onDidChangeSelectedNotebooks.event, - onDidReceiveMessage: new EventEmitter().event, - dispose: () => {}, - asWebviewUri: (uri: Uri) => uri, - postMessage: () => Promise.resolve(true), - updateNotebookAffinity: () => {}, - createNotebookCellExecution: () => ({}) as any, - createNotebookExecution: () => ({}) as any, - notebookType: _view, - rendererScripts: _rendererScripts || [] - }; - console.log('MOCK createNotebookController RETURNING controller with id:', mockControllerObject.id); - return mockControllerObject; - }; - console.log( - 'mockedVSCode.notebooks.createNotebookController:', - typeof (mockedVSCode as any).notebooks.createNotebookController - ); + installMockedCreateNotebookController(onDidChangeSelectedNotebooks.event); when(languageService.getSupportedLanguages(anything())).thenReturn([PYTHON_LANGUAGE]); when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(true); when(mockedVSCodeNamespaces.workspace.onDidCloseNotebookDocument).thenReturn(onDidCloseNotebookDocument.event); @@ -903,4 +925,429 @@ suite(`Notebook Controller`, function () { }); }); }); + + suite('executeQueuedCells', function () { + let vscodeController: VSCodeNotebookController; + let notifyQueueCompleteSpy: sinon.SinonSpy; + let createNotebookCellExecutionStub: sinon.SinonStub; + let mockExecution: { + appendOutput: sinon.SinonStub; + clearOutput: sinon.SinonStub; + end: sinon.SinonStub; + replaceOutput: sinon.SinonStub; + appendOutputItems: sinon.SinonStub; + start: sinon.SinonStub; + token: { isCancellationRequested: boolean }; + }; + + setup(function () { + crateMockedPythonApi(disposables); + stubAgentDependencies(serviceContainer, 'test-key'); + when(serviceContainer.tryGet(anything())).thenReturn(undefined); + // handleExecution reports execute_notebook whenever the batch covers every code cell, + // which an agent-only batch does; these tests assert on the queue, not on telemetry. + when(serviceContainer.get(ITelemetryService)).thenReturn( + instance(mock()) + ); + + mockExecution = { + appendOutput: sinon.stub().resolves(), + clearOutput: sinon.stub().resolves(), + end: sinon.stub(), + replaceOutput: sinon.stub().resolves(), + appendOutputItems: sinon.stub().resolves(), + start: sinon.stub(), + token: { isCancellationRequested: false } + }; + createNotebookCellExecutionStub = sinon.stub().callsFake(() => { + mockExecution.end = sinon.stub(); + + return mockExecution; + }); + + installMockedCreateNotebookController(onDidChangeSelectedNotebooks.event, createNotebookCellExecutionStub); + + notifyQueueCompleteSpy = sinon.spy(notebookCellExecutions, 'notifyQueueComplete'); + + vscodeController = new VSCodeNotebookController( + instance(kernelConnection), + 'test-controller-id', + 'jupyter-notebook', + instance(kernelProvider), + instance(context), + disposables, + instance(languageService), + instance(configService), + instance(extensionChecker), + instance(serviceContainer), + displayDataProvider + ); + }); + + teardown(function () { + notifyQueueCompleteSpy.restore(); + sinon.restore(); + }); + + // The connected kernel is a plain object, not `instance(mock())`: a ts-mockito proxy + // answers `then` with a function, so awaiting the connect promise would never settle. + // Anything left unstubbed throws inside the per-cell try and would abort the batch for the wrong reason. + function stubKernelForExecution(kernelExecution: Partial): void { + const neverFires = () => new Disposable(() => undefined); + const connectedKernel = { + controller: { + id: 'test-controller-id', + createNotebookCellExecution: (cell: NotebookCell) => + vscodeController.controller.createNotebookCellExecution(cell) + }, + disposing: false, + onDisposed: neverFires, + onStatusChanged: neverFires + } as unknown as IKernel; + + const oldConnectToNotebook = KernelConnector.connectToNotebookKernel; + KernelConnector.connectToNotebookKernel = async () => connectedKernel; + disposables.push(new Disposable(() => (KernelConnector.connectToNotebookKernel = oldConnectToNotebook))); + + when(serviceContainer.get(LastCellExecutionTracker)).thenReturn( + instance(mock()) + ); + when(kernelProvider.getKernelExecution(anything())).thenReturn(kernelExecution as INotebookKernelExecution); + } + + test('a failed kernel segment stops the agent cell and the trailing segment', async function () { + // Catches: executeKernelCells swallowing a KernelError, so Run All continues past a failed cell. + const { + notebook, + cells: [failingCell, agentCell, trailingCell] + } = createMockNotebookWithCells([ + { metadata: { id: 'code-1' }, text: 'raise ValueError()' }, + { metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, text: 'Test prompt' }, + { metadata: { id: 'code-2' }, text: 'print(2)' } + ]); + + const executedIndexes: number[] = []; + stubKernelForExecution({ + failed: false, + executeCell: async (cell: NotebookCell) => { + executedIndexes.push(cell.index); + throw new KernelError({ ename: 'ValueError', evalue: 'boom', traceback: [] }); + } + }); + + await vscodeController.controller.executeHandler( + [failingCell, agentCell, trailingCell], + notebook, + vscodeController.controller + ); + + assert.deepStrictEqual(executedIndexes, [0], 'the trailing segment must not run after a failure'); + assert.isFalse( + createNotebookCellExecutionStub.getCalls().some((call) => call.args[0] === agentCell), + 'the agent cell must not start after a failed segment' + ); + }); + + test('an interrupted kernel segment stops the agent cell even though its cells resolve', async function () { + // Catches: relying on a rejection alone - cancelled cell executions resolve + // (CellExecution.completedDueToCancellation), so only the queue's verdict shows the interrupt. + const { + notebook, + cells: [interruptedCell, agentCell, trailingCell] + } = createMockNotebookWithCells([ + { metadata: { id: 'code-1' }, text: 'time.sleep(30)' }, + { metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, text: 'Test prompt' }, + { metadata: { id: 'code-2' }, text: 'print(2)' } + ]); + + const executedIndexes: number[] = []; + stubKernelForExecution({ + failed: true, + executeCell: async (cell: NotebookCell) => { + executedIndexes.push(cell.index); + } + }); + + await vscodeController.controller.executeHandler( + [interruptedCell, agentCell, trailingCell], + notebook, + vscodeController.controller + ); + + assert.deepStrictEqual(executedIndexes, [0], 'the trailing segment must not run after an interrupt'); + assert.isFalse( + createNotebookCellExecutionStub.getCalls().some((call) => call.args[0] === agentCell), + 'the agent cell must not start after an interrupt' + ); + }); + + test('an interrupt during the agent cell stops the trailing segment', async function () { + // Catches: reading the stop from a rejection - executeAgentCell reports one by ending its + // cell and returning, so the batch sees a run that finished normally. + const { + notebook, + cells: [agentCell, trailingCell] + } = createMockNotebookWithCells([ + { metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, text: 'Test prompt' }, + { metadata: { id: 'code-1' }, text: 'print(1)' } + ]); + + const executedIndexes: number[] = []; + stubKernelForExecution({ + failed: false, + executeCell: async (cell: NotebookCell) => { + executedIndexes.push(cell.index); + } + }); + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(undefined); + + const { interruptHandler } = vscodeController.controller; + + if (!interruptHandler) { + assert.fail('the controller must install an interrupt handler for Stop to reach the agent'); + } + + // The agent's execution object is created as its run starts, which is where Stop lands. + createNotebookCellExecutionStub.callsFake((cell: NotebookCell) => { + if (cell === agentCell) { + Promise.resolve(interruptHandler(notebook)).catch((ex) => + logger.error('Failed to interrupt the notebook', ex) + ); + } + mockExecution.end = sinon.stub(); + + return mockExecution; + }); + + await vscodeController.controller.executeHandler( + [agentCell, trailingCell], + notebook, + vscodeController.controller + ); + + assert.deepStrictEqual(executedIndexes, [], 'the trailing segment must not run after an interrupt'); + }); + + test('a clean kernel segment still runs the agent cell and the trailing segment', async function () { + // Catches: aborting the batch when nothing failed (e.g. consulting the queue verdict too early). + const { + notebook, + cells: [firstCell, agentCell, trailingCell] + } = createMockNotebookWithCells([ + { metadata: { id: 'code-1' }, text: 'x = 1' }, + { metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, text: 'Test prompt' }, + { metadata: { id: 'code-2' }, text: 'print(2)' } + ]); + + const executedIndexes: number[] = []; + stubKernelForExecution({ + failed: false, + executeCell: async (cell: NotebookCell) => { + executedIndexes.push(cell.index); + } + }); + + await vscodeController.controller.executeHandler( + [firstCell, agentCell, trailingCell], + notebook, + vscodeController.controller + ); + + assert.deepStrictEqual(executedIndexes, [0, 2], 'both kernel segments must run when nothing failed'); + assert.isTrue( + createNotebookCellExecutionStub.getCalls().some((call) => call.args[0] === agentCell), + 'the agent cell must run between the segments' + ); + }); + + test('agent-only batch fires notifyQueueComplete (arms deferred snapshot save)', async function () { + // Catches: agent-only runs never reach CellExecutionQueue, so snapshot save never arms. + const { + notebook: agentNotebook, + cells: [agentCell] + } = createMockNotebookWithCells([ + { + metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, + text: 'Test prompt' + } + ]); + + const notebookUri = agentNotebook.uri.toString(); + + let queueCompletionUri: string | undefined; + const queueListener = notebookCellExecutions.onDidCompleteQueueExecution((event) => { + queueCompletionUri = event.notebookUri; + }); + disposables.push(new Disposable(() => queueListener.dispose())); + + const executeHandler = vscodeController.controller.executeHandler; + assert.isDefined(executeHandler); + + await executeHandler([agentCell], agentNotebook, vscodeController.controller); + + assert.isTrue(notifyQueueCompleteSpy.calledOnce, 'notifyQueueComplete must run after agent-only execution'); + assert.strictEqual(notifyQueueCompleteSpy.firstCall.args[0], notebookUri); + assert.strictEqual(queueCompletionUri, notebookUri); + assert.isTrue(createNotebookCellExecutionStub.calledOnce, 'agent cell should run through executeAgentCell'); + }); + + test('an agent batch and a later kernel-only batch each fire exactly one completion', async function () { + // Catches: tying completion to an "this batch ran an agent cell" flag — CellExecutionQueue no + // longer notifies, so a plain Run would arm no deferred snapshot save. + const { + notebook: agentNotebook, + cells: [agentCell, codeCell] + } = createMockNotebookWithCells([ + { + metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, + text: 'Test prompt' + }, + { + metadata: { id: 'code-1' }, + text: 'print(1)' + } + ]); + + stubKernelForExecution({ failed: false, executeCell: async () => undefined }); + + await vscodeController.controller.executeHandler([agentCell], agentNotebook, vscodeController.controller); + await vscodeController.controller.executeHandler([codeCell], agentNotebook, vscodeController.controller); + + assert.deepStrictEqual( + notifyQueueCompleteSpy.getCalls().map((call) => call.args[0]), + [agentNotebook.uri.toString(), agentNotebook.uri.toString()], + 'each gesture fires one completion for its own notebook' + ); + }); + + test('a kernel-only batch fires the completion itself', async function () { + // Catches: leaving completion to CellExecutionQueue, which no longer announces it — the + // deferred snapshot save would never arm for an ordinary run. + const { + notebook: codeNotebook, + cells: [codeCell] + } = createMockNotebookWithCells([ + { + metadata: { id: 'code-block-1' }, + text: 'x = 1' + } + ]); + + stubKernelForExecution({ failed: false, executeCell: async () => undefined }); + + await vscodeController.controller.executeHandler([codeCell], codeNotebook, vscodeController.controller); + + assert.isTrue(notifyQueueCompleteSpy.calledOnce, 'a kernel-only batch must fire one completion'); + assert.strictEqual(notifyQueueCompleteSpy.firstCall.args[0], codeNotebook.uri.toString()); + }); + + test('a run that re-enters once per generated cell fires exactly one completion', async function () { + // Catches the N+1 snapshot saves an agent run produced: each generated cell is dispatched + // through `notebook.cell.execute`, which lands back in this handler while the outer batch is + // still in flight. Completion belongs to the gesture, not to every nested run. + const { + notebook, + cells: [runCell, generatedFirst, generatedSecond] + } = createMockNotebookWithCells([ + { metadata: { id: 'code-1' }, text: 'x = 1' }, + { metadata: { id: 'generated-1' }, text: 'print(1)' }, + { metadata: { id: 'generated-2' }, text: 'print(2)' } + ]); + + const executedIds: string[] = []; + stubKernelForExecution({ + failed: false, + executeCell: async (cell: NotebookCell) => { + executedIds.push(cell.metadata.id as string); + + if (cell !== runCell) { + return; + } + + await vscodeController.controller.executeHandler( + [generatedFirst], + notebook, + vscodeController.controller + ); + await vscodeController.controller.executeHandler( + [generatedSecond], + notebook, + vscodeController.controller + ); + } + }); + + await vscodeController.controller.executeHandler([runCell], notebook, vscodeController.controller); + + assert.deepStrictEqual( + executedIds, + ['code-1', 'generated-1', 'generated-2'], + 'both re-entrant runs must have executed' + ); + assert.strictEqual( + notifyQueueCompleteSpy.callCount, + 1, + 'the gesture owns the completion; the runs nested inside it must not fire their own' + ); + }); + + test('a batch that fails still fires its own completion and does not silence the next one', async function () { + // Catches: skipping the completion when the batch unwinds — an interrupted run must still + // snapshot what it produced, and must not leave the re-entrancy depth above zero. + const { + notebook, + cells: [failingCell, laterCell] + } = createMockNotebookWithCells([ + { metadata: { id: 'code-1' }, text: 'raise ValueError()' }, + { metadata: { id: 'code-2' }, text: 'x = 1' } + ]); + + stubKernelForExecution({ + failed: false, + executeCell: async (cell: NotebookCell) => { + if (cell === failingCell) { + throw new KernelError({ ename: 'ValueError', evalue: 'boom', traceback: [] }); + } + } + }); + + await vscodeController.controller.executeHandler([failingCell], notebook, vscodeController.controller); + await vscodeController.controller.executeHandler([laterCell], notebook, vscodeController.controller); + + assert.strictEqual(notifyQueueCompleteSpy.callCount, 2, 'a failed batch must not silence later runs'); + }); + + test('a rejected scratch-cell cleanup neither escapes nor strands the completion', async function () { + // Catches: clearing prior scratch cells outside the frame that owns the re-entrancy depth — + // a rejected workspace edit would escape before the depth is handed back, silencing every + // completion for the rest of the session. + const { + notebook, + cells: [agentCell, , laterCell] + } = createMockNotebookWithCells([ + { metadata: { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }, text: 'Test prompt' }, + { + metadata: { agent_source_block_id: 'agent-block-1', id: 'scratch-1', is_ephemeral: true }, + text: 'print(1)' + }, + { metadata: { id: 'code-1' }, text: 'x = 1' } + ]); + + stubKernelForExecution({ failed: false, executeCell: async () => undefined }); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReject(new Error('edit failed')); + + let escaped: unknown; + + try { + await vscodeController.controller.executeHandler([agentCell], notebook, vscodeController.controller); + } catch (ex) { + escaped = ex; + } + + await vscodeController.controller.executeHandler([laterCell], notebook, vscodeController.controller); + + assert.isUndefined(escaped, 'a rejected cleanup edit must not escape the execute handler'); + assert.strictEqual(notifyQueueCompleteSpy.callCount, 2, 'a cleanup failure must not strand the depth'); + }); + }); }); diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts new file mode 100644 index 0000000000..d4f9632d25 --- /dev/null +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -0,0 +1,527 @@ +import { + CancellationError, + CancellationToken, + NotebookCell, + NotebookCellData, + NotebookCellOutput, + NotebookCellOutputItem, + NotebookController, + NotebookDocument, + NotebookEdit, + NotebookRange, + WorkspaceEdit, + commands, + workspace +} from 'vscode'; + +import { AgentBlock, DeepnoteBlock, extractOutputsText } from '@deepnote/blocks'; +import { + AgentBlockContext, + AgentStreamEvent, + executeAgentBlock, + serializeNotebookContextFromBlocks +} from '@deepnote/runtime-core'; + +import { translateCellDisplayOutput } from '../../kernels/execution/helpers'; +import { IEncryptedStorage } from '../../platform/common/application/types'; +import type { IDisposable } from '../../platform/common/types'; +import { createDeferred } from '../../platform/common/utils/async'; +import { dispose } from '../../platform/common/utils/lifecycle'; +import { uuidUtils } from '../../platform/common/uuid'; +import { ServiceContainer } from '../../platform/ioc/container'; +import { logger } from '../../platform/logging'; +import { Cancellation } from '../../platform/common/cancellation'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; +import { IDeepnoteNotebookManager } from '../types'; +import { + generateBlockId, + generateSortingKey, + getBlockId, + getEphemeralCellAgentSourceBlockId, + isAgentCell +} from './dataConversionUtils'; +import { DeepnoteDataConverter } from './deepnoteDataConverter'; +import { getOrPromptOpenAiApiKey } from './deepnoteSecretStore'; + +/** Project MCP servers and integrations from the `.deepnote` file (CLI ExecutionEngine parity). Callers must gate on `workspace.isTrusted` — MCP spawn is arbitrary command execution. */ +function getProjectAgentContext(notebook: NotebookDocument): Pick { + const projectId = notebook.metadata?.deepnoteProjectId as string | undefined; + const notebookId = notebook.metadata?.deepnoteNotebookId as string | undefined; + + if (!projectId || !notebookId) { + return { mcpServers: [], integrations: [] }; + } + + const manager = ServiceContainer.instance.tryGet(IDeepnoteNotebookManager); + const project = manager?.getProjectForNotebook(projectId, notebookId)?.project; + const mcpServers = project?.settings?.mcpServers ?? []; + const integrations = project?.integrations ?? []; + + if (mcpServers.length > 0) { + logger.info( + `Agent cell: using ${mcpServers.length} project MCP server(s): ${mcpServers.map((s) => s.name).join(', ')}` + ); + } + + if (integrations.length > 0) { + logger.info( + `Agent cell: using ${integrations.length} project integration(s): ${integrations + .map((i) => i.name) + .join(', ')}` + ); + } + + return { mcpServers, integrations }; +} + +// Wording matches @deepnote/runtime-core ExecutionEngine so backend and extension runs look the same to the agent. +const MARKDOWN_BLOCK_ADDED_TEXT = 'Markdown block added.'; +const NO_OUTPUT_TEXT = '(no output)'; + +function notebookCellDataFromCell(cell: NotebookCell): NotebookCellData { + return { + kind: cell.kind, + value: cell.document.getText(), + languageId: cell.document.languageId, + metadata: cell.metadata, + outputs: [...(cell.outputs || [])] + }; +} + +function toError(error: unknown): Error { + return error instanceof Error ? error : new Error(String(error)); +} + +export function serializeNotebookContext({ + cells, + notebookName +}: { + cells: NotebookCell[]; + notebookName: string; +}): string { + const converter = new DeepnoteDataConverter(); + + const blocks = cells.reduce((acc, cell) => { + try { + const block = converter.convertCellToBlock(notebookCellDataFromCell(cell), cell.index); + acc.push(block); + } catch (error) { + logger.error(`Error converting cell to block: ${error}`); + } + return acc; + }, []); + + return serializeNotebookContextFromBlocks({ blocks, notebookName }); +} + +function joinMultilineString(value: unknown): unknown { + return Array.isArray(value) && value.every((entry) => typeof entry === 'string') ? value.join('') : value; +} + +/** Join nbformat line arrays before `extractOutputsText` — `String(array)` inserts commas between lines. */ +function normalizeOutputsForTextExtraction(outputs: unknown[]): unknown[] { + return outputs.map((output) => { + const candidate = output as { output_type?: unknown; text?: unknown; data?: unknown } | null; + + if (candidate?.output_type === 'stream') { + return { ...candidate, text: joinMultilineString(candidate.text) }; + } + + if ( + (candidate?.output_type === 'execute_result' || candidate?.output_type === 'display_data') && + candidate.data != null && + typeof candidate.data === 'object' + ) { + const data = Object.fromEntries( + Object.entries(candidate.data).map(([mime, value]) => [ + mime, + mime.startsWith('text/') ? joinMultilineString(value) : value + ]) + ); + + return { ...candidate, data }; + } + + return output; + }); +} + +export function describeExecutionOutputs(outputs: unknown[]): string { + return extractOutputsText(normalizeOutputsForTextExtraction(outputs), { includeTraceback: true }) || NO_OUTPUT_TEXT; +} + +export interface ExecuteAgentCellOptions { + executeAgentBlockFn?: typeof executeAgentBlock; +} + +/** + * True for both the host's own cancellation and the `AbortError` that runtime-core raises from + * `signal.throwIfAborted()`. `isCancellationError` covers only the former. + */ +function isStopped(error: unknown): boolean { + return error instanceof CancellationError || (error instanceof Error && error.name === 'AbortError'); +} + +/** + * Runs an agent block into the cell output and inserts generated cells below. + * Call `removeEphemeralCellsForAgentBlocks` on the batch first. Never rejects — errors become stderr on the cell. + * + * `token` stops the run, bridged to the `AbortSignal` runtime-core forwards to `agent.stream`, so the + * in-flight model request is aborted rather than left to finish. Throwing from a tool callback cannot + * stop it: runtime-core catches that and hands the model an `Execution error: …` string to retry. + */ +export async function executeAgentCell( + cell: NotebookCell, + controller: NotebookController, + encryptedStorage: IEncryptedStorage, + token: CancellationToken, + options?: ExecuteAgentCellOptions +): Promise { + const executeAgentBlockFn = options?.executeAgentBlockFn ?? executeAgentBlock; + const execution = controller.createNotebookCellExecution(cell); + const stopController = new AbortController(); + const stopSubscription = token.onCancellationRequested(() => stopController.abort()); + + // The agent runs off the kernel, so nothing announces it on the internal shim — the only source + // SnapshotService and the execute_cell analytics read. Without this the run is invisible to both. + const endExecution = (success: boolean) => { + stopSubscription.dispose(); + execution.end(success, Date.now()); + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Idle); + }; + + execution.start(Date.now()); + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Executing); + + try { + await execution.clearOutput(); + + const prompt = cell.document.getText(); + + // runtime-core awaits each event; append deltas only (O(n) over the EH boundary). + const output = new NotebookCellOutput([NotebookCellOutputItem.stdout(`[Agent] Planning next steps...`)]); + await execution.replaceOutput([output]); + + const dataConverter = new DeepnoteDataConverter(); + const deepnoteBlock = dataConverter.convertCellToBlock(notebookCellDataFromCell(cell), cell.index); + const agentBlock: AgentBlock | null = deepnoteBlock.type === 'agent' ? deepnoteBlock : null; + + if (agentBlock == null) { + throw new Error('Cell is not an agent cell'); + } + + const staleCellCount = cell.notebook + .getCells() + .filter((c) => getEphemeralCellAgentSourceBlockId(c) === agentBlock.id).length; + + if (staleCellCount > 0) { + throw new Error( + `Agent block ${agentBlock.id} still has ${staleCellCount} generated cell(s) from its previous run` + ); + } + + const openAiToken = await getOrPromptOpenAiApiKey(encryptedStorage); + + let lastAgentEventType: AgentStreamEvent['type'] | undefined; + + // Caller must clear scratch cells before the batch; context serialization does not filter them. + const notebookContext = serializeNotebookContext({ + cells: cell.notebook.getCells().filter((c) => c.index !== cell.index), + notebookName: (cell.notebook.metadata?.deepnoteNotebookName as string | undefined) ?? '' + }); + + const context: AgentBlockContext = { + openAiToken, + ...getProjectAgentContext(cell.notebook), + notebookContext, + signal: stopController.signal, + addMarkdownBlock: async ({ content }: { content: string }) => { + Cancellation.throwIfCanceled(token); + + try { + await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'markdown', content); + + return MARKDOWN_BLOCK_ADDED_TEXT; + } catch (error) { + return `Failed to add markdown block: ${toError(error).message}`; + } + }, + addAndExecuteCodeBlock: async ({ code }: { code: string }) => { + Cancellation.throwIfCanceled(token); + + try { + const insertedCell = await insertEphemeralCell( + cell.notebook, + cell.index, + agentBlock.id, + 'code', + code + ); + + const { success, outputs, error } = await executeEphemeralCell(insertedCell, token); + const outputText = error ?? describeExecutionOutputs(outputs); + + return success ? `Output:\n${outputText}` : `Execution failed:\n${outputText}`; + } catch (error) { + if (isStopped(error)) { + throw error; + } + + return `Execution error: ${toError(error).message}`; + } + }, + onAgentEvent: async (event: AgentStreamEvent) => { + // Runs in runtime-core's own stream loop, which has no catch — the one place the host + // can end the run rather than merely refuse it. + Cancellation.throwIfCanceled(token); + + logger.trace(`Agent event: ${event.type}`); + + let delta = lastAgentEventType != null && lastAgentEventType !== event.type ? `\n\n` : ''; + + switch (event.type) { + case 'tool_called': + delta += `[Agent] Tool called: ${event.toolName}`; + break; + case 'tool_output': + delta += `[Agent] Tool output: ${event.toolName}\n`; + delta += `[Agent] Tool output length: ${event.output?.length}`; + break; + case 'text_delta': + if (lastAgentEventType !== 'text_delta') { + delta += `[Agent] Text:\n`; + } + delta += event.text; + break; + case 'reasoning_delta': + if (lastAgentEventType !== 'reasoning_delta') { + delta += `[Agent] Reasoning:\n`; + } + delta += event.text; + break; + default: + event satisfies never; + } + lastAgentEventType = event.type; + + await execution.appendOutputItems(NotebookCellOutputItem.stdout(delta), 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}`); + + endExecution(true); + } catch (error) { + if (isStopped(error)) { + logger.info('Agent cell execution stopped'); + + const stoppedOutput = new NotebookCellOutput([NotebookCellOutputItem.stderr('[Agent] Stopped')]); + + await execution.appendOutput([stoppedOutput]).then(undefined, () => undefined); + endExecution(false); + + return; + } + + // logger.error does not print stacks unless isJupyterError — log stack explicitly. + logger.error('Agent cell execution failed', error); + if (error instanceof Error) { + if (error.cause) { + 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); + endExecution(false); + } +} + +function getInsertIndexAfterAgentCell( + notebook: NotebookDocument, + agentCellIndex: number, + agentBlockId: string +): number { + let index = agentCellIndex + 1; + + while (index < notebook.cellCount) { + if (getEphemeralCellAgentSourceBlockId(notebook.cellAt(index)) === agentBlockId) { + index++; + } else { + break; + } + } + + return index; +} + +/** Inserts an ephemeral cell after the agent; returns the created cell resolved by block id (`cellAt` clamps bad indices). */ +async function insertEphemeralCell( + notebook: NotebookDocument, + agentCellIndex: number, + 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])]); + + if (!(await workspace.applyEdit(edit))) { + throw new Error(`Failed to insert ephemeral ${blockType} cell for agent block ${agentBlockId}`); + } + + const insertedCell = notebook.getCells().find((c) => getBlockId(c) === block.id); + + if (!insertedCell) { + throw new Error(`Inserted ephemeral ${blockType} cell ${block.id} not found in notebook`); + } + + return insertedCell; +} + +export const EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS = 5 * 60 * 1000; + +export interface EphemeralCellExecutionResult { + success: boolean; + outputs: unknown[]; + executionCount: number | null; + /** Why the run failed, when the failure wasn't the cell's own output (cancellation, timeout). */ + error?: string; +} + +export async function executeEphemeralCell( + cell: NotebookCell, + token: CancellationToken +): Promise { + // Cancel before dispatch — a rejected wait alone still runs the cell in the kernel. + if (token.isCancellationRequested) { + throw new CancellationError(); + } + + const completionDeferred = createDeferred(); + const disposables: IDisposable[] = []; + + disposables.push( + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell === cell && e.state === NotebookCellExecutionState.Idle) { + completionDeferred.resolve(); + } + }) + ); + + 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; + + // Race dispatch with Idle wait — command can hang while the cell is still running. + await Promise.all([ + commands.executeCommand('notebook.cell.execute', { + ranges: [{ start: cellIndex, end: cellIndex + 1 }], + document: cell.notebook.uri + }), + completionDeferred.promise + ]); + + return { + success: cell.executionSummary?.success === true, + outputs: cell.outputs.map(translateCellDisplayOutput), + executionCount: cell.executionSummary?.executionOrder ?? null + }; + } catch (error) { + if (error instanceof CancellationError) { + throw error; + } + + // Surface timeout/cancel reason — "(no output)" makes the agent retry while the cell still runs. + return { + success: false, + outputs: [], + executionCount: null, + error: error instanceof Error ? error.message : String(error) + }; + } finally { + dispose(disposables); + clearTimeout(timeout); + } +} + +/** + * Removes prior-run scratch cells owned by agents in `cells` and returns the batch without them. + * Required before `executeAgentCell` — otherwise stale generated code would run. Only agents in + * `cells` are scoped so standalone ephemeral runs stay untouched. Edit failures are logged, not thrown. + */ +export async function removeEphemeralCellsForAgentBlocks( + notebook: NotebookDocument, + cells: NotebookCell[] +): Promise { + const agentBlockIds = new Set( + cells + .filter(isAgentCell) + .map(getBlockId) + .filter((id): id is string => typeof id === 'string') + ); + + if (agentBlockIds.size === 0) { + return cells; + } + + const isOwnedScratch = (cell: NotebookCell) => { + const owner = getEphemeralCellAgentSourceBlockId(cell); + + return owner !== undefined && agentBlockIds.has(owner); + }; + + const remainingCells = cells.filter((cell) => !isOwnedScratch(cell)); + const deletions: NotebookEdit[] = []; + + for (let i = notebook.cellCount - 1; i >= 0; i--) { + if (isOwnedScratch(notebook.cellAt(i))) { + deletions.push(NotebookEdit.deleteCells(new NotebookRange(i, i + 1))); + } + } + + if (deletions.length === 0) { + return remainingCells; + } + + const edit = new WorkspaceEdit(); + edit.set(notebook.uri, deletions); + + if (await workspace.applyEdit(edit)) { + logger.info(`Removed ${deletions.length} ephemeral cell(s) for ${agentBlockIds.size} agent block(s)`); + } else { + logger.error(`Failed to remove ephemeral cells for agent blocks ${[...agentBlockIds].join(', ')}`); + } + + return remainingCells; +} diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts new file mode 100644 index 0000000000..fe18fe187f --- /dev/null +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -0,0 +1,794 @@ +import { expect } from 'chai'; +import * as sinon from 'sinon'; +import { anything, capture, instance, mock, reset, verify, when } from 'ts-mockito'; +import { + CancellationError, + CancellationToken, + CancellationTokenSource, + Disposable, + NotebookCell, + NotebookCellData, + NotebookCellOutput, + NotebookCellOutputItem, + NotebookController, + NotebookDocument, + WorkspaceEdit +} from 'vscode'; + +import type { AgentBlock } from '@deepnote/blocks'; +import type { AgentBlockContext } from '@deepnote/runtime-core'; + +import { IEncryptedStorage } from '../../platform/common/application/types'; +import type { IDisposable } from '../../platform/common/types'; +import { dispose } from '../../platform/common/utils/lifecycle'; +import { ServiceContainer } from '../../platform/ioc/container'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; +import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; +import { + describeExecutionOutputs, + EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS, + executeAgentCell, + executeEphemeralCell, + removeEphemeralCellsForAgentBlocks +} from './agentCellExecutionHandler'; +import { IDeepnoteNotebookManager } from '../types'; +import { createDeepnoteFile, createDeepnoteProject, createMockCell, createMockNotebook } from './deepnoteTestHelpers'; + +// Key namespacing is EncryptedStorage's job and is covered in deepnoteSecretStore.unit.test.ts. +function createEncryptedStorageFake(secretStorage: Map): IEncryptedStorage { + const encryptedStorage = mock(); + + when(encryptedStorage.store(anything(), anything(), anything())).thenCall( + (_service: string, key: string, value: string | undefined) => { + if (value === undefined) { + secretStorage.delete(key); + } else { + secretStorage.set(key, value); + } + + return Promise.resolve(); + } + ); + when(encryptedStorage.retrieve(anything(), anything())).thenCall((_service: string, key: string) => + Promise.resolve(secretStorage.get(key)) + ); + + return instance(encryptedStorage); +} + +// getProjectAgentContext still resolves the notebook manager off the static container. +function stubServiceContainerInstance(): ServiceContainer { + const serviceContainer = mock(); + + sinon.stub(ServiceContainer, 'instance').get(() => instance(serviceContainer)); + + return serviceContainer; +} + +// Mocked WorkspaceEdit.set drops edits — capture on prototype and apply to `cells`. +// `appliedEdits` is local because the shared workspace mock is never reset between tests. +function applyNotebookEditsTo(cells: NotebookCell[], notebook: NotebookDocument) { + type RecordedEdit = { range: { start: number; end: number }; newCells?: NotebookCellData[] }; + let recordedEdits: RecordedEdit[] = []; + let appliedEdits = 0; + + sinon.stub(WorkspaceEdit.prototype, 'set').callsFake((_uri, edits) => { + recordedEdits = edits as unknown as RecordedEdit[]; + }); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { + appliedEdits++; + + for (const notebookEdit of recordedEdits) { + const { start, end } = notebookEdit.range; + const deleteCount = end - start; + const newCellData = notebookEdit.newCells; + + if (!newCellData || newCellData.length === 0) { + if (deleteCount > 0) { + cells.splice(start, deleteCount); + } + continue; + } + + const inserted = newCellData.map((cellData) => { + const created = createMockCell({ + text: cellData.value, + metadata: cellData.metadata + }); + (created as { notebook: NotebookDocument }).notebook = notebook; + + return created; + }); + + cells.splice(start, deleteCount, ...inserted); + } + cells.forEach((cell, index) => ((cell as { index: number }).index = index)); + recordedEdits = []; + + return Promise.resolve(true); + }); + + return { appliedEdits: () => appliedEdits }; +} + +suite('AgentCellExecutionHandler', () => { + const secretStorage = new Map(); + let disposables: IDisposable[] = []; + + suite('describeExecutionOutputs', () => { + test('joins nbformat line arrays in stream text', () => { + const output = { + output_type: 'stream', + name: 'stdout', + text: ['hello\n', 'world\n'] + }; + + expect(describeExecutionOutputs([output])).to.equal('hello\nworld\n'); + }); + + // String(line[]) joins with commas — breaks DataFrame text/plain for the agent. + test('joins nbformat line arrays in execute_result and display_data text/plain', () => { + const executeResult = { + output_type: 'execute_result', + data: { 'text/plain': [' a b\n', '0 1 4\n', '1 2 5'] }, + metadata: {}, + execution_count: 1 + }; + const displayData = { + output_type: 'display_data', + data: { 'text/plain': ['line one\n', 'line two'] }, + metadata: {} + }; + + expect(describeExecutionOutputs([executeResult])).to.equal(' a b\n0 1 4\n1 2 5'); + expect(describeExecutionOutputs([displayData])).to.equal('line one\nline two'); + }); + + test('leaves single-line text/plain untouched', () => { + const output = { + output_type: 'execute_result', + data: { 'text/plain': ['42'] }, + metadata: {}, + execution_count: 1 + }; + + expect(describeExecutionOutputs([output])).to.equal('42'); + }); + + test('reports no output for an empty output list', () => { + expect(describeExecutionOutputs([])).to.equal('(no output)'); + }); + }); + + suite('executeAgentCell', () => { + let mockExecution: { + appendOutput: sinon.SinonStub; + clearOutput: sinon.SinonStub; + end: sinon.SinonStub; + replaceOutput: sinon.SinonStub; + appendOutputItems: sinon.SinonStub; + start: sinon.SinonStub; + }; + let mockController: NotebookController; + let executeAgentBlockStub: sinon.SinonStub; + let mockServiceContainer: ServiceContainer; + let encryptedStorage: IEncryptedStorage; + let neverCancelled: CancellationToken; + + setup(() => { + secretStorage.clear(); + secretStorage.set('openAiApiKey', 'test-key'); + encryptedStorage = createEncryptedStorageFake(secretStorage); + + const neverCancelledSource = new CancellationTokenSource(); + + neverCancelled = neverCancelledSource.token; + disposables.push(neverCancelledSource); + mockServiceContainer = stubServiceContainerInstance(); + disposables.push(new Disposable(() => sinon.restore())); + + mockExecution = { + appendOutput: sinon.stub().resolves(), + clearOutput: sinon.stub().resolves(), + end: sinon.stub(), + replaceOutput: sinon.stub().resolves(), + appendOutputItems: sinon.stub().resolves(), + start: sinon.stub() + }; + + mockController = { + createNotebookCellExecution: sinon.stub().returns(mockExecution) + } as unknown as NotebookController; + + executeAgentBlockStub = sinon.stub().resolves({ finalOutput: 'done' }); + }); + + teardown(() => { + disposables = dispose(disposables); + reset(mockedVSCodeNamespaces.commands); + // Restore applyEdit default; don't reset() the shared workspace mock. + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(true)); + secretStorage.clear(); + }); + + function createAgentCell(text: string = 'Test prompt') { + return createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text + }); + } + + function createAgentCellInMutableNotebook(cells: NotebookCell[] = [], agentBlockId = 'agent-block-1') { + const notebook = createMockNotebook({ cells }); + const agentCell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' }, id: agentBlockId }, + text: 'Test prompt' + }); + + (agentCell as { notebook: typeof notebook }).notebook = notebook; + (agentCell as { index: number }).index = 0; + cells.unshift(agentCell); + + applyNotebookEditsTo(cells, notebook); + + return { agentCell, cells, notebook }; + } + + function getStdoutChunkText(callIndex: number): string { + const item = mockExecution.appendOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; + + return Buffer.from(item.data).toString('utf-8'); + } + + test('creates execution, clears output, sets planning output, and ends successfully', async () => { + const cell = createAgentCell('Analyze data'); + + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect((mockController.createNotebookCellExecution as sinon.SinonStub).calledOnceWith(cell)).to.be.true; + expect(mockExecution.start.calledOnce).to.be.true; + expect(mockExecution.clearOutput.calledOnce).to.be.true; + expect(mockExecution.clearOutput.calledBefore(mockExecution.replaceOutput)).to.be.true; + 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...'); + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.true; + }); + + // SnapshotService and execute_cell analytics read this shim, not the raw NotebookCellExecution — + // without these events a Run All containing an agent block never matches its own code-cell count. + test('reports the run on the execution-state shim so SnapshotService can see it', async () => { + const cell = createAgentCell('Analyze data'); + const seenStates: NotebookCellExecutionState[] = []; + + disposables.push( + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell === cell) { + seenStates.push(e.state); + } + }) + ); + + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(seenStates).to.deep.equal([NotebookCellExecutionState.Executing, NotebookCellExecutionState.Idle]); + }); + + test('reports Idle on the shim even when the run fails', async () => { + mockExecution.clearOutput.rejects(new Error('Something went wrong')); + + const cell = createAgentCell(); + const seenStates: NotebookCellExecutionState[] = []; + + disposables.push( + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell === cell) { + seenStates.push(e.state); + } + }) + ); + + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(seenStates).to.deep.equal([NotebookCellExecutionState.Executing, NotebookCellExecutionState.Idle]); + }); + + // Incremental deltas only — full transcript per event is O(n²) over the EH boundary. + test('streams text_delta events via appendOutputItems with incremental stdout chunks', async () => { + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.onAgentEvent?.({ type: 'text_delta', text: 'first' }); + await context.onAgentEvent?.({ type: 'text_delta', text: ' second' }); + + return { finalOutput: 'first second' }; + }); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(mockExecution.appendOutputItems.callCount).to.equal(2); + + const item = mockExecution.appendOutputItems.firstCall.args[0] as NotebookCellOutputItem; + expect(item.mime).to.equal('application/vnd.code.notebook.stdout'); + expect(getStdoutChunkText(0)).to.equal('[Agent] Text:\nfirst'); + expect(getStdoutChunkText(1)).to.equal(' 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: '' }; + }); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { + executeAgentBlockFn: executeAgentBlockStub + }); + + const chunk2 = getStdoutChunkText(1); + expect(chunk2).to.include('\n\n'); + expect(chunk2).to.include('[Agent] Tool called: search'); + }); + + test('fails execution and writes clearOutput error to stderr', async () => { + mockExecution.clearOutput.rejects(new Error('Something went wrong')); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { + 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[]; + 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, encryptedStorage, neverCancelled, { + 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, encryptedStorage, neverCancelled, { + 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'); + }); + + // Caller clears prior ephemeral output; dirty notebook would duplicate agent context. + test('refuses to run rather than clearing the previous run itself', async () => { + const previousResult = createMockCell({ + text: 'print("previous run")', + metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }, + index: 1 + }); + const { agentCell, cells } = createAgentCellInMutableNotebook([previousResult]); + + await executeAgentCell(agentCell, mockController, encryptedStorage, neverCancelled, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(executeAgentBlockStub.called).to.be.false; + expect(mockExecution.end.firstCall.args[0]).to.be.false; + expect(cells).to.include(previousResult); + + const [outputs] = mockExecution.appendOutput.firstCall.args as [NotebookCellOutput[]]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('previous run'); + }); + + test('inserts ephemeral markdown cells after the agent cell in order', async () => { + const { agentCell, cells } = createAgentCellInMutableNotebook(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.addMarkdownBlock({ content: 'first' }); + await context.addMarkdownBlock({ content: 'second' }); + + return { finalOutput: '' }; + }); + + await executeAgentCell(agentCell, mockController, encryptedStorage, neverCancelled, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(cells.map((cell) => cell.document.getText())).to.deep.equal(['Test prompt', 'first', 'second']); + expect(cells[1].metadata?.is_ephemeral).to.be.true; + expect(cells[1].metadata?.agent_source_block_id).to.equal(agentCell.metadata?.id); + expect(cells[2].metadata?.is_ephemeral).to.be.true; + expect(cells[2].metadata?.agent_source_block_id).to.equal(agentCell.metadata?.id); + }); + + // cellAt clamps — failed insert must not run an existing cell at that index. + test('fails the tool call without executing anything when the insert edit is rejected', async () => { + const { agentCell } = createAgentCellInMutableNotebook(); + let toolResult: string | undefined; + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(false)); + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + toolResult = await context.addAndExecuteCodeBlock({ code: 'print(1)' }); + + return { finalOutput: '' }; + }); + + await executeAgentCell(agentCell, mockController, encryptedStorage, neverCancelled, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(toolResult).to.include('Execution error'); + expect(toolResult).to.include('Failed to insert ephemeral code cell'); + verify(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).never(); + }); + + test('passes project MCP servers and integrations to the agent', async () => { + const integrations = [{ id: 'warehouse', name: 'Warehouse', type: 'postgres' }]; + const mcpServers = [{ name: 'files', command: 'mcp-files', args: [] }]; + const notebookManager = mock(); + + when(mockServiceContainer.tryGet(IDeepnoteNotebookManager)).thenReturn( + instance(notebookManager) + ); + when(notebookManager.getProjectForNotebook('project-1', 'notebook-1')).thenReturn( + createDeepnoteFile({ project: createDeepnoteProject({ integrations, settings: { mcpServers } }) }) + ); + + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test prompt', + notebookMetadata: { deepnoteProjectId: 'project-1', deepnoteNotebookId: 'notebook-1' } + }); + + await executeAgentCell(cell, mockController, encryptedStorage, neverCancelled, { + executeAgentBlockFn: executeAgentBlockStub + }); + + const context = executeAgentBlockStub.firstCall.args[1] as AgentBlockContext; + expect(context.mcpServers).to.deep.equal(mcpServers); + expect(context.integrations).to.deep.equal(integrations); + }); + + suite('cancellation', () => { + let runTokenSource: CancellationTokenSource; + + setup(() => { + runTokenSource = new CancellationTokenSource(); + disposables.push(runTokenSource); + }); + + test('refuses to add a generated cell once stopped', async () => { + const { agentCell, cells } = createAgentCellInMutableNotebook(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + runTokenSource.cancel(); + await context.addMarkdownBlock({ content: 'after the stop' }); + + return { finalOutput: '' }; + }); + + await executeAgentCell(agentCell, mockController, encryptedStorage, runTokenSource.token, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(cells.map((cell) => cell.document.getText())).to.deep.equal(['Test prompt']); + }); + + test('stops at the next stream event rather than running to completion', async () => { + const cell = createAgentCell(); + let eventsAfterStop = 0; + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + runTokenSource.cancel(); + await context.onAgentEvent?.({ type: 'text_delta', text: 'first' }); + eventsAfterStop += 1; + await context.onAgentEvent?.({ type: 'text_delta', text: 'second' }); + + return { finalOutput: '' }; + }); + + await executeAgentCell(cell, mockController, encryptedStorage, runTokenSource.token, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(eventsAfterStop).to.equal(0); + }); + + test('reports a stop as stopped rather than as a failed run', async () => { + const cell = createAgentCell(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + runTokenSource.cancel(); + await context.onAgentEvent?.({ type: 'text_delta', text: 'first' }); + + return { finalOutput: '' }; + }); + + await executeAgentCell(cell, mockController, encryptedStorage, runTokenSource.token, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(mockExecution.end.firstCall.args[0]).to.be.false; + + const [outputs] = mockExecution.appendOutput.firstCall.args as [NotebookCellOutput[]]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('Stopped'); + expect(text).to.not.include('Canceled'); + }); + + test('a run that is never stopped still completes', async () => { + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, encryptedStorage, runTokenSource.token, { + executeAgentBlockFn: executeAgentBlockStub + }); + + expect(mockExecution.end.firstCall.args[0]).to.be.true; + }); + }); + }); + + suite('removeEphemeralCellsForAgentBlocks', () => { + teardown(() => { + sinon.restore(); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(true)); + }); + + function createAgentCell(agentBlockId: string) { + return createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' }, id: agentBlockId }, + text: 'Test prompt' + }); + } + + function createEphemeralCell(agentBlockId: string, text: string) { + return createMockCell({ + metadata: { is_ephemeral: true, agent_source_block_id: agentBlockId }, + text + }); + } + + function createMutableNotebook(cells: NotebookCell[]) { + const notebook = createMockNotebook({ cells }); + + cells.forEach((cell, index) => { + (cell as { notebook: NotebookDocument }).notebook = notebook; + (cell as { index: number }).index = index; + }); + + return { notebook, ...applyNotebookEditsTo(cells, notebook) }; + } + + test('drops the previous run from the batch and deletes it from the notebook', async () => { + const agentCell = createAgentCell('agent-block-1'); + const previousResult = createEphemeralCell('agent-block-1', 'print("previous run")'); + const cells = [agentCell, previousResult]; + const { notebook } = createMutableNotebook(cells); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [...cells]); + + expect(batch).to.deep.equal([agentCell]); + expect(cells).to.deep.equal([agentCell]); + }); + + test('keeps another agent and ordinary cells', async () => { + const agentCell = createAgentCell('agent-block-1'); + const ownResult = createEphemeralCell('agent-block-1', 'own'); + const otherAgentResult = createEphemeralCell('agent-block-2', 'other agent'); + const userCell = createMockCell({ text: 'user code', metadata: {} }); + const cells = [agentCell, ownResult, otherAgentResult, userCell]; + const { notebook } = createMutableNotebook(cells); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [...cells]); + + expect(batch).to.deep.equal([agentCell, otherAgentResult, userCell]); + expect(cells).to.deep.equal([agentCell, otherAgentResult, userCell]); + }); + + // Generated cell may execute alone via notebook.cell.execute — must stay in batch. + test('leaves an ephemeral cell whose agent is not in the batch', async () => { + const agentCell = createAgentCell('agent-block-1'); + const generatedCell = createEphemeralCell('agent-block-1', 'print("just generated")'); + const cells = [agentCell, generatedCell]; + const { notebook, appliedEdits } = createMutableNotebook(cells); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [generatedCell]); + + expect(batch).to.deep.equal([generatedCell]); + expect(cells).to.deep.equal([agentCell, generatedCell]); + expect(appliedEdits()).to.equal(0); + }); + + test('applies no edit when the batch has no agent cell', async () => { + const userCell = createMockCell({ text: 'user code', metadata: {} }); + const cells = [userCell]; + const { notebook, appliedEdits } = createMutableNotebook(cells); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [...cells]); + + expect(batch).to.deep.equal([userCell]); + expect(appliedEdits()).to.equal(0); + }); + + // Rejected delete edit must not block running ordinary cells in the batch. + test('still drops the previous run from the batch when the edit is rejected', async () => { + const agentCell = createAgentCell('agent-block-1'); + const previousResult = createEphemeralCell('agent-block-1', 'print("previous run")'); + const userCell = createMockCell({ text: 'user code', metadata: {} }); + const cells = [agentCell, previousResult, userCell]; + const { notebook } = createMutableNotebook(cells); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(false)); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [...cells]); + + expect(batch).to.deep.equal([agentCell, userCell]); + expect(cells).to.deep.equal([agentCell, previousResult, userCell]); + }); + }); + + suite('executeEphemeralCell', () => { + suite('with active cancellation token', () => { + let tokenSource: CancellationTokenSource; + + setup(() => { + tokenSource = new CancellationTokenSource(); + }); + + teardown(() => { + tokenSource.dispose(); + 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 }); + + (cell as { index: number }).index = currentIndex; + + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenCall(async () => { + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Idle); + }); + + await executeEphemeralCell(cell, tokenSource.token); + + const [commandName, commandArg] = capture( + mockedVSCodeNamespaces.commands.executeCommand as (cmd: string, arg: unknown) => Thenable + ).last(); + + expect(commandName).to.equal('notebook.cell.execute'); + expect(commandArg).to.deep.equal({ + ranges: [{ start: currentIndex, end: currentIndex + 1 }], + document: cell.notebook.uri + }); + }); + + test('reports the failure reason instead of swallowing it', async () => { + const cell = createMockCell({ index: 0 }); + + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenReject( + new Error('kernel is dead') + ); + + const result = await executeEphemeralCell(cell, tokenSource.token); + + expect(result.success).to.be.false; + expect(result.error).to.equal('kernel is dead'); + }); + }); + + // Pre-cancelled token must skip executeCommand, not only the idle wait. + suite('with pre-cancelled token', () => { + let tokenSource: CancellationTokenSource; + + setup(() => { + tokenSource = new CancellationTokenSource(); + tokenSource.cancel(); + }); + + teardown(() => { + tokenSource.dispose(); + reset(mockedVSCodeNamespaces.commands); + }); + + test('throws without dispatching to the kernel when the token is pre-cancelled', async () => { + const cell = createMockCell({ index: 0 }); + + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); + + try { + await executeEphemeralCell(cell, tokenSource.token); + expect.fail('Should have thrown'); + } catch (e) { + expect(e).to.be.instanceOf(CancellationError); + } + + verify(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).never(); + }); + }); + + // Timeout must fire even when executeCommand never resolves. + suite('with fake timers', () => { + let clock: sinon.SinonFakeTimers; + let tokenSource: CancellationTokenSource; + + setup(() => { + clock = sinon.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }); + tokenSource = new CancellationTokenSource(); + }); + + teardown(() => { + clock.restore(); + tokenSource.dispose(); + reset(mockedVSCodeNamespaces.commands); + }); + + test('times out while the dispatch is still pending', async () => { + const cell = createMockCell({ index: 0 }); + + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenCall( + () => new Promise(() => undefined) + ); + + const resultPromise = executeEphemeralCell(cell, tokenSource.token); + await clock.tickAsync(EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS); + + const result = await resultPromise; + + expect(result.success).to.be.false; + expect(result.error).to.equal('Ephemeral cell execution timed out'); + }); + }); + }); +}); diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts new file mode 100644 index 0000000000..d922a78ca4 --- /dev/null +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -0,0 +1,244 @@ +import { + CancellationToken, + Disposable, + EventEmitter, + NotebookCell, + NotebookCellStatusBarItem, + NotebookCellStatusBarItemProvider, + NotebookEdit, + NotebookRange, + WorkspaceEdit, + commands, + l10n, + notebooks, + window, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { getBlockId, getEphemeralCellAgentSourceBlockId, isAgentCell } from './dataConversionUtils'; + +/** Same key as `agentBlockSchema` / `executeAgentBlock`. */ +export const AGENT_MODEL_METADATA_KEY = 'deepnote_agent_model'; + +/** Persisted default — absent key becomes `undefined` and breaks openai() model selection. */ +export const AGENT_MODEL_AUTO = 'auto'; + +const AGENT_MODEL_OPTIONS = [AGENT_MODEL_AUTO, 'gpt-5.6-sol', 'gpt-5.6-terra', 'gpt-5.6-luna']; + +const CLEAR_EPHEMERAL_BLOCKS_COMMAND = 'deepnote.clearEphemeralBlocks'; +const SWITCH_AGENT_MODEL_COMMAND = 'deepnote.switchAgentModel'; + +const AGENT_INDICATOR_PRIORITY = 100; +const MODEL_PICKER_PRIORITY = 90; +const CLEAR_EPHEMERAL_PRIORITY = 80; + +@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(SWITCH_AGENT_MODEL_COMMAND, async (cell?: NotebookCell) => { + if (!cell) { + throw new Error(`${SWITCH_AGENT_MODEL_COMMAND} requires the cell it was invoked from`); + } + + await this.switchModel(cell); + }) + ); + + this.disposables.push( + commands.registerCommand(CLEAR_EPHEMERAL_BLOCKS_COMMAND, async (cell?: NotebookCell) => { + if (!cell) { + throw new Error(`${CLEAR_EPHEMERAL_BLOCKS_COMMAND} requires the cell it was invoked from`); + } + + await this.clearEphemeralBlocks(cell); + }) + ); + + this.disposables.push(this._onDidChangeCellStatusBarItems); + } + + public dispose(): void { + this.disposables.forEach((disposable) => disposable.dispose()); + } + + public provideCellStatusBarItems( + cell: NotebookCell, + token: CancellationToken + ): NotebookCellStatusBarItem[] | undefined { + if (token.isCancellationRequested) { + return undefined; + } + + if (!isAgentCell(cell)) { + return undefined; + } + + const metadata = cell.metadata as Record | undefined; + const model = this.getModel(metadata); + + const items = [this.createAgentIndicatorItem(), this.createModelPickerItem(cell, model)]; + + if (this.getCellsToClear(cell).length > 0) { + items.push(this.createClearEphemeralItem(cell)); + } + + return items; + } + + /** Deletes the ephemeral cells this agent block generated, after a modal confirmation. */ + private async clearEphemeralBlocks(cell: NotebookCell): Promise { + if (!isAgentCell(cell)) { + return; + } + + const cellsToClear = this.getCellsToClear(cell); + + if (cellsToClear.length === 0) { + return; + } + + const confirmation = await window.showWarningMessage( + l10n.t('Clear {0} ephemeral block(s) from this notebook?', cellsToClear.length), + { modal: true }, + l10n.t('Clear') + ); + + if (confirmation !== l10n.t('Clear')) { + return; + } + + // Descending so each deletion's index still addresses the cell it was computed from. + const deletions = [...cellsToClear] + .sort((a, b) => b.index - a.index) + .map((target) => NotebookEdit.deleteCells(new NotebookRange(target.index, target.index + 1))); + + const edit = new WorkspaceEdit(); + edit.set(cell.notebook.uri, deletions); + + if (!(await workspace.applyEdit(edit))) { + void window.showErrorMessage(l10n.t('Failed to clear ephemeral blocks')); + } + } + + private createAgentIndicatorItem(): NotebookCellStatusBarItem { + return { + text: `$(hubot) ${l10n.t('Agent Block')}`, + alignment: 1, + priority: AGENT_INDICATOR_PRIORITY, + tooltip: l10n.t('Deepnote Agent Block\nAI-powered block that autonomously generates code and analysis') + }; + } + + private createClearEphemeralItem(cell: NotebookCell): NotebookCellStatusBarItem { + return { + text: `$(trash) ${l10n.t('Clear ephemeral blocks')}`, + alignment: 1, + priority: CLEAR_EPHEMERAL_PRIORITY, + tooltip: l10n.t('Remove the ephemeral blocks generated by this agent block'), + command: { + title: l10n.t('Clear ephemeral blocks'), + command: CLEAR_EPHEMERAL_BLOCKS_COMMAND, + arguments: [cell] + } + }; + } + + private createModelPickerItem(cell: NotebookCell, model: string): NotebookCellStatusBarItem { + return { + text: `$(symbol-enum) ${l10n.t('Model: {0}', model)}`, + alignment: 1, + priority: MODEL_PICKER_PRIORITY, + tooltip: l10n.t('AI Model: {0}\nClick to change', model), + command: { + title: l10n.t('Switch Model'), + command: SWITCH_AGENT_MODEL_COMMAND, + arguments: [cell] + } + }; + } + + /** Ephemeral cells this agent block generated; empty when it has no block id or has not run. */ + private getCellsToClear(cell: NotebookCell): NotebookCell[] { + const agentBlockId = getBlockId(cell); + + if (!agentBlockId) { + return []; + } + + return cell.notebook + .getCells() + .filter((candidate) => getEphemeralCellAgentSourceBlockId(candidate) === agentBlockId); + } + + private getModel(metadata: Record | undefined): string { + const value = metadata?.[AGENT_MODEL_METADATA_KEY]; + if (typeof value === 'string' && value) { + return value; + } + + return AGENT_MODEL_AUTO; + } + + private async switchModel(cell: NotebookCell): Promise { + if (!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; + } + + await this.updateCellMetadata(cell, { [AGENT_MODEL_METADATA_KEY]: selected.label }); + } + + private async updateCellMetadata(cell: NotebookCell, updates: Record): Promise { + const updatedMetadata = { ...cell.metadata, ...updates }; + + 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..7adb119af1 --- /dev/null +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -0,0 +1,463 @@ +import { assert, expect } from 'chai'; +import * as sinon from 'sinon'; +import { anything, verify, when } from 'ts-mockito'; +import { CancellationToken, NotebookCell, NotebookEdit, WorkspaceEdit } from 'vscode'; + +import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock'; +import { AgentCellStatusBarProvider } from './agentCellStatusBarProvider'; +import { createMockCell, createMockNotebookWithCells } from './deepnoteTestHelpers'; + +suite('AgentCellStatusBarProvider', () => { + let provider: AgentCellStatusBarProvider; + let mockToken: CancellationToken; + + const commandHandlers = new Map Promise>(); + + // Records every command registration. Call AFTER resetVSCodeMocks(), which regenerates the mocks. + function activateCapturingCommands(): void { + commandHandlers.clear(); + when( + mockedVSCodeNamespaces.notebooks.registerNotebookCellStatusBarItemProvider(anything(), anything()) + ).thenReturn({ dispose: () => undefined }); + when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument).thenReturn(() => ({ + dispose: () => undefined + })); + when(mockedVSCodeNamespaces.commands.registerCommand(anything(), anything())).thenCall( + (id: string, callback: (cell?: NotebookCell) => Promise) => { + commandHandlers.set(id, callback); + + return { dispose: () => undefined }; + } + ); + + provider.activate(); + } + + function handlerFor(id: string): (cell?: NotebookCell) => Promise { + const handler = commandHandlers.get(id); + + if (!handler) { + throw new Error(`No handler captured for '${id}'; call activateCapturingCommands() first.`); + } + + return handler; + } + + setup(() => { + mockToken = { + isCancellationRequested: false, + 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(2); + }); + + test('Should return undefined for any cell that is not an agent block', () => { + const nonAgentCells: Record> = { + 'code cell': { __deepnotePocket: { type: 'code' } }, + 'sql cell': { __deepnotePocket: { type: 'sql' } }, + 'markdown cell': { __deepnotePocket: { type: 'markdown' } }, + 'cell without a pocket': {} + }; + + for (const [description, metadata] of Object.entries(nonAgentCells)) { + const items = provider.provideCellStatusBarItems(createMockCell({ metadata }), mockToken); + + expect(items, description).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 and no command', () => { + 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); + expect(items[0].command).to.be.undefined; + }); + }); + + suite('Model Picker', () => { + test('Should display default model picker for agent cell without model metadata', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: auto'); + expect(items[1].text).to.include('$(symbol-enum)'); + expect(items[1].command).to.not.be.undefined; + const cmd = items[1].command as any; + expect(cmd.command).to.equal('deepnote.switchAgentModel'); + expect(items[1].priority).to.equal(90); + }); + + test('Should display configured model from metadata', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_agent_model: 'gpt-5.6-sol', + deepnote_max_iterations: 50 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: gpt-5.6-sol'); + }); + + test('Should display "auto" when model is empty string', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_agent_model: '' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: auto'); + }); + }); + + // Driven through the registered command rather than switchModel directly, so each case also + // covers the wiring the status bar item actually goes through. + suite('Model Switching', () => { + let capturedEdit: { index: number; metadata: Record } | undefined; + let invokeCommand: (cell?: NotebookCell) => Promise; + + setup(() => { + resetVSCodeMocks(); + capturedEdit = undefined; + + // The vscode mock's NotebookEdit.updateCellMetadata discards its metadata argument, so + // capturing the WorkspaceEdit cannot show what was written. Stub the static instead. + sinon.stub(NotebookEdit, 'updateCellMetadata').callsFake((index: number, metadata) => { + capturedEdit = { index, metadata: metadata as Record }; + + return {} as NotebookEdit; + }); + + activateCapturingCommands(); + invokeCommand = handlerFor('deepnote.switchAgentModel'); + }); + + teardown(() => { + sinon.restore(); + resetVSCodeMocks(); + }); + + function agentCell(): NotebookCell { + return createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent', id: 'pocket-1' }, + id: 'block-1', + deepnote_agent_model: 'gpt-5.6-sol' + }, + index: 2 + }); + } + + function pick(label: string | undefined) { + when(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).thenReturn( + Promise.resolve(label === undefined ? undefined : ({ label } as any)) + ); + } + + test('Should write the picked model without dropping the cell’s other metadata', async () => { + // Catches: an inverted spread in updateCellMetadata, which makes the switch a silent + // no-op while still calling applyEdit — so a call-count assertion would not notice. + pick('gpt-5.6-terra'); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(true)); + + await invokeCommand(agentCell()); + + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).once(); + expect(capturedEdit!.index).to.equal(2); + expect(capturedEdit!.metadata).to.deep.equal({ + __deepnotePocket: { type: 'agent', id: 'pocket-1' }, + id: 'block-1', + deepnote_agent_model: 'gpt-5.6-terra' + }); + }); + + test('Should not edit the notebook when the current model is re-picked', async () => { + // Catches: losing the `selected.label === currentModel` guard, which dirties the + // document on a no-op selection. + pick('gpt-5.6-sol'); + + await invokeCommand(agentCell()); + + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('Should not edit the notebook when the picker is dismissed', async () => { + // Catches: losing the `!selected` guard, which throws on `selected.label` when the + // user presses Escape. + pick(undefined); + + await invokeCommand(agentCell()); + + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('Should report an error when the workspace edit is rejected', async () => { + // Catches: dropping the `if (!success)` branch, which loses the model change silently. + pick('gpt-5.6-luna'); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(false)); + + let statusBarRefreshed = false; + provider.onDidChangeCellStatusBarItems(() => { + statusBarRefreshed = true; + }); + + await invokeCommand(agentCell()); + + verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); + expect(statusBarRefreshed, 'a rejected edit must not refresh the status bar').to.be.false; + }); + + test('Should ignore a non-agent cell', async () => { + // Catches: dropping the isAgentCell guard, which would offer the model picker on any cell. + pick('gpt-5.6-luna'); + + await invokeCommand(createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } })); + + verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('Should reject when invoked without a cell', async () => { + // Catches: falling back to the selected cell, which rewrites a block the user never clicked. + pick('gpt-5.6-terra'); + + await assert.isRejected(invokeCommand(undefined), /requires the cell it was invoked from/); + + verify(mockedVSCodeNamespaces.window.showQuickPick(anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + }); + + // Driven through the registered command rather than clearEphemeralBlocks directly, so each case + // also covers the wiring the status bar item actually goes through. + suite('Clearing ephemeral blocks', () => { + let invokeCommand: (cell?: NotebookCell) => Promise; + + setup(() => { + resetVSCodeMocks(); + activateCapturingCommands(); + invokeCommand = handlerFor('deepnote.clearEphemeralBlocks'); + }); + + teardown(() => { + sinon.restore(); + resetVSCodeMocks(); + }); + + // Mocked WorkspaceEdit.set drops the edits, so capture them on the prototype and replay the + // deletions against `cells` — an ascending delete order corrupts the survivors, not the count. + function applyDeletionsTo(cells: NotebookCell[]): void { + let recordedEdits: { range: { start: number; end: number } }[] = []; + + sinon.stub(WorkspaceEdit.prototype, 'set').callsFake((_uri, edits) => { + recordedEdits = edits as unknown as { range: { start: number; end: number } }[]; + }); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { + for (const { range } of recordedEdits) { + cells.splice(range.start, range.end - range.start); + } + + return Promise.resolve(true); + }); + } + + function confirmWith(label: string | undefined): void { + when(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).thenReturn( + Promise.resolve(label as any) + ); + } + + function agentBlock(text: string, blockId: string) { + return { text, metadata: { __deepnotePocket: { type: 'agent' }, id: blockId } }; + } + + function ephemeralCell(text: string, agentSourceBlockId?: string) { + return { + text, + metadata: { + is_ephemeral: true, + ...(agentSourceBlockId ? { agent_source_block_id: agentSourceBlockId } : {}) + } + }; + } + + test('Should delete every cell this agent block generated and nothing else', async () => { + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph A1', 'agent-block-1'), + ephemeralCell('eph A2', 'agent-block-1'), + { text: 'user code', metadata: {} }, + agentBlock('agent B', 'agent-block-2'), + ephemeralCell('eph B1', 'agent-block-2'), + ephemeralCell('orphan') + ]); + applyDeletionsTo(cells); + confirmWith('Clear'); + + await invokeCommand(cells[0]); + + expect(cells.map((cell) => cell.document.getText())).to.deep.equal([ + 'agent A', + 'user code', + 'agent B', + 'eph B1', + 'orphan' + ]); + }); + + test('Should report how many blocks the clear removes', async () => { + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph A1', 'agent-block-1'), + ephemeralCell('eph A2', 'agent-block-1') + ]); + applyDeletionsTo(cells); + confirmWith('Clear'); + + await invokeCommand(cells[0]); + + verify( + mockedVSCodeNamespaces.window.showWarningMessage( + 'Clear 2 ephemeral block(s) from this notebook?', + anything(), + anything() + ) + ).once(); + }); + + test('Should not edit the notebook when the confirmation is dismissed', async () => { + // Catches: applying the edit before the modal is answered, which deletes on a cancel. + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph A1', 'agent-block-1') + ]); + applyDeletionsTo(cells); + confirmWith(undefined); + + await invokeCommand(cells[0]); + + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + expect(cells).to.have.lengthOf(2); + }); + + test('Should not prompt for an agent block that generated nothing', async () => { + // Catches: prompting to clear 0 blocks when the agent has not run. + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph B1', 'agent-block-2') + ]); + applyDeletionsTo(cells); + confirmWith('Clear'); + + await invokeCommand(cells[0]); + + verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('Should ignore a cell that is not an agent block', async () => { + // Catches: dropping the isAgentCell guard, which would clear from any cell whose id + // happens to own ephemeral children. + const { cells } = createMockNotebookWithCells([ + { text: 'user code', metadata: { id: 'agent-block-1' } }, + ephemeralCell('eph A1', 'agent-block-1') + ]); + applyDeletionsTo(cells); + confirmWith('Clear'); + + await invokeCommand(cells[0]); + + verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + + test('Should report an error when the workspace edit is rejected', async () => { + // Catches: dropping the `if (!applyEdit)` branch, which loses the clear silently. + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph A1', 'agent-block-1') + ]); + confirmWith('Clear'); + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenReturn(Promise.resolve(false)); + + await invokeCommand(cells[0]); + + verify(mockedVSCodeNamespaces.window.showErrorMessage(anything())).once(); + }); + + suite('Clear button', () => { + test('Should offer the button on an agent block that generated ephemeral cells', () => { + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph A1', 'agent-block-1') + ]); + const items = provider.provideCellStatusBarItems(cells[0], mockToken)!; + + expect(items).to.have.lengthOf(3); + expect(items[2].text).to.include('$(trash)'); + expect(items[2].text).to.include('Clear ephemeral blocks'); + expect(items[2].alignment).to.equal(1); + expect(items[2].priority).to.equal(80); + + const command = items[2].command as { command: string; arguments: unknown[] }; + expect(command.command).to.equal('deepnote.clearEphemeralBlocks'); + expect(command.arguments).to.deep.equal([cells[0]]); + }); + + test('Should hide the button on an agent block that owns no ephemeral cells', () => { + // Catches: an always-visible button, which prompts to clear 0 blocks. + const { cells } = createMockNotebookWithCells([ + agentBlock('agent A', 'agent-block-1'), + ephemeralCell('eph B1', 'agent-block-2') + ]); + const items = provider.provideCellStatusBarItems(cells[0], mockToken)!; + + expect(items).to.have.lengthOf(2); + }); + }); + + test('Should reject when invoked without a cell', async () => { + // Catches: falling back to the selected cell, which clears a run the user never clicked. + confirmWith('Clear'); + + await assert.isRejected(invokeCommand(undefined), /requires the cell it was invoked from/); + + verify(mockedVSCodeNamespaces.window.showWarningMessage(anything(), anything(), anything())).never(); + verify(mockedVSCodeNamespaces.workspace.applyEdit(anything())).never(); + }); + }); +}); diff --git a/src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts b/src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts new file mode 100644 index 0000000000..8ed762b1b1 --- /dev/null +++ b/src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts @@ -0,0 +1,34 @@ +import { inject, injectable } from 'inversify'; +import { commands, l10n, window } from 'vscode'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { IEncryptedStorage } from '../../platform/common/application/types'; +import { IExtensionContext } from '../../platform/common/types'; +import { clearOpenAiApiKey, promptForOpenAiApiKey } from './deepnoteSecretStore'; + +@injectable() +export class AgentOpenAiApiKeyCommandHandler implements IExtensionSyncActivationService { + constructor( + @inject(IEncryptedStorage) private readonly encryptedStorage: IEncryptedStorage, + @inject(IExtensionContext) private readonly extensionContext: IExtensionContext + ) {} + + public activate(): void { + this.extensionContext.subscriptions.push( + commands.registerCommand('deepnote.setOpenAiApiKey', () => this.setApiKey()), + commands.registerCommand('deepnote.clearOpenAiApiKey', () => this.clearApiKey()) + ); + } + + private async setApiKey(): Promise { + const key = await promptForOpenAiApiKey(this.encryptedStorage); + if (key) { + void window.showInformationMessage(l10n.t('OpenAI API key has been saved.')); + } + } + + private async clearApiKey(): Promise { + await clearOpenAiApiKey(this.encryptedStorage); + void window.showInformationMessage(l10n.t('OpenAI API key has been cleared.')); + } +} diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.ts b/src/notebooks/deepnote/converters/agentBlockConverter.ts new file mode 100644 index 0000000000..416e01a91d --- /dev/null +++ b/src/notebooks/deepnote/converters/agentBlockConverter.ts @@ -0,0 +1,25 @@ +import type { DeepnoteBlock } from '@deepnote/blocks'; +import { NotebookCellData, NotebookCellKind } from 'vscode'; + +import type { BlockConverter } from './blockConverter'; + +/** Agent blocks as plaintext code cells; pocket metadata from DeepnoteDataConverter. */ +export class AgentBlockConverter implements BlockConverter { + applyChangesToBlock(block: DeepnoteBlock, cell: NotebookCellData): void { + block.content = cell.value; + } + + 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..65cc4cc488 --- /dev/null +++ b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts @@ -0,0 +1,144 @@ +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('accepts agent type case-insensitively', () => { + assert.strictEqual(converter.canConvert('agent'), true); + 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('normalizes missing or empty content to empty cell value', () => { + const emptyBlock: DeepnoteBlock = { + blockGroup: 'test-group', + content: '', + id: 'agent-block-456', + sortingKey: 'a1', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + const undefinedContentBlock: DeepnoteBlock = { + blockGroup: 'test-group', + id: 'agent-block-789', + sortingKey: 'a2', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + assert.strictEqual(converter.convertToCell(emptyBlock).value, ''); + assert.strictEqual(converter.convertToCell(undefinedContentBlock).value, ''); + }); + + 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'); + }); + }); + + suite('applyChangesToBlock', () => { + test('updates block content without modifying other block properties', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Old prompt', + id: 'agent-block-789', + metadata: { + deepnote_agent_model: 'gpt-5.6-sol', + 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-5.6-sol', + custom: 'value' + }); + }); + + 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, ''); + }); + }); +}); diff --git a/src/notebooks/deepnote/dataConversionUtils.ts b/src/notebooks/deepnote/dataConversionUtils.ts index 1b30484770..ea9948eaa8 100644 --- a/src/notebooks/deepnote/dataConversionUtils.ts +++ b/src/notebooks/deepnote/dataConversionUtils.ts @@ -2,6 +2,10 @@ * Utility functions for Deepnote block ID and sorting key generation */ +import { NotebookCell, NotebookCellData } from 'vscode'; + +import type { Pocket } from '../../platform/deepnote/pocket'; + export function parseJsonWithFallback(value: string, fallback?: unknown): unknown | null { try { return JSON.parse(value); @@ -22,6 +26,35 @@ export function generateBlockId(): string { return id; } +/** Agent block cell. Lives here so importers avoid `@deepnote/runtime-core`. */ +export function isAgentCell(cell: NotebookCell): boolean { + const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; + + return pocket?.type === 'agent'; +} + +/** Agent-generated scratch cell (`metadata.is_ephemeral`). */ +export function isEphemeralCell(cell: NotebookCell | NotebookCellData): boolean { + return cell.metadata?.is_ephemeral === true; +} + +/** + * Serialized block id, or undefined. Prefer `__deepnoteBlockId` — VS Code may rewrite `id`. + * Missing id makes callers mint a new one and reassign the block on save. + */ +export function getBlockId(cell: NotebookCell | NotebookCellData): string | undefined { + return ( + (cell.metadata?.__deepnoteBlockId as string | undefined) || + (cell.metadata?.id as string | undefined) || + (cell.metadata?.deepnoteBlockId as string | undefined) + ); +} + +/** Owning agent block id when `isEphemeralCell`; otherwise undefined. */ +export function getEphemeralCellAgentSourceBlockId(cell: NotebookCell): string | undefined { + return isEphemeralCell(cell) ? (cell.metadata?.agent_source_block_id as string | undefined) : undefined; +} + /** * Generate sorting key based on index (format: a0, a1, ..., a99, b0, b1, ...) */ diff --git a/src/notebooks/deepnote/dataConversionUtils.unit.test.ts b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts new file mode 100644 index 0000000000..ee90e0c0b9 --- /dev/null +++ b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts @@ -0,0 +1,80 @@ +import { expect } from 'chai'; + +import { getBlockId, getEphemeralCellAgentSourceBlockId, isAgentCell } from './dataConversionUtils'; +import { createMockCell } from './deepnoteTestHelpers'; + +suite('DataConversionUtils', () => { + suite('isAgentCell', () => { + test('returns true for cell with agent pocket type', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + + expect(isAgentCell(cell)).to.be.true; + }); + + test('returns false for non-agent pocket types', () => { + const codeCell = createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } }); + const markdownCell = createMockCell({ metadata: { __deepnotePocket: { type: 'markdown' } } }); + + expect(isAgentCell(codeCell)).to.be.false; + expect(isAgentCell(markdownCell)).to.be.false; + }); + + test('returns false when pocket type is not agent', () => { + const noPocketCell = createMockCell({ metadata: {} }); + const noMetadataCell = createMockCell({ metadata: undefined }); + + expect(isAgentCell(noPocketCell)).to.be.false; + expect(isAgentCell(noMetadataCell)).to.be.false; + }); + }); + + suite('getBlockId', () => { + test('prefers the backup id VS Code cannot rewrite', () => { + const cell = createMockCell({ metadata: { __deepnoteBlockId: 'backup-id', id: 'rewritten-id' } }); + + expect(getBlockId(cell)).to.equal('backup-id'); + }); + + test('falls back to id when the backup is absent', () => { + const cell = createMockCell({ metadata: { id: 'block-id' } }); + + expect(getBlockId(cell)).to.equal('block-id'); + }); + + // Fallback-cell metadata; minting a new id would reassign the block on save. + test('falls back to the legacy deepnoteBlockId when both are absent', () => { + const cell = createMockCell({ metadata: { deepnoteBlockId: 'legacy-id' } }); + + expect(getBlockId(cell)).to.equal('legacy-id'); + }); + + test('ranks the legacy name below both current ones', () => { + const cell = createMockCell({ metadata: { id: 'block-id', deepnoteBlockId: 'legacy-id' } }); + + expect(getBlockId(cell)).to.equal('block-id'); + }); + + test('returns undefined for a cell that was never serialized', () => { + const cell = createMockCell({ metadata: {} }); + + expect(getBlockId(cell)).to.be.undefined; + }); + }); + + suite('getEphemeralCellAgentSourceBlockId', () => { + test('returns the agent block that generated the cell', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' } }); + + expect(getEphemeralCellAgentSourceBlockId(cell)).to.equal('agent-block-1'); + }); + + // agent_source_block_id alone does not mark a cell for agent cleanup. + test('returns undefined when the cell is not ephemeral or ordinary', () => { + const withSourceOnly = createMockCell({ metadata: { agent_source_block_id: 'agent-block-1' } }); + const ordinaryCell = createMockCell({ metadata: {} }); + + expect(getEphemeralCellAgentSourceBlockId(withSourceOnly)).to.be.undefined; + expect(getEphemeralCellAgentSourceBlockId(ordinaryCell)).to.be.undefined; + }); + }); +}); diff --git a/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.ts b/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.ts index 3ab4a7a2c8..0bb6386189 100644 --- a/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.ts +++ b/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.ts @@ -2,7 +2,7 @@ import { inject, injectable } from 'inversify'; import { Disposable, NotebookCellKind, workspace } from 'vscode'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; -import { ITelemetryService } from '../../platform/analytics/types'; +import { ITelemetryService, TelemetryEventProperties } from '../../platform/analytics/types'; import { IDisposableRegistry } from '../../platform/common/types'; import { isDeepnoteNotebook } from '../../platform/common/utils'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; @@ -11,10 +11,12 @@ import { toTelemetryIntegrationType } from '../../platform/notebooks/deepnote/integrationTypes'; import { IDeepnoteNotebookManager } from '../types'; +import { isEphemeralCell } from './dataConversionUtils'; /** * Tracks cell executions, plus the plain code/markdown insertions from VS Code's built-in - * "+ Code" / "+ Markdown" controls that never reach an extension command. + * "+ Code" / "+ Markdown" controls that never reach an extension command, plus the scratch + * cells the agent writes and runs on the user's behalf. */ @injectable() export class DeepnoteCellExecutionAnalytics implements IExtensionSyncActivationService { @@ -39,6 +41,18 @@ export class DeepnoteCellExecutionAnalytics implements IExtensionSyncActivationS } for (const cell of change.addedCells) { + const blockType = cell.kind === NotebookCellKind.Code ? 'code' : 'markdown'; + + // Agent scratch cells stamp a pocket like any typed block, but no command + // inserts them, so this is the only place they can be counted. + if (isEphemeralCell(cell)) { + this.analytics.trackEvent({ + eventName: 'add_block', + properties: { blockType, isEphemeral: true } + }); + continue; + } + // Typed Deepnote blocks stamp a pocket on insert and are already counted by // DeepnoteNotebookCommandListener. if (cell.metadata?.__deepnotePocket?.type) { @@ -47,7 +61,7 @@ export class DeepnoteCellExecutionAnalytics implements IExtensionSyncActivationS this.analytics.trackEvent({ eventName: 'add_block', - properties: { blockType: cell.kind === NotebookCellKind.Code ? 'code' : 'markdown' } + properties: { blockType, isEphemeral: false } }); } } @@ -67,7 +81,12 @@ export class DeepnoteCellExecutionAnalytics implements IExtensionSyncActivationS const languageId = e.cell.document.languageId; const cellType = languageId === 'sql' ? 'sql' : languageId === 'markdown' ? 'markdown' : 'code'; - const properties: { cellType: 'sql' | 'markdown' | 'code'; integrationType?: string } = { cellType }; + // The agent runs its generated cells through `notebook.cell.execute`, which re-enters the + // kernel path; unmarked they are indistinguishable here from a user pressing Run. + const properties: TelemetryEventProperties['execute_cell'] = { + cellType, + isEphemeral: isEphemeralCell(e.cell) + }; if (cellType === 'sql') { // The status-bar switch updates only this key, so the __deepnotePocket copy can go stale. diff --git a/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.unit.test.ts b/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.unit.test.ts index dc72115043..bdb08f633a 100644 --- a/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteCellExecutionAnalytics.unit.test.ts @@ -31,6 +31,17 @@ suite('DeepnoteCellExecutionAnalytics', () => { } as unknown as NotebookCell; } + /** Agent scratch cells go through the converter, so they carry a pocket like any typed block. */ + function ephemeralCell(kind: NotebookCellKind): NotebookCell { + return { + kind, + metadata: { + __deepnotePocket: { type: kind === NotebookCellKind.Code ? 'code' : 'markdown' }, + is_ephemeral: true + } + } as unknown as NotebookCell; + } + function executingCell(languageId: string, sqlIntegrationId?: string, notebookType = 'deepnote'): NotebookCell { return { document: { languageId }, @@ -106,7 +117,9 @@ suite('DeepnoteCellExecutionAnalytics', () => { row.expected.forEach((blockType) => verify( - telemetry.trackEvent(deepEqual({ eventName: 'add_block', properties: { blockType } })) + telemetry.trackEvent( + deepEqual({ eventName: 'add_block', properties: { blockType, isEphemeral: false } }) + ) ).once() ); verify(telemetry.trackEvent(anything())).times(row.expected.length); @@ -118,6 +131,22 @@ suite('DeepnoteCellExecutionAnalytics', () => { verify(telemetry.trackEvent(anything())).never(); }); + + test('counts agent scratch cells, which no command reports, as ephemeral', () => { + fireContentChange([ephemeralCell(NotebookCellKind.Code), ephemeralCell(NotebookCellKind.Markup)], []); + + verify( + telemetry.trackEvent( + deepEqual({ eventName: 'add_block', properties: { blockType: 'code', isEphemeral: true } }) + ) + ).once(); + verify( + telemetry.trackEvent( + deepEqual({ eventName: 'add_block', properties: { blockType: 'markdown', isEphemeral: true } }) + ) + ).once(); + verify(telemetry.trackEvent(anything())).twice(); + }); }); suite('execute_cell', () => { @@ -171,12 +200,30 @@ suite('DeepnoteCellExecutionAnalytics', () => { ); verify( - telemetry.trackEvent(deepEqual({ eventName: 'execute_cell', properties: { ...row.expected } })) + telemetry.trackEvent( + deepEqual({ + eventName: 'execute_cell', + properties: { ...row.expected, isEphemeral: false } + }) + ) ).once(); verify(telemetry.trackEvent(anything())).once(); }); }); + test('an agent-generated cell reports isEphemeral true', () => { + const cell = { ...executingCell('python'), metadata: { is_ephemeral: true } } as unknown as NotebookCell; + + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Executing); + + verify( + telemetry.trackEvent( + deepEqual({ eventName: 'execute_cell', properties: { cellType: 'code', isEphemeral: true } }) + ) + ).once(); + verify(telemetry.trackEvent(anything())).once(); + }); + test('ignores Pending and Idle transitions, and non-Deepnote notebooks', () => { const cell = executingCell('python'); diff --git a/src/notebooks/deepnote/deepnoteDataConverter.ts b/src/notebooks/deepnote/deepnoteDataConverter.ts index 9f71700a71..42e1eb3eaa 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.ts @@ -1,8 +1,9 @@ import { isExecutableBlock, type DeepnoteBlock } from '@deepnote/blocks'; import { NotebookCellData, NotebookCellKind, NotebookCellOutput, NotebookCellOutputItem } from 'vscode'; -import { generateBlockId, generateSortingKey } from './dataConversionUtils'; +import { generateBlockId, generateSortingKey, getBlockId } from './dataConversionUtils'; import type { DeepnoteOutput } from '../../platform/deepnote/deepnoteTypes'; +import { AgentBlockConverter } from './converters/agentBlockConverter'; import { ConverterRegistry } from './converters/converterRegistry'; import { BlockConverter } from './converters/blockConverter'; import { CodeBlockConverter } from './converters/codeBlockConverter'; @@ -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()); @@ -437,7 +439,7 @@ export class DeepnoteDataConverter { private createFallbackBlock(cell: NotebookCellData, index: number): DeepnoteBlock { const meta = cell.metadata as Record | undefined; - const preservedId = (meta?.__deepnoteBlockId ?? meta?.id ?? meta?.deepnoteBlockId) as string | undefined; + const preservedId = getBlockId(cell); const preservedSortingKey = (meta?.sortingKey ?? meta?.deepnoteSortingKey) as string | undefined; const preservedBlockGroup = meta?.blockGroup as string | undefined; @@ -494,19 +496,22 @@ export class DeepnoteDataConverter { } } - // Check if this is a stream output - const stdoutItem = output.items.find((item) => item.mime === 'application/vnd.code.notebook.stdout'); - const stderrItem = output.items.find((item) => item.mime === 'application/vnd.code.notebook.stderr'); - - if (stdoutItem || stderrItem) { - const item = stdoutItem || stderrItem; - const text = new TextDecoder().decode(item!.data); - - return { - name: stderrItem ? 'stderr' : 'stdout', + // Streamed deltas are appended as new items on one output, so every item must be joined. + const streamItems = output.items.filter( + (item) => + item.mime === 'application/vnd.code.notebook.stdout' || + item.mime === 'application/vnd.code.notebook.stderr' + ); + + if (streamItems.length > 0) { + const decoder = new TextDecoder(); + const streamOutput: DeepnoteOutput = { + name: streamItems[0].mime === 'application/vnd.code.notebook.stderr' ? 'stderr' : 'stdout', output_type: 'stream', - text - } as DeepnoteOutput; + text: streamItems.map((item) => decoder.decode(item.data)).join('') + }; + + return streamOutput; } // Rich output (execute_result or display_data) diff --git a/src/notebooks/deepnote/deepnoteDataConverter.unit.test.ts b/src/notebooks/deepnote/deepnoteDataConverter.unit.test.ts index c1ac9b7b94..b7198127b0 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.unit.test.ts @@ -282,6 +282,102 @@ suite('DeepnoteDataConverter', () => { assert.strictEqual(new TextDecoder().decode(outputs[0].items[0].data), 'Hello world\n'); }); + test('joins every streamed stdout item into one stream output', () => { + const cells: NotebookCellData[] = [ + { + kind: NotebookCellKind.Code, + value: 'summarize the data', + languageId: 'plaintext', + metadata: { + __deepnotePocket: { + type: 'agent', + sortingKey: 'a0' + }, + id: 'agent-block-1' + }, + outputs: [ + new NotebookCellOutput([ + NotebookCellOutputItem.stdout('[Agent] Planning next steps...'), + NotebookCellOutputItem.stdout('\n\n[Agent] Tool called: add_code'), + NotebookCellOutputItem.stdout('\n\n[Agent] Text:\nDone.') + ]) + ] + } + ]; + + const blocks = converter.convertCellsToBlocks(cells); + + assert.deepStrictEqual((blocks[0] as ExecutableBlock).outputs, [ + { + name: 'stdout', + output_type: 'stream', + text: '[Agent] Planning next steps...\n\n[Agent] Tool called: add_code\n\n[Agent] Text:\nDone.' + } + ]); + }); + + test('converts a single stderr item to a stderr stream output', () => { + const cells: NotebookCellData[] = [ + { + kind: NotebookCellKind.Code, + value: 'raise ValueError()', + languageId: 'python', + metadata: { + __deepnotePocket: { + type: 'code', + sortingKey: 'a0' + }, + id: 'block-1' + }, + outputs: [new NotebookCellOutput([NotebookCellOutputItem.stderr('Agent execution failed: boom')])] + } + ]; + + const blocks = converter.convertCellsToBlocks(cells); + + assert.deepStrictEqual((blocks[0] as ExecutableBlock).outputs, [ + { + name: 'stderr', + output_type: 'stream', + text: 'Agent execution failed: boom' + } + ]); + }); + + test('keeps every item of a mixed stdout/stderr output and labels it by the first', () => { + const cells: NotebookCellData[] = [ + { + kind: NotebookCellKind.Code, + value: 'print("a")', + languageId: 'python', + metadata: { + __deepnotePocket: { + type: 'code', + sortingKey: 'a0' + }, + id: 'block-1' + }, + outputs: [ + new NotebookCellOutput([ + NotebookCellOutputItem.stdout('out-1\n'), + NotebookCellOutputItem.stderr('err-1\n'), + NotebookCellOutputItem.stdout('out-2\n') + ]) + ] + } + ]; + + const blocks = converter.convertCellsToBlocks(cells); + + assert.deepStrictEqual((blocks[0] as ExecutableBlock).outputs, [ + { + name: 'stdout', + output_type: 'stream', + text: 'out-1\nerr-1\nout-2\n' + } + ]); + }); + test('converts error output', () => { const deepnoteOutputs: DeepnoteOutput[] = [ { diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index c5abf5a442..b6c3f47339 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -12,12 +12,14 @@ import { } from 'vscode'; import { inject, injectable, optional } from 'inversify'; import type { DeepnoteBlock } from '@deepnote/blocks'; +import fastDeepEqual from 'fast-deep-equal'; import { IControllerRegistration } from '../controllers/types'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import { IDisposableRegistry } from '../../platform/common/types'; import { logger } from '../../platform/logging'; import { IDeepnoteNotebookManager } from '../types'; +import { getBlockId, isEphemeralCell } from './dataConversionUtils'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; import { DeepnoteNotebookSerializer } from './deepnoteSerializer'; @@ -158,21 +160,34 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic } /** - * Checks whether the source code content has actually changed between the - * live notebook and the new cells from disk. If only outputs differ (disk - * has fewer/no outputs), it's an auto-save of stripped content — skip reload. + * Checks whether anything the file carries has actually changed between the live notebook and + * the new cells from disk. Outputs and execution state are ignored: in snapshot mode the main + * file has them stripped, so our own auto-save would otherwise look like an external edit. */ private contentActuallyChanged(notebook: NotebookDocument, newCells: NotebookCellData[]): boolean { - const liveCells = notebook.getCells(); + // Ephemeral cells aren't persisted; counting them looks like an external delete. + const liveCells = notebook.getCells().filter((cell) => !isEphemeralCell(cell)); if (liveCells.length !== newCells.length) { return true; } - return liveCells.some( - (live, i) => - live.kind !== newCells[i].kind || - live.document.languageId !== newCells[i].languageId || - live.document.getText() !== newCells[i].value - ); + + return liveCells.some((live, i) => { + const fromDisk = newCells[i]; + + if ( + live.kind !== fromDisk.kind || + live.document.languageId !== fromDisk.languageId || + live.document.getText() !== fromDisk.value || + getBlockId(live) !== getBlockId(fromDisk) + ) { + return true; + } + + const liveBlock = this.persistedBlock(live, i); + const diskBlock = this.persistedBlock(fromDisk, i); + + return liveBlock.type !== diskBlock.type || !fastDeepEqual(liveBlock.metadata, diskBlock.metadata); + }); } /** @@ -279,14 +294,14 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic const liveCells = notebook.getCells(); const liveOutputsByBlockId = new Map(); for (const liveCell of liveCells) { - const blockId = this.getBlockIdFromMetadata(liveCell.metadata); + const blockId = getBlockId(liveCell); if (blockId && liveCell.outputs.length > 0) { liveOutputsByBlockId.set(blockId, liveCell.outputs); } } for (const cell of newCells) { - const blockId = this.getBlockIdFromMetadata(cell.metadata); + const blockId = getBlockId(cell); if (blockId && (!cell.outputs || cell.outputs.length === 0)) { const liveOutputs = liveOutputsByBlockId.get(blockId); if (liveOutputs) { @@ -300,7 +315,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic edits.push(NotebookEdit.replaceCells(new NotebookRange(0, notebook.cellCount), newCells)); for (let i = 0; i < newCells.length; i++) { - const blockId = this.getBlockIdFromMetadata(newCells[i].metadata); + const blockId = getBlockId(newCells[i]); if (blockId) { edits.push( NotebookEdit.updateCellMetadata(i, { @@ -373,15 +388,20 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic blockIdFromFallback: boolean; }> = []; + // Ephemeral cells live in the document but are stripped from the file, so a live index cannot + // address originalBlocks. Track the file's own cursor, advanced only by persisted cells. + let fileBlockIndex = 0; + for (let i = 0; i < liveCells.length; i++) { try { const cell = liveCells[i]; - let blockId = this.getBlockIdFromMetadata(cell.metadata); + const originalBlock = isEphemeralCell(cell) ? undefined : originalBlocks?.[fileBlockIndex++]; + let blockId = getBlockId(cell); let blockIdFromFallback = false; // Fallback to original project blocks when metadata was lost - if (!blockId && originalBlocks) { - blockId = originalBlocks[i]?.id; + if (!blockId && originalBlock) { + blockId = originalBlock.id; blockIdFromFallback = true; } @@ -389,7 +409,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic continue; } - const fallbackType = originalBlocks?.[i]?.type; + const fallbackType = originalBlock?.type; const blockType = ((cell.metadata?.type as string) ?? fallbackType ?? 'code') as DeepnoteBlock['type']; const newOutputs = this.converter.transformOutputsForVsCode( snapshotOutputs.get(blockId)!, @@ -502,10 +522,6 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic logger.info(`[FileChangeWatcher] Updated notebook outputs from external snapshot: ${notebook.uri.path}`); } - private getBlockIdFromMetadata(metadata: Record | undefined): string | undefined { - return (metadata?.__deepnoteBlockId ?? metadata?.id) as string | undefined; - } - private handleFileChange(uri: Uri): void { // Deterministic self-write check — no timers involved if (this.consumeSelfWrite(uri)) { @@ -646,6 +662,26 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic return true; } + /** + * The block a cell would be written as, produced by the very conversion the serializer runs on + * save. Comparing that instead of raw cell metadata keeps the check honest in both directions: + * whatever the write path derives, normalizes or strips ends up on the block rather than in + * `block.metadata`, while everything an external editor can put in the file survives. + * Outputs are left out — they are deliberately not compared and converting them is the + * expensive part. + */ + private persistedBlock(cell: NotebookCell | NotebookCellData, index: number): DeepnoteBlock { + const isLiveCell = 'document' in cell; + const projection = new NotebookCellData( + cell.kind, + isLiveCell ? cell.document.getText() : cell.value, + isLiveCell ? cell.document.languageId : cell.languageId + ); + projection.metadata = { ...cell.metadata }; + + return this.converter.convertCellToBlock(projection, index); + } + private selfWriteKey(uri: Uri): string { return uri.with({ query: '', fragment: '' }).toString(); } diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index cd9eb4d818..18f92991ff 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -1,8 +1,16 @@ -import type { DeepnoteFile } from '@deepnote/blocks'; +import type { DeepnoteBlock, DeepnoteFile } from '@deepnote/blocks'; import { assert } from 'chai'; import * as sinon from 'sinon'; import { anything, instance, mock, when } from 'ts-mockito'; -import { Disposable, EventEmitter, FileSystemWatcher, NotebookCellKind, NotebookDocument, Uri } from 'vscode'; +import { + Disposable, + EventEmitter, + FileSystemWatcher, + NotebookCellKind, + NotebookDocument, + NotebookEdit, + Uri +} from 'vscode'; import type { IControllerRegistration } from '../controllers/types'; import type { IDisposableRegistry } from '../../platform/common/types'; @@ -150,6 +158,65 @@ project: content: print("hello") `; + // The same block after our own save: the serializer recomputes contentHash and the package + // normalizes sortingKey (generateSortingKey -> String(index).padStart(6, '0')). + const rewrittenBySaveYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00Z' +project: + id: e132b172-b114-410e-8331-011517db664f + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + type: code + sortingKey: '000000' + blockGroup: '1' + contentHash: 'sha256:0badc0de1' + content: print("hello") +`; + + const agentYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00Z' +project: + id: e132b172-b114-410e-8331-011517db664f + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-1 + type: agent + sortingKey: '000000' + blockGroup: '1' + content: summarize the dataframe + metadata: + deepnote_agent_model: gpt-5 +`; + + const renamedBlockYaml = ` +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00Z' +project: + id: e132b172-b114-410e-8331-011517db664f + name: Test Project + notebooks: + - id: notebook-1 + name: Notebook 1 + blocks: + - id: block-9 + type: code + sortingKey: '000000' + blockGroup: '1' + content: print("hello") +`; + test('should skip reload when content matches notebook cells', async () => { const uri = Uri.file('/workspace/test.deepnote'); // Create a notebook whose cell content already matches validYaml @@ -178,6 +245,39 @@ project: assert.strictEqual(applyEditCount, 0, 'applyEdit should not be called when cells match'); }); + test('should skip reload when the live notebook only adds ephemeral cells', async () => { + const uri = Uri.file('/workspace/test.deepnote'); + // Ephemeral cells are not serialized — extra live cells must not look like an external edit. + const notebook = createMockNotebook({ + uri, + cells: [ + { + metadata: { id: 'block-1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")', languageId: 'python' } + }, + { + metadata: { id: 'eph-1', is_ephemeral: true, agent_source_block_id: 'agent-1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("agent generated")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(validYaml); + + onDidChangeFile.fire(uri); + + await waitFor(() => readFileCalls > 0); + await new Promise((resolve) => setTimeout(resolve, autoSaveGraceMs)); + + assert.strictEqual(applyEditCount, 0, 'ephemeral-only difference should not trigger a reload'); + assert.strictEqual(saveCount, 0, 'ephemeral-only difference should not trigger a save'); + }); + test('should reload on external change', async () => { const uri = Uri.file('/workspace/test.deepnote'); const notebook = createMockNotebook({ uri, cellCount: 0 }); @@ -400,6 +500,109 @@ project: assert.strictEqual(applyEditCount, 0, 'applyEdit should NOT be called for auto-save (same source)'); }); + test('should skip reload when only save-rewritten block fields differ', async function () { + this.timeout(8000); + const uri = Uri.file('/workspace/test.deepnote'); + // State right after our own save: contentHash recomputed, sortingKey normalized, outputs + // stripped from the main file. None of that is an edit, and reloading would drop the + // agent's ephemeral cells. + const notebook = createMockNotebook({ + uri, + cells: [ + { + metadata: { + id: 'block-1', + __deepnoteBlockId: 'block-1', + __hadOutputs: true, + __deepnotePocket: { + blockGroup: '1', + contentHash: 'sha256:deadbeef', + sortingKey: 'a0', + type: 'code' + } + }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(rewrittenBySaveYaml); + + onDidChangeFile.fire(uri); + + await waitFor(() => readFileCalls > 0); + await new Promise((resolve) => setTimeout(resolve, autoSaveGraceMs)); + + assert.strictEqual(applyEditCount, 0, 'a no-op save must not reload'); + assert.strictEqual(saveCount, 0, 'a no-op save must not trigger another save'); + }); + + test('should reload when only block metadata changed on disk', async function () { + this.timeout(8000); + const uri = Uri.file('/workspace/test.deepnote'); + // Live cell holds the model the file carried before an external editor rewrote it; + // source, kind and language are identical on both sides. + const notebook = createMockNotebook({ + uri, + cells: [ + { + metadata: { + deepnote_agent_model: 'gpt-4o', + id: 'block-1', + __deepnoteBlockId: 'block-1', + __hadOutputs: false, + __deepnotePocket: { blockGroup: '1', sortingKey: '000000', type: 'agent' } + }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'summarize the dataframe', languageId: 'plaintext' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(agentYaml); + + onDidChangeFile.fire(uri); + + await waitFor(() => applyEditCount > 0); + + assert.isAtLeast(applyEditCount, 1, 'a metadata-only external edit must reload'); + }); + + test('should reload when a block id changed on disk', async function () { + this.timeout(8000); + const uri = Uri.file('/workspace/test.deepnote'); + const notebook = createMockNotebook({ + uri, + cells: [ + { + metadata: { + id: 'block-1', + __deepnoteBlockId: 'block-1', + __hadOutputs: false, + __deepnotePocket: { blockGroup: '1', sortingKey: '000000', type: 'code' } + }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(renamedBlockYaml); + + onDidChangeFile.fire(uri); + + await waitFor(() => applyEditCount > 0); + + assert.isAtLeast(applyEditCount, 1, 'an external block-id change must reload'); + }); + suite('normalized one-shot self-write markers', () => { // YAML matching `print("hello")` lives in `validYaml` above. const helloYaml = validYaml; @@ -1136,6 +1339,154 @@ project: fallbackOnDidCreate.dispose(); }); + test('should resolve the block for a metadata-less cell that sits below an ephemeral cell', async () => { + // Catches: indexing the main file's block list by live cell index. Ephemeral agent cells + // exist in the document but are stripped from the file, so every cell below one is offset — + // the metadata-less cell would adopt the wrong block's id and outputs, and that id gets + // written back, leaving two cells claiming one block. + const deepnoteFile: DeepnoteFile = { + version: '1.0', + metadata: { createdAt: '2025-01-01T00:00:00Z' }, + project: { + id: 'e132b172-b114-410e-8331-011517db664f', + name: 'Test Project', + notebooks: [ + { + id: 'notebook-1', + name: 'Notebook 1', + blocks: [] + } + ] + } + }; + // Force casting without metadata + deepnoteFile.project.notebooks[0].blocks = [ + { id: 'block-1', type: 'code', sortingKey: 'a0' } as DeepnoteBlock, + { id: 'block-2', type: 'code', sortingKey: 'a1' } as DeepnoteBlock + ]; + + const mockedManager = mock(); + when(mockedManager.getProjectForNotebook('e132b172-b114-410e-8331-011517db664f', 'notebook-1')).thenReturn( + deepnoteFile + ); + + const offsetDisposables: IDisposableRegistry = []; + const offsetOnDidChange = new EventEmitter(); + const offsetOnDidCreate = new EventEmitter(); + const offsetFsWatcher = mock(); + when(offsetFsWatcher.onDidChange).thenReturn(offsetOnDidChange.event); + when(offsetFsWatcher.onDidCreate).thenReturn(offsetOnDidCreate.event); + when(offsetFsWatcher.dispose()).thenReturn(); + when(mockedVSCodeNamespaces.workspace.createFileSystemWatcher(anything())).thenReturn( + instance(offsetFsWatcher) + ); + + let offsetApplyEditCount = 0; + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { + offsetApplyEditCount++; + return Promise.resolve(true); + }); + + // The vscode mock's NotebookEdit.updateCellMetadata drops its metadata argument, so the + // WorkspaceEdit cannot reveal which block id was written. Stub the static to capture it. + const metadataWrites: Array<{ index: number; metadata: Record }> = []; + sinon.stub(NotebookEdit, 'updateCellMetadata').callsFake((index: number, metadata) => { + metadataWrites.push({ index, metadata: metadata as Record }); + + return {} as NotebookEdit; + }); + + const offsetWatcher = new DeepnoteFileChangeWatcher( + offsetDisposables, + instance(mockedManager), + instance(mockSnapshotService) + ); + offsetWatcher.activate(); + + const notebook = createMockNotebook({ + uri: Uri.file('/workspace/test.deepnote'), + metadata: { + deepnoteProjectId: 'e132b172-b114-410e-8331-011517db664f', + deepnoteNotebookId: 'notebook-1' + }, + cells: [ + { + metadata: { id: 'block-1', type: 'code' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("first")' } + }, + { + // Agent scratch cell: lives in the document, never persisted to the file + metadata: { id: 'eph-1', type: 'code', is_ephemeral: true, agent_source_block_id: 'agent-1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("scratch")' } + }, + { + metadata: { type: 'code' }, // No id — VS Code lost it + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("second")' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + + const newOutputs = new Map([ + [ + 'block-1', + [ + { + output_type: 'execute_result', + data: { 'text/plain': 'First Output' }, + execution_count: 1 + } + ] + ], + [ + 'block-2', + [ + { + output_type: 'execute_result', + data: { 'text/plain': 'Second Output' }, + execution_count: 2 + } + ] + ] + ]); + when(mockSnapshotService.readSnapshot(anything(), anything())).thenReturn(Promise.resolve(newOutputs)); + + offsetOnDidChange.fire( + Uri.file( + '/workspace/snapshots/my-project_e132b172-b114-410e-8331-011517db664f_latest.snapshot.deepnote' + ) + ); + + await waitFor(() => offsetApplyEditCount > 0); + await waitFor(() => metadataWrites.length > 0); + + const writeForLastCell = metadataWrites.find((w) => w.index === 2); + assert.isDefined(writeForLastCell, 'the metadata-less cell at live index 2 should receive a block id'); + assert.strictEqual( + writeForLastCell!.metadata.__deepnoteBlockId, + 'block-2', + 'live index 2 is the second *persisted* cell, so it must resolve to block-2' + ); + assert.notStrictEqual( + writeForLastCell!.metadata.__deepnoteBlockId, + 'block-1', + 'block-1 already belongs to live index 0 — two cells must never claim one block' + ); + + for (const d of offsetDisposables) { + d.dispose(); + } + offsetOnDidChange.dispose(); + offsetOnDidCreate.dispose(); + }); + test('should only update cells whose outputs changed (per-cell updates)', async () => { const snapshotUri = Uri.file( '/workspace/snapshots/my-project_e132b172-b114-410e-8331-011517db664f_latest.snapshot.deepnote' diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index b1669cb86e..699a0d77a1 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -1105,9 +1105,9 @@ 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 + // Environment picker only; execution goes through the real controller on retry. controller.executeHandler = async (cells, doc) => { logger.info( `Placeholder controller execute handler called for ${getDisplayPath(doc.uri)} with ${ @@ -1115,6 +1115,12 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } cells` ); + if (!workspace.isTrusted) { + logger.info(`Workspace is not trusted, skipping environment setup for ${getDisplayPath(doc.uri)}`); + + return; + } + // Create a cancellation token that cancels when the notebook is closed const cts = new CancellationTokenSource(); const closeListener = workspace.onDidCloseNotebookDocument((closedDoc) => { @@ -1133,38 +1139,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return; } - // Environment is now configured, execute the cells through the kernel - const docNotebookKey = getNotebookKey(doc.uri); - const realController = this.notebookControllers.get(docNotebookKey); - - if (!realController) { - logger.error(`No controller found after environment configuration for ${docNotebookKey}`); - - return; - } - - logger.info(`Executing ${cells.length} cells through kernel after environment configuration`); - - // Get or create a kernel for this notebook with the new connection - const kernel = this.kernelProvider.getOrCreate(doc, { - metadata: realController.connection, - controller: realController.controller, - resourceUri: doc.uri - }); - - // Execute cells through the kernel - const kernelExecution = this.kernelProvider.getKernelExecution(kernel); - - for (const cell of cells) { - 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`); + void window.showInformationMessage(l10n.t('Environment ready. Run the cells again to execute them.')); } catch (error) { if (isCancellationError(error)) { logger.info(`Environment setup cancelled for ${getDisplayPath(doc.uri)}`); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts index 8cb56c4aeb..04ca77dc45 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -3,6 +3,7 @@ import * as sinon from 'sinon'; import { anything, deepEqual, instance, mock, verify, when } from 'ts-mockito'; import { DeepnoteKernelAutoSelector } from './deepnoteKernelAutoSelector.node'; import { createMockChildProcess } from '../../kernels/deepnote/deepnoteTestHelpers.node'; +import { createMockCell } from './deepnoteTestHelpers'; import { ServerHandleRegistry } from '../../kernels/deepnote/deepnoteServerHandleRegistry.node'; import { IDeepnoteEnvironmentManager, @@ -21,7 +22,7 @@ import { IConfigurationService } from '../../platform/common/types'; import { IDeepnoteNotebookManager } from '../types'; import { IKernelProvider, IKernel, IJupyterKernelSpec } from '../../kernels/types'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; -import { CancellationError, NotebookDocument, Uri, NotebookController, CancellationToken } from 'vscode'; +import { CancellationError, EventEmitter, NotebookDocument, Uri, NotebookController, CancellationToken } from 'vscode'; import { DeepnoteToolkitMissingError } from '../../platform/errors/deepnoteKernelErrors'; import { DeepnoteEnvironment } from '../../kernels/deepnote/environments/deepnoteEnvironment'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; @@ -1094,6 +1095,67 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { }); }); + suite('Placeholder controller execution', () => { + function createPlaceholder() { + const placeholder = { + supportsExecutionOrder: false, + supportedLanguages: [] as string[], + updateNotebookAffinity: sandbox.stub(), + dispose: sandbox.stub(), + createNotebookCellExecution: sandbox.stub() + } as unknown as NotebookController; + + when( + mockedVSCodeNamespaces.notebooks!.createNotebookController(anything(), anything(), anything()) + ).thenReturn(placeholder); + + const onDidCloseNotebookDocument = new EventEmitter(); + when(mockedVSCodeNamespaces.workspace.onDidCloseNotebookDocument).thenReturn( + onDidCloseNotebookDocument.event + ); + + const internals = selector as unknown as { + createPlaceholderController(notebook: NotebookDocument): NotebookController; + }; + + internals.createPlaceholderController(mockNotebook); + + return placeholder; + } + + const agentCell = createMockCell({ index: 0, metadata: { __deepnotePocket: { type: 'agent' } } }); + const codeCell = createMockCell({ index: 1 }); + + test('configures the environment and executes nothing', async () => { + when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(true); + const placeholder = createPlaceholder(); + const ensureEnvironment = sandbox + .stub(selector, 'ensureEnvironmentConfiguredBeforeExecution') + .resolves(true); + + await placeholder.executeHandler!([agentCell, codeCell], mockNotebook, placeholder); + + assert.isTrue(ensureEnvironment.calledOnce, 'should prompt for an environment'); + assert.isTrue( + (placeholder.createNotebookCellExecution as sinon.SinonStub).notCalled, + 'placeholder must not create executions' + ); + verify(mockKernelProvider.getOrCreate(anything(), anything())).never(); + }); + + test('does nothing at all in an untrusted workspace', async () => { + when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(false); + const placeholder = createPlaceholder(); + const ensureEnvironment = sandbox + .stub(selector, 'ensureEnvironmentConfiguredBeforeExecution') + .resolves(true); + + await placeholder.executeHandler!([agentCell, codeCell], mockNotebook, placeholder); + + assert.isTrue(ensureEnvironment.notCalled, 'should not prompt in an untrusted workspace'); + }); + }); + /** * Every exec in the installer can now be killed, so cancellation surfaces here immediately * instead of after pip has finished anyway. A user-initiated Stop is not a failure and diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts index 27865247d0..e65717a4e0 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.ts @@ -38,6 +38,8 @@ import { } from './deepnoteSchemas'; import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes'; import { Pocket } from '../../platform/deepnote/pocket'; +import { AGENT_MODEL_AUTO, AGENT_MODEL_METADATA_KEY } from './agentCellStatusBarProvider'; +import { generateBlockId, isAgentCell } from './dataConversionUtils'; export const INPUT_BLOCK_TYPES = [ 'input-text', @@ -165,6 +167,7 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation } private registerCommands(): void { + this.disposableRegistry.push(commands.registerCommand(Commands.AddAgentBlock, () => this.addAgentBlock())); this.disposableRegistry.push(commands.registerCommand(Commands.AddSqlBlock, () => this.addSqlBlock())); this.disposableRegistry.push( commands.registerCommand(Commands.AddBigNumberChartBlock, () => this.addBigNumberChartBlock()) @@ -229,6 +232,76 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation ); } + /** + * Inserts an empty agent block below the selection. A notebook may hold at most one; a second + * request reports that and leaves the notebook untouched. + * + * Unlike the other block commands this mints the block id up front: an agent block without one + * gets a fresh random id on every `convertCellToBlock`, so each run would stamp its generated + * cells with a different owner and the stale-run cleanup would never match them. + */ + public async addAgentBlock(): Promise { + const editor = window.activeNotebookEditor; + if (!editor) { + throw new Error(l10n.t('No active notebook editor found')); + } + const document = editor.notebook; + const selection = editor.selection; + const agentBlockExistsMessage = l10n.t('This notebook already contains an agent block.'); + + if (document.getCells().some(isAgentCell)) { + void window.showInformationMessage(agentBlockExistsMessage); + + return; + } + + const blockId = generateBlockId(); + + let alreadyHasAgentBlock = false; + let insertIndex = 0; + + const result = await notebookUpdaterUtils.chainWithPendingUpdates(document, (edit) => { + // Repeated inside the serialized callback: a concurrent invocation passes the check above + // while its edit is still queued, and only here has that edit already applied. + if (document.getCells().some(isAgentCell)) { + alreadyHasAgentBlock = true; + + return; + } + + insertIndex = selection ? selection.end : document.cellCount; + + const newCell = new NotebookCellData(NotebookCellKind.Code, '', 'plaintext'); + newCell.metadata = { + __deepnotePocket: { + type: 'agent' + }, + id: blockId, + __deepnoteBlockId: blockId, + [AGENT_MODEL_METADATA_KEY]: AGENT_MODEL_AUTO + }; + const nbEdit = NotebookEdit.insertCells(insertIndex, [newCell]); + edit.set(document.uri, [nbEdit]); + }); + + if (alreadyHasAgentBlock) { + void window.showInformationMessage(agentBlockExistsMessage); + + return; + } + + if (result !== true) { + throw new Error(l10n.t('Failed to insert agent block')); + } + + this.trackAddBlock('agent'); + + const notebookRange = new NotebookRange(insertIndex, insertIndex + 1); + editor.revealRange(notebookRange, NotebookEditorRevealType.Default); + editor.selection = notebookRange; + await commands.executeCommand('notebook.cell.edit'); + } + public async addSqlBlock(): Promise { const editor = window.activeNotebookEditor; if (!editor) { @@ -590,6 +663,6 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation } private trackAddBlock(blockType: string): void { - this.analytics.trackEvent({ eventName: 'add_block', properties: { blockType } }); + this.analytics.trackEvent({ eventName: 'add_block', properties: { blockType, isEphemeral: false } }); } } diff --git a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts index dd54a44e5e..f75e9ece90 100644 --- a/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts @@ -1,6 +1,6 @@ import { assert } from 'chai'; import * as sinon from 'sinon'; -import { when, reset, anything, mock, instance } from 'ts-mockito'; +import { when, reset, anything, deepEqual, mock, instance, verify } from 'ts-mockito'; import { NotebookCell, NotebookDocument, @@ -8,8 +8,7 @@ import { NotebookRange, NotebookCellKind, NotebookCellData, - WorkspaceEdit, - Uri + WorkspaceEdit } from 'vscode'; import { @@ -21,11 +20,10 @@ import { formatInputBlockCellContent, getInputBlockLanguage } from './inputBlock import { ITelemetryService } from '../../platform/analytics/types'; import { IConfigurationService, IDisposable } from '../../platform/common/types'; import * as notebookUpdater from '../../kernels/execution/notebookUpdater'; -import { createMockedNotebookDocument } from '../../test/datascience/editor-integration/helpers'; import { WrappedError } from '../../platform/errors/types'; import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes'; import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; -import { createMockCell } from './deepnoteTestHelpers'; +import { createMockCell, createMockNotebookWithCells } from './deepnoteTestHelpers'; suite('DeepnoteNotebookCommandListener', () => { let commandListener: DeepnoteNotebookCommandListener; @@ -362,9 +360,9 @@ suite('DeepnoteNotebookCommandListener', () => { }); /** - * Helper to create mock NotebookCell with metadata + * Helper to create NotebookCellData with metadata, for seeding createMockEditor. */ - function createMockCell(content: string, metadata?: Record): NotebookCellData { + function createMockCellData(content: string, metadata?: Record): NotebookCellData { const cell = new NotebookCellData(NotebookCellKind.Code, content, 'json'); if (metadata != null) { cell.metadata = metadata; @@ -373,22 +371,32 @@ suite('DeepnoteNotebookCommandListener', () => { } /** - * Helper to create mock NotebookEditor and NotebookDocument + * Helper to create mock NotebookEditor and NotebookDocument. + * + * Built on createMockNotebookWithCells rather than createMockedNotebookDocument because the + * latter drops NotebookCellData.metadata, which the block commands read. */ function createMockEditor( cellDataArray: NotebookCellData[], selection?: NotebookRange ): { - editor: NotebookEditor; + // revealRange is narrowed to the stub it actually is, so assertions on it need no cast. + editor: NotebookEditor & { revealRange: sinon.SinonStub }; document: NotebookDocument; } { - const uri = Uri.file('/test/notebook.ipynb'); - const document = createMockedNotebookDocument(cellDataArray, {}, uri); + const { notebook: document } = createMockNotebookWithCells( + cellDataArray.map((data) => ({ + kind: data.kind, + languageId: data.languageId, + text: data.value, + metadata: data.metadata + })) + ); const editorSelection = selection != null ? selection : new NotebookRange(0, cellDataArray.length > 0 ? 1 : 0); - const editor: NotebookEditor = { + const editor: NotebookEditor & { revealRange: sinon.SinonStub } = { notebook: document, selection: editorSelection, selections: [editorSelection], @@ -454,7 +462,7 @@ suite('DeepnoteNotebookCommandListener', () => { { description: 'should add input-text block after selection when selection exists', blockType: 'input-text', - existingCells: [createMockCell('{}')], + existingCells: [createMockCellData('{}')], selection: new NotebookRange(0, 1), expectedInsertIndex: 1, expectedVariableName: 'input_1', @@ -563,8 +571,8 @@ suite('DeepnoteNotebookCommandListener', () => { description: 'should generate correct variable name when existing inputs exist', blockType: 'input-text', existingCells: [ - createMockCell('{ "deepnote_variable_name": "input_1" }'), - createMockCell('{ "deepnote_variable_name": "input_2" }') + createMockCellData('{ "deepnote_variable_name": "input_1" }'), + createMockCellData('{ "deepnote_variable_name": "input_2" }') ], selection: new NotebookRange(1, 2), expectedInsertIndex: 2, @@ -574,7 +582,7 @@ suite('DeepnoteNotebookCommandListener', () => { { description: 'should insert at selection.end when selection is in the middle', blockType: 'input-text', - existingCells: [createMockCell('{}'), createMockCell('{}'), createMockCell('{}')], + existingCells: [createMockCellData('{}'), createMockCellData('{}'), createMockCellData('{}')], selection: new NotebookRange(1, 2), expectedInsertIndex: 2, expectedVariableName: 'input_1', @@ -583,7 +591,7 @@ suite('DeepnoteNotebookCommandListener', () => { { description: 'should handle large variable numbers correctly', blockType: 'input-text', - existingCells: [createMockCell('{ "deepnote_variable_name": "input_99" }')], + existingCells: [createMockCellData('{ "deepnote_variable_name": "input_99" }')], selection: undefined, expectedInsertIndex: 1, expectedVariableName: 'input_100', @@ -664,11 +672,8 @@ suite('DeepnoteNotebookCommandListener', () => { }); // Verify reveal and selection were set - assert.isTrue( - (editor.revealRange as sinon.SinonStub).calledOnce, - 'Should reveal the new cell range' - ); - const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + assert.isTrue(editor.revealRange.calledOnce, 'Should reveal the new cell range'); + const revealCall = editor.revealRange.firstCall; assert.equal(revealCall.args[0].start, expectedInsertIndex, 'Should reveal correct range start'); assert.equal(revealCall.args[0].end, expectedInsertIndex + 1, 'Should reveal correct range end'); }); @@ -755,8 +760,8 @@ suite('DeepnoteNotebookCommandListener', () => { ); // Verify reveal and selection were set - assert.isTrue((editor.revealRange as sinon.SinonStub).calledOnce, 'Should reveal the new cell range'); - const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + assert.isTrue(editor.revealRange.calledOnce, 'Should reveal the new cell range'); + const revealCall = editor.revealRange.firstCall; assert.equal(revealCall.args[0].start, 0, 'Should reveal correct range start'); assert.equal(revealCall.args[0].end, 1, 'Should reveal correct range end'); assert.equal(revealCall.args[1], 0, 'Should use NotebookEditorRevealType.Default (value 0)'); @@ -764,7 +769,7 @@ suite('DeepnoteNotebookCommandListener', () => { test('should add SQL block after selection when selection exists', async () => { // Setup mocks - const existingCells = [createMockCell('{}'), createMockCell('{}')]; + const existingCells = [createMockCellData('{}'), createMockCellData('{}')]; const selection = new NotebookRange(1, 2); const { editor } = createMockEditor(existingCells, selection); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -787,8 +792,8 @@ suite('DeepnoteNotebookCommandListener', () => { test('should generate correct variable name when existing df variables exist', async () => { // Setup mocks with existing df variables const existingCells = [ - createMockCell('{ "deepnote_variable_name": "df_1" }'), - createMockCell('{ "deepnote_variable_name": "df_2" }') + createMockCellData('{ "deepnote_variable_name": "df_1" }'), + createMockCellData('{ "deepnote_variable_name": "df_2" }') ]; const { editor } = createMockEditor(existingCells, undefined); const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -807,8 +812,8 @@ suite('DeepnoteNotebookCommandListener', () => { test('should ignore input variables when generating df variable name', async () => { // Setup mocks with input variables (should not affect df numbering) const existingCells = [ - createMockCell('{ "deepnote_variable_name": "input_10" }'), - createMockCell('{ "deepnote_variable_name": "df_2" }') + createMockCellData('{ "deepnote_variable_name": "input_10" }'), + createMockCellData('{ "deepnote_variable_name": "df_2" }') ]; const { editor } = createMockEditor(existingCells, undefined); const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -849,6 +854,202 @@ suite('DeepnoteNotebookCommandListener', () => { }); }); + suite('addAgentBlock', () => { + function insertedCell(getCapturedNotebookEdits: () => any[] | null) { + const edits = getCapturedNotebookEdits()!; + assert.equal(edits.length, 1, 'Should have one notebook edit'); + + const notebookEdit = edits[0] as any; + assert.equal(notebookEdit.newCells.length, 1, 'Should insert one cell'); + + return notebookEdit.newCells[0]; + } + + test('should add an empty plaintext agent block at the end when no selection exists', async () => { + const { editor, document } = createMockEditor([], undefined); + const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + assert.isTrue(chainStub.calledOnce, 'chainWithPendingUpdates should be called once'); + assert.equal(chainStub.firstCall.args[0], document, 'Should be called with correct document'); + + const newCell = insertedCell(getCapturedNotebookEdits); + assert.equal(newCell.kind, NotebookCellKind.Code, 'Should be a code cell'); + assert.equal(newCell.languageId, 'plaintext', 'Should have plaintext language'); + assert.equal(newCell.value, '', 'Should have empty content'); + assert.equal(newCell.metadata.__deepnotePocket.type, 'agent', 'Should have agent pocket type'); + + assert.isTrue(editor.revealRange.calledOnce, 'Should reveal the new cell range'); + const revealCall = editor.revealRange.firstCall; + assert.equal(revealCall.args[0].start, 0, 'Should reveal correct range start'); + assert.equal(revealCall.args[0].end, 1, 'Should reveal correct range end'); + }); + + test('should add the agent block after the selection when one exists', async () => { + const existingCells = [createMockCellData('{}'), createMockCellData('{}')]; + const { editor, document } = createMockEditor(existingCells, new NotebookRange(1, 2)); + const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + insertedCell(getCapturedNotebookEdits); + assert.equal(chainStub.firstCall.args[0], document, 'Should edit the active document'); + + const revealCall = editor.revealRange.firstCall; + assert.equal(revealCall.args[0].start, 2, 'Should insert below the selection'); + assert.equal(revealCall.args[0].end, 3, 'Should select only the new cell'); + }); + + test('should mint a block id under both keys so runs keep a stable owner', async () => { + // Catches: an id-less agent block, which gets a fresh random id on every + // convertCellToBlock — its generated cells would never be matched back to it. + const { editor } = createMockEditor([], undefined); + const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + const { metadata } = insertedCell(getCapturedNotebookEdits); + assert.match(metadata.id, /^[0-9a-f]{32}$/, 'Should mint a 32-char hex block id'); + assert.equal(metadata.__deepnoteBlockId, metadata.id, 'Backup id key must match id'); + }); + + test('should give each notebook its own agent block id', async () => { + // Catches: a hoisted/constant id, which would make two agent blocks fight over the + // same generated cells. Two notebooks, because one notebook only ever gets one block. + const { editor } = createMockEditor([], undefined); + const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + const first = insertedCell(getCapturedNotebookEdits).metadata.id; + + const { editor: otherEditor } = createMockEditor([], undefined); + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(otherEditor); + + await commandListener.addAgentBlock(); + const second = insertedCell(getCapturedNotebookEdits).metadata.id; + + assert.notEqual(first, second, 'Each agent block needs its own id'); + }); + + test('should persist the default model rather than leaving the key absent', async () => { + // Catches: omitting deepnote_agent_model, which reaches openai() as undefined. + const { editor } = createMockEditor([], undefined); + const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + const { metadata } = insertedCell(getCapturedNotebookEdits); + assert.equal(metadata.deepnote_agent_model, 'auto', 'Should persist the auto default'); + }); + + test('should refuse a second agent block and leave the notebook untouched', async () => { + const { editor } = createMockEditor([ + createMockCellData('existing agent', { __deepnotePocket: { type: 'agent' }, id: 'agent-block-1' }), + createMockCellData('user code') + ]); + const { chainStub } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + assert.isFalse(chainStub.called, 'Must not edit a notebook that already has an agent block'); + verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).once(); + assert.isFalse(editor.revealRange.called, 'Must not reveal anything'); + verify(mockTelemetryService.trackEvent(anything())).never(); + }); + + test('should report the added agent block to analytics', async () => { + const { editor } = createMockEditor([], undefined); + mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + verify( + mockTelemetryService.trackEvent( + deepEqual({ eventName: 'add_block', properties: { blockType: 'agent', isEphemeral: false } }) + ) + ).once(); + }); + + test('should insert only one agent block when two invocations race', async () => { + // Catches: an existence check that runs before the queued update — both invocations + // pass it while neither edit has applied yet. + const { editor, document } = createMockEditor([], undefined); + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(editor); + when(mockedVSCodeNamespaces.commands.executeCommand(anything())).thenResolve(undefined as any); + + const cells = document.getCells(); + const insertedCells: NotebookCellData[] = []; + let pending: Promise = Promise.resolve(); + + // Mirrors chainWithPendingUpdates: a callback runs only once the previous edit applied. + sandbox + .stub(notebookUpdater.notebookUpdaterUtils, 'chainWithPendingUpdates') + .callsFake((_doc: NotebookDocument, callback: (edit: WorkspaceEdit) => void) => { + const applied = pending.then(() => { + const edit = new WorkspaceEdit(); + sandbox.stub(edit, 'set').callsFake((_uri, edits) => { + for (const newCell of (edits[0] as any).newCells as NotebookCellData[]) { + insertedCells.push(newCell); + cells.push( + createMockCell({ + metadata: newCell.metadata, + index: cells.length, + notebook: document + }) + ); + } + }); + callback(edit); + + return true; + }); + pending = applied; + + return applied; + }); + + await Promise.all([commandListener.addAgentBlock(), commandListener.addAgentBlock()]); + + assert.equal(insertedCells.length, 1, 'Should insert exactly one agent block'); + assert.equal(cells.length, 1, 'Notebook should end up with a single cell'); + }); + + test('should still add the block when other cells carry no agent pocket', async () => { + // Catches: a guard that trips on any cell, blocking the first agent block outright. + const { editor } = createMockEditor([ + createMockCellData('user code', { __deepnotePocket: { type: 'code' }, id: 'code-block-1' }), + createMockCellData('scratch', { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }) + ]); + const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); + + await commandListener.addAgentBlock(); + + assert.isTrue(chainStub.calledOnce, 'Should insert the first agent block'); + assert.equal(insertedCell(getCapturedNotebookEdits).metadata.__deepnotePocket.type, 'agent'); + verify(mockedVSCodeNamespaces.window.showInformationMessage(anything())).never(); + }); + + test('should throw error when no active editor exists', async () => { + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(undefined); + + await assert.isRejected(commandListener.addAgentBlock(), Error, 'No active notebook editor found'); + }); + + test('should throw error when chainWithPendingUpdates fails', async () => { + const { editor } = createMockEditor([], undefined); + when(mockedVSCodeNamespaces.window.activeNotebookEditor).thenReturn(editor); + + sandbox.replace( + notebookUpdater.notebookUpdaterUtils, + 'chainWithPendingUpdates', + sinon.stub().resolves(false) + ); + + await assert.isRejected(commandListener.addAgentBlock(), Error, 'Failed to insert agent block'); + }); + }); + suite('addBigNumberChartBlock', () => { test('should add big number block at the end when no selection exists', async () => { // Setup mocks @@ -887,8 +1088,8 @@ suite('DeepnoteNotebookCommandListener', () => { assert.equal(newCell.metadata.__deepnotePocket.type, 'big-number', 'Should have big-number type'); // Verify reveal and selection were set - assert.isTrue((editor.revealRange as sinon.SinonStub).calledOnce, 'Should reveal the new cell range'); - const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + assert.isTrue(editor.revealRange.calledOnce, 'Should reveal the new cell range'); + const revealCall = editor.revealRange.firstCall; assert.equal(revealCall.args[0].start, 0, 'Should reveal correct range start'); assert.equal(revealCall.args[0].end, 1, 'Should reveal correct range end'); assert.equal(revealCall.args[1], 0, 'Should use NotebookEditorRevealType.Default (value 0)'); @@ -896,7 +1097,7 @@ suite('DeepnoteNotebookCommandListener', () => { test('should add big number block after selection when selection exists', async () => { // Setup mocks - const existingCells = [createMockCell('{}'), createMockCell('{}')]; + const existingCells = [createMockCellData('{}'), createMockCellData('{}')]; const selection = new NotebookRange(0, 1); const { editor } = createMockEditor(existingCells, selection); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -918,7 +1119,7 @@ suite('DeepnoteNotebookCommandListener', () => { test('should insert at correct position in the middle of notebook', async () => { // Setup mocks - const existingCells = [createMockCell('{}'), createMockCell('{}'), createMockCell('{}')]; + const existingCells = [createMockCellData('{}'), createMockCellData('{}'), createMockCellData('{}')]; const selection = new NotebookRange(1, 2); const { editor } = createMockEditor(existingCells, selection); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -1023,8 +1224,8 @@ suite('DeepnoteNotebookCommandListener', () => { assert.equal(newCell.metadata.__deepnotePocket.type, 'visualization', 'Should have visualization type'); // Verify reveal and selection were set - assert.isTrue((editor.revealRange as sinon.SinonStub).calledOnce, 'Should reveal the new cell range'); - const revealCall = (editor.revealRange as sinon.SinonStub).firstCall; + assert.isTrue(editor.revealRange.calledOnce, 'Should reveal the new cell range'); + const revealCall = editor.revealRange.firstCall; assert.equal(revealCall.args[0].start, 0, 'Should reveal correct range start'); assert.equal(revealCall.args[0].end, 1, 'Should reveal correct range end'); assert.equal(revealCall.args[1], 0, 'Should use NotebookEditorRevealType.Default (value 0)'); @@ -1032,7 +1233,7 @@ suite('DeepnoteNotebookCommandListener', () => { test('should add chart block after selection when selection exists', async () => { // Setup mocks - const existingCells = [createMockCell('{}'), createMockCell('{}')]; + const existingCells = [createMockCellData('{}'), createMockCellData('{}')]; const selection = new NotebookRange(0, 1); const { editor } = createMockEditor(existingCells, selection); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -1055,8 +1256,8 @@ suite('DeepnoteNotebookCommandListener', () => { test('should use hardcoded variable name df_1', async () => { // Setup mocks with existing df variables const existingCells = [ - createMockCell('{ "deepnote_variable_name": "df_1" }'), - createMockCell('{ "variable": "df_2" }') + createMockCellData('{ "deepnote_variable_name": "df_1" }'), + createMockCellData('{ "variable": "df_2" }') ]; const { editor } = createMockEditor(existingCells, undefined); const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -1076,9 +1277,9 @@ suite('DeepnoteNotebookCommandListener', () => { test('should always use df_1 regardless of existing variables', async () => { // Setup mocks with various existing variables const existingCells = [ - createMockCell('{ "deepnote_variable_name": "input_10" }'), - createMockCell('{ "deepnote_variable_name": "df_5" }'), - createMockCell('{ "variable": "df_2" }') + createMockCellData('{ "deepnote_variable_name": "input_10" }'), + createMockCellData('{ "deepnote_variable_name": "df_5" }'), + createMockCellData('{ "variable": "df_2" }') ]; const { editor } = createMockEditor(existingCells, undefined); const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); @@ -1097,7 +1298,7 @@ suite('DeepnoteNotebookCommandListener', () => { test('should insert at correct position in the middle of notebook', async () => { // Setup mocks - const existingCells = [createMockCell('{}'), createMockCell('{}'), createMockCell('{}')]; + const existingCells = [createMockCellData('{}'), createMockCellData('{}'), createMockCellData('{}')]; const selection = new NotebookRange(1, 2); const { editor } = createMockEditor(existingCells, selection); const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor); diff --git a/src/notebooks/deepnote/deepnoteSecretStore.ts b/src/notebooks/deepnote/deepnoteSecretStore.ts new file mode 100644 index 0000000000..7a4844b51f --- /dev/null +++ b/src/notebooks/deepnote/deepnoteSecretStore.ts @@ -0,0 +1,47 @@ +import { l10n, window } from 'vscode'; + +import { IEncryptedStorage } from '../../platform/common/application/types'; + +const AGENT_SERVICE_NAME = 'deepnote-agent'; +const OPENAI_API_KEY = 'openAiApiKey'; + +async function retrieveOpenAiApiKey(storage: IEncryptedStorage): Promise { + const value = await storage.retrieve(AGENT_SERVICE_NAME, OPENAI_API_KEY); + + return value && value.length > 0 ? value : undefined; +} + +export async function clearOpenAiApiKey(storage: IEncryptedStorage): Promise { + await storage.store(AGENT_SERVICE_NAME, OPENAI_API_KEY, undefined); +} + +export async function getOrPromptOpenAiApiKey(storage: IEncryptedStorage): Promise { + const value = (await retrieveOpenAiApiKey(storage)) ?? (await promptForOpenAiApiKey(storage)); + + if (!value) { + throw new Error( + l10n.t('OpenAI API key is not set. Use the command "Deepnote: Set OpenAI API Key" to configure it.') + ); + } + + return value; +} + +export async function promptForOpenAiApiKey(storage: IEncryptedStorage): Promise { + const input = await window.showInputBox({ + prompt: l10n.t('Enter your OpenAI API key'), + placeHolder: l10n.t('sk-...'), + password: true, + ignoreFocusOut: true + }); + + const trimmed = input?.trim(); + + if (!trimmed) { + return undefined; + } + + await storage.store(AGENT_SERVICE_NAME, OPENAI_API_KEY, trimmed); + + return trimmed; +} diff --git a/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts new file mode 100644 index 0000000000..d935661c12 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts @@ -0,0 +1,119 @@ +import { assert } from 'chai'; +import { anything, instance, mock, when } from 'ts-mockito'; + +import { IEncryptedStorage } from '../../platform/common/application/types'; +import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; +import { clearOpenAiApiKey, getOrPromptOpenAiApiKey, promptForOpenAiApiKey } from './deepnoteSecretStore'; + +// The real EncryptedStorage namespaces secrets as `${service}.${key}`; the fake mirrors that so a +// store/retrieve mismatch between the two service names would surface here. +const STORED_KEY = 'deepnote-agent.openAiApiKey'; + +suite('deepnoteSecretStore', () => { + let storageData: Map; + let encryptedStorage: IEncryptedStorage; + let promptCount: number; + + function whenPromptReturns(value: string | undefined) { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenCall(() => { + promptCount++; + + return Promise.resolve(value); + }); + } + + setup(() => { + storageData = new Map(); + promptCount = 0; + encryptedStorage = mock(); + + when(encryptedStorage.store(anything(), anything(), anything())).thenCall( + (service: string, key: string, value: string | undefined) => { + if (value === undefined) { + storageData.delete(`${service}.${key}`); + } else { + storageData.set(`${service}.${key}`, value); + } + + return Promise.resolve(); + } + ); + + when(encryptedStorage.retrieve(anything(), anything())).thenCall((service: string, key: string) => + Promise.resolve(storageData.get(`${service}.${key}`)) + ); + + whenPromptReturns(undefined); + }); + + suite('promptForOpenAiApiKey', () => { + test('trims, stores and returns the entered key', async () => { + whenPromptReturns(' sk-entered '); + + const value = await promptForOpenAiApiKey(instance(encryptedStorage)); + + assert.strictEqual(value, 'sk-entered'); + assert.strictEqual(storageData.get(STORED_KEY), 'sk-entered'); + }); + + for (const input of [undefined, ' ']) { + test(`stores nothing and returns undefined when the user enters ${JSON.stringify(input)}`, async () => { + whenPromptReturns(input); + + const value = await promptForOpenAiApiKey(instance(encryptedStorage)); + + assert.isUndefined(value); + assert.isFalse(storageData.has(STORED_KEY)); + }); + } + }); + + suite('clearOpenAiApiKey', () => { + test('deletes the stored key', async () => { + storageData.set(STORED_KEY, 'sk-stored'); + + await clearOpenAiApiKey(instance(encryptedStorage)); + + assert.isFalse(storageData.has(STORED_KEY)); + }); + }); + + suite('getOrPromptOpenAiApiKey', () => { + test('returns the stored key without prompting', async () => { + storageData.set(STORED_KEY, 'sk-stored'); + + const value = await getOrPromptOpenAiApiKey(instance(encryptedStorage)); + + assert.strictEqual(value, 'sk-stored'); + assert.strictEqual(promptCount, 0); + }); + + test('prompts and returns the entered key when nothing is stored', async () => { + whenPromptReturns('sk-prompted'); + + const value = await getOrPromptOpenAiApiKey(instance(encryptedStorage)); + + assert.strictEqual(value, 'sk-prompted'); + assert.strictEqual(promptCount, 1); + }); + + test('treats an empty stored value as missing and prompts', async () => { + storageData.set(STORED_KEY, ''); + whenPromptReturns('sk-prompted'); + + const value = await getOrPromptOpenAiApiKey(instance(encryptedStorage)); + + assert.strictEqual(value, 'sk-prompted'); + assert.strictEqual(promptCount, 1); + }); + + test('throws when nothing is stored and the user cancels the prompt', async () => { + try { + await getOrPromptOpenAiApiKey(instance(encryptedStorage)); + assert.fail('Should have thrown'); + } catch (e) { + assert.include((e as Error).message, 'OpenAI API key is not set'); + } + }); + }); +}); diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 5d029bcf12..8e57fa7596 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -7,6 +7,7 @@ import { workspace, type CancellationToken, type NotebookData, type NotebookSeri import { logger } from '../../platform/logging'; import { IDeepnoteNotebookManager } from '../types'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; +import { getBlockId, isEphemeralCell } from './dataConversionUtils'; import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; import { SnapshotService } from './snapshots/snapshotService'; import { computeHash } from '../../platform/common/crypto'; @@ -230,11 +231,16 @@ 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`); + 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' @@ -244,13 +250,17 @@ 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`); + // An id the cell still carries is its real identity, so only cells that arrived without one may be + // re-identified. convertCellsToBlocks maps 1:1 in order, so index i lines up with nonEphemeralCells[i]. + const recoverableBlocks = new Set(blocks.filter((_, index) => !getBlockId(nonEphemeralCells[index]))); + // Try to recover block IDs from original blocks when VS Code fails to preserve metadata // This uses content-based matching as a fallback when metadata.id is missing - this.recoverBlockIdsFromOriginal(blocks, notebook.blocks ?? []); + this.recoverBlockIdsFromOriginal(blocks, notebook.blocks ?? [], recoverableBlocks); // Log block IDs after conversion and recovery for (let i = 0; i < blocks.length; i++) { @@ -258,7 +268,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()) { @@ -498,8 +508,13 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { * Uses content-based matching as a fallback strategy to recover id, sortingKey, and blockGroup. * @param blocks Blocks converted from cells (may have generated values if metadata was lost) * @param originalBlocks Original blocks from the stored project + * @param recoverableBlocks Blocks whose cell carried no id; only these may take an original's identity */ - private recoverBlockIdsFromOriginal(blocks: DeepnoteBlock[], originalBlocks: DeepnoteBlock[]): void { + private recoverBlockIdsFromOriginal( + blocks: DeepnoteBlock[], + originalBlocks: DeepnoteBlock[], + recoverableBlocks: Set + ): void { // Build a map of original blocks by content for quick lookup // Key: content (trimmed), Value: array of blocks with that content (in case of duplicates) const contentToOriginalBlocks = new Map(); @@ -527,6 +542,12 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { let recoveredCount = 0; for (const block of blocks) { + // A cell that carried an id already owns its identity - an agent block mints one when it is inserted - + // so it must never adopt the id of a block deleted in the same save. + if (!recoverableBlocks.has(block)) { + continue; + } + // Skip if this block already has an original ID if (claimedIds.has(block.id)) { continue; diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index 6af7e32ee7..8a00f04a1a 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -214,6 +214,71 @@ project: 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', 'notebook-1', projectData); + + 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('correct-sibling save (Chunk 2 anti-regression)', () => { const sharedProjectId = 'shared-project'; const nbA = 'sibling-a'; @@ -663,6 +728,282 @@ project: 'Block ID should be newly generated when content differs' ); }); + + test('should keep a minted agent block ID when a deleted block had the same content', async () => { + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2023-01-01T00:00:00Z', + modifiedAt: '2023-01-02T00:00:00Z' + }, + project: { + id: 'project-agent-id', + name: 'Agent ID Test', + notebooks: [ + { + id: 'notebook-1', + name: 'Test Notebook', + blocks: [ + { + blockGroup: 'deleted-group', + id: 'deleted-block-id', + content: '', + sortingKey: 'a0', + metadata: {}, + type: 'code' + } + ], + executionMode: 'block', + isModule: false + } + ], + settings: {} + } + }; + + manager.storeOriginalProject('project-agent-id', 'notebook-1', projectData); + + // The empty code block was deleted and an empty agent block added in the same save + const notebookData = { + cells: [ + { + kind: 2, + value: '', + languageId: 'plaintext', + metadata: { + id: 'minted-agent-id', + __deepnoteBlockId: 'minted-agent-id', + __deepnotePocket: { type: 'agent' } + } + } + ], + metadata: { + deepnoteProjectId: 'project-agent-id', + deepnoteNotebookId: 'notebook-1' + } + }; + + const result = await serializer.serializeNotebook(notebookData as any, {} as any); + const yamlString = new TextDecoder().decode(result); + const parsedResult = deserializeDeepnoteFile(yamlString); + + const notebook = parsedResult.project.notebooks.find((nb) => nb.id === 'notebook-1'); + assert.isDefined(notebook); + assert.strictEqual(notebook!.blocks[0].id, 'minted-agent-id', 'Agent block should keep its minted ID'); + assert.notStrictEqual( + notebook!.blocks[0].blockGroup, + 'deleted-group', + 'Agent block should not inherit the deleted block blockGroup' + ); + }); + + test('should keep an ID the cell carried even when the deleted block has the same type', async () => { + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2023-01-01T00:00:00Z', + modifiedAt: '2023-01-02T00:00:00Z' + }, + project: { + id: 'project-same-type', + name: 'Same Type Test', + notebooks: [ + { + id: 'notebook-1', + name: 'Test Notebook', + blocks: [ + { + blockGroup: 'deleted-group', + id: 'deleted-code-id', + content: '', + sortingKey: 'a0', + metadata: {}, + type: 'code' + } + ], + executionMode: 'block', + isModule: false + } + ], + settings: {} + } + }; + + manager.storeOriginalProject('project-same-type', 'notebook-1', projectData); + + const notebookData = { + cells: [ + { + kind: 2, + value: '', + languageId: 'python', + metadata: { + id: 'minted-code-id', + __deepnoteBlockId: 'minted-code-id' + } + } + ], + metadata: { + deepnoteProjectId: 'project-same-type', + deepnoteNotebookId: 'notebook-1' + } + }; + + const result = await serializer.serializeNotebook(notebookData as any, {} as any); + const yamlString = new TextDecoder().decode(result); + const parsedResult = deserializeDeepnoteFile(yamlString); + + const notebook = parsedResult.project.notebooks.find((nb) => nb.id === 'notebook-1'); + assert.isDefined(notebook); + assert.strictEqual(notebook!.blocks[0].id, 'minted-code-id', 'Block should keep the ID its cell carried'); + }); + + test('should not recover an ID that another cell still carries', async () => { + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2023-01-01T00:00:00Z', + modifiedAt: '2023-01-02T00:00:00Z' + }, + project: { + id: 'project-claimed-id', + name: 'Claimed ID Test', + notebooks: [ + { + id: 'notebook-1', + name: 'Test Notebook', + blocks: [ + { + blockGroup: 'group-kept', + id: 'kept-id', + content: '', + sortingKey: 'a0', + metadata: {}, + type: 'code' + }, + { + blockGroup: 'group-stripped', + id: 'stripped-id', + content: '', + sortingKey: 'a1', + metadata: {}, + type: 'code' + } + ], + executionMode: 'block', + isModule: false + } + ], + settings: {} + } + }; + + manager.storeOriginalProject('project-claimed-id', 'notebook-1', projectData); + + // Both blocks are empty, so content matching alone cannot tell them apart + const notebookData = { + cells: [ + { + kind: 2, + value: '', + languageId: 'python', + metadata: { + id: 'kept-id', + __deepnoteBlockId: 'kept-id', + __deepnotePocket: { type: 'code', sortingKey: 'a0', blockGroup: 'group-kept' } + } + }, + { + kind: 2, + value: '', + languageId: 'python', + metadata: {} + } + ], + metadata: { + deepnoteProjectId: 'project-claimed-id', + deepnoteNotebookId: 'notebook-1' + } + }; + + const result = await serializer.serializeNotebook(notebookData as any, {} as any); + const yamlString = new TextDecoder().decode(result); + const parsedResult = deserializeDeepnoteFile(yamlString); + + const notebook = parsedResult.project.notebooks.find((nb) => nb.id === 'notebook-1'); + assert.isDefined(notebook); + assert.strictEqual(notebook!.blocks[0].id, 'kept-id', 'Cell that carried an ID should keep it'); + assert.strictEqual(notebook!.blocks[1].id, 'stripped-id', 'Stripped cell should take the remaining ID'); + }); + + test('should recover IDs by cell position after ephemeral cells are dropped', async () => { + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2023-01-01T00:00:00Z', + modifiedAt: '2023-01-02T00:00:00Z' + }, + project: { + id: 'project-ephemeral-offset', + name: 'Ephemeral Offset Test', + notebooks: [ + { + id: 'notebook-1', + name: 'Test Notebook', + blocks: [ + { + blockGroup: 'group-1', + id: 'real-block-id', + content: 'print("kept")', + sortingKey: 'a0', + metadata: {}, + type: 'code' + } + ], + executionMode: 'block', + isModule: false + } + ], + settings: {} + } + }; + + manager.storeOriginalProject('project-ephemeral-offset', 'notebook-1', projectData); + + const notebookData = { + cells: [ + { + kind: 2, + value: 'print("scratch")', + languageId: 'python', + metadata: { + id: 'ephemeral-block-id', + is_ephemeral: true, + agent_source_block_id: 'agent-1' + } + }, + { + kind: 2, + value: 'print("kept")', + languageId: 'python', + metadata: {} + } + ], + metadata: { + deepnoteProjectId: 'project-ephemeral-offset', + deepnoteNotebookId: 'notebook-1' + } + }; + + const result = await serializer.serializeNotebook(notebookData as any, {} as any); + const yamlString = new TextDecoder().decode(result); + const parsedResult = deserializeDeepnoteFile(yamlString); + + const notebook = parsedResult.project.notebooks.find((nb) => nb.id === 'notebook-1'); + assert.isDefined(notebook); + assert.strictEqual(notebook!.blocks.length, 1, 'Ephemeral cell should be excluded'); + assert.strictEqual(notebook!.blocks[0].id, 'real-block-id', 'Stripped cell should recover its ID'); + }); }); suite('integration scenarios', () => { diff --git a/src/notebooks/deepnote/deepnoteTestHelpers.ts b/src/notebooks/deepnote/deepnoteTestHelpers.ts index c28ecc04bf..4abbc16b4c 100644 --- a/src/notebooks/deepnote/deepnoteTestHelpers.ts +++ b/src/notebooks/deepnote/deepnoteTestHelpers.ts @@ -4,7 +4,10 @@ import { NotebookCellKind, NotebookCellOutput, NotebookDocument, + Position, TextDocument, + Range, + TextLine, Uri, WorkspaceFolder } from 'vscode'; @@ -29,6 +32,8 @@ export interface CreateMockCellOptions { notebookUri?: Uri; notebookMetadata?: Record; index?: number; + mime?: string; + notebook?: NotebookDocument; } /** @@ -47,6 +52,7 @@ export interface CreateMockNotebookOptions { notebookType?: string; uri?: Uri; metadata?: Record; + cells?: NotebookCell[]; } /** @@ -56,13 +62,45 @@ export interface CreateMockNotebookOptions { * @returns A mock NotebookDocument */ export function createMockNotebook(options?: CreateMockNotebookOptions): NotebookDocument { - const { notebookType = 'deepnote', uri = Uri.file('/test/notebook.deepnote'), metadata = {} } = options ?? {}; + const { + notebookType = 'deepnote', + uri = Uri.file('/test/notebook.deepnote'), + metadata = {}, + cells = [] + } = options ?? {}; return { uri, notebookType, - metadata - } as NotebookDocument; + metadata, + get cellCount() { + return cells.length; + }, + // Mirrors VS Code: the index is clamped to the notebook rather than throwing. + cellAt: (index: number) => cells[Math.min(Math.max(index, 0), cells.length - 1)] ?? ({} as NotebookCell), + getCells: () => cells, + version: 1, + isDirty: false, + isUntitled: false, + isClosed: false, + save: async () => true + } satisfies NotebookDocument; +} + +/** + * Builds one mock notebook and cells that share it (correct `index` and `notebook` references). + */ +export function createMockNotebookWithCells( + cellOptions: Omit[] +): { cells: NotebookCell[]; notebook: NotebookDocument } { + const cells: NotebookCell[] = []; + const notebook = createMockNotebook({ cells }); + + for (let index = 0; index < cellOptions.length; index++) { + cells.push(createMockCell({ ...cellOptions[index], index, notebook })); + } + + return { cells, notebook }; } /** @@ -101,24 +139,28 @@ export function createMockCell(options?: CreateMockCellOptions): NotebookCell { outputs = [], notebookType = 'deepnote', notebookUri = Uri.file('/test/notebook.deepnote'), - index = 0 + index = 0, + mime = 'text/plain' } = opts; // Preserve explicit undefined for metadata fields - const metadata = Object.prototype.hasOwnProperty.call(opts, 'metadata') ? opts.metadata : {}; + const metadata = 'metadata' in opts ? opts.metadata ?? {} : {}; const notebookMetadata = Object.prototype.hasOwnProperty.call(opts, 'notebookMetadata') ? opts.notebookMetadata : {}; - const notebook = createMockNotebook({ - notebookType, - uri: notebookUri, - metadata: notebookMetadata - }); + const notebook = + opts.notebook ?? + createMockNotebook({ + notebookType, + uri: notebookUri, + metadata: notebookMetadata + }); + const resolvedUri = notebook.uri; - const cellPath = `${notebookUri.path}#cell${index}`; + const cellPath = `${resolvedUri.path}#cell${index}`; - const document = { + const document: TextDocument = { uri: Uri.file(cellPath), fileName: cellPath, isUntitled: false, @@ -130,23 +172,25 @@ export function createMockCell(options?: CreateMockCellOptions): NotebookCell { save: async () => true, eol: 1, lineCount: 1, - lineAt: () => ({ text: '' }) as unknown, + lineAt: () => ({ text: '' }) as unknown as TextLine, offsetAt: () => 0, - positionAt: () => ({}) as unknown, - validateRange: () => ({}) as unknown, - validatePosition: () => ({}) as unknown, - getWordRangeAtPosition: () => undefined - } as unknown as TextDocument; + positionAt: () => new Position(0, 0), + validateRange: () => new Range(new Position(0, 0), new Position(0, 0)), + validatePosition: () => new Position(0, 0), + getWordRangeAtPosition: () => undefined, + encoding: 'utf-8' + }; return { index, + mime, notebook, kind, document, metadata, outputs, executionSummary: undefined - } as unknown as NotebookCell; + }; } /** A Deepnote code block (whole-file YAML shape); override any field. */ diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts new file mode 100644 index 0000000000..f336454226 --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts @@ -0,0 +1,76 @@ +import { + CancellationToken, + Disposable, + EventEmitter, + NotebookCell, + NotebookCellStatusBarItem, + NotebookCellStatusBarItemProvider, + l10n, + notebooks, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { isEphemeralCell } from './dataConversionUtils'; + +const EPHEMERAL_INDICATOR_PRIORITY = 1000; + +@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 ?? null)]; + } + + private createEphemeralIndicatorItem(agentSourceBlockId: string | null): NotebookCellStatusBarItem { + const tooltipLines = [l10n.t('Auto-generated ephemeral block')]; + if (agentSourceBlockId) { + tooltipLines.push(l10n.t('Source agent block: {0}', agentSourceBlockId)); + } + + 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..42b38a4617 --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts @@ -0,0 +1,121 @@ +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 undefined when is_ephemeral is false', () => { + const cell = createMockCell({ metadata: { is_ephemeral: false } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined when is_ephemeral is not set', () => { + 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 is_ephemeral is a non-boolean truthy value', () => { + const cell = createMockCell({ metadata: { is_ephemeral: 'true' } }); + 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: { is_ephemeral: true } }); + const items = provider.provideCellStatusBarItems(cell, cancelledToken); + + expect(items).to.be.undefined; + }); + }); + + suite('Status Bar Item Properties', () => { + test('Should set ephemeral status bar item properties', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + // Catches: an actionable item creeping back in — clearing is the agent block's button, + // and an ephemeral cell only labels itself. + expect(items).to.have.lengthOf(1); + expect(items[0].text).to.include('$(sparkle)'); + expect(items[0].text).to.include('Ephemeral'); + expect(items[0].alignment).to.equal(1); + expect(items[0].priority).to.equal(1000); + expect(items[0].command).to.be.undefined; + }); + }); + + suite('Tooltip', () => { + test('Should describe ephemeral tooltip without source block when agent_source_block_id is absent', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[0].tooltip).to.include('Auto-generated ephemeral block'); + expect(items[0].tooltip).to.not.include('Source agent block'); + }); + + test('Should include agent source block ID in tooltip when present', () => { + const cell = createMockCell({ + metadata: { + is_ephemeral: true, + agent_source_block_id: 'a0000000000000000000000000000004' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[0].tooltip).to.include('a0000000000000000000000000000004'); + expect(items[0].tooltip).to.include('Source agent block'); + }); + }); + + suite('Coexistence with other cell types', () => { + test('Should return items for ephemeral cells regardless of pocket type', () => { + const pocketTypes = ['agent', 'code', 'markdown'] as const; + + for (const type of pocketTypes) { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type }, + is_ephemeral: true, + ...(type === 'agent' ? { agent_source_block_id: 'source-id' } : {}) + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.not.be.undefined; + } + }); + }); +}); diff --git a/src/notebooks/deepnote/runtimeCore.web.ts b/src/notebooks/deepnote/runtimeCore.web.ts new file mode 100644 index 0000000000..01cc943a8f --- /dev/null +++ b/src/notebooks/deepnote/runtimeCore.web.ts @@ -0,0 +1,7 @@ +/** Web stands in for @deepnote/runtime-core, which needs Node built-ins. Agent blocks run on desktop only. */ +const unsupported = (): never => { + throw new Error('Deepnote agent blocks are not supported in the web extension host.'); +}; + +export const executeAgentBlock = unsupported; +export const serializeNotebookContextFromBlocks = unsupported; diff --git a/src/notebooks/deepnote/snapshots/snapshotService.ts b/src/notebooks/deepnote/snapshots/snapshotService.ts index d39603d6b7..ce64336d8d 100644 --- a/src/notebooks/deepnote/snapshots/snapshotService.ts +++ b/src/notebooks/deepnote/snapshots/snapshotService.ts @@ -118,6 +118,7 @@ function generateTimestamp(): string { @injectable() export class SnapshotService implements ISnapshotMetadataService, IExtensionSyncActivationService { private readonly converter = new DeepnoteDataConverter(); + private readonly endedExecutionSessions = new Set(); private readonly environmentStates = new Map(); private readonly fileWrittenCallbacks: ((uri: Uri) => void)[] = []; private readonly pendingSnapshotSaves = new Map< @@ -193,11 +194,21 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync this, this.disposables ); + + notebookCellExecutions.onDidStartQueueExecution( + (e) => this.retireFinishedExecutionSession(e.notebookUri), + this, + this.disposables + ); } async captureEnvironmentBeforeExecution(notebookUri: string): Promise { logger.info(`[Snapshot] captureEnvironmentBeforeExecution called for ${notebookUri}`); + // Covers queues opened outside a controller run — the interactive window, execution resumed + // after a reload — which never signal a run start. + this.retireFinishedExecutionSession(notebookUri); + // Seed the session start at capture time so `startedAt` reflects capture, not the first cell. this.tracker.ensureExecutionState(notebookUri, Date.now()); @@ -216,6 +227,7 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync } clearExecutionState(notebookUri: string): void { + this.endedExecutionSessions.delete(notebookUri); this.tracker.clear(notebookUri); this.environmentStates.delete(notebookUri); @@ -701,6 +713,10 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync private async onExecutionComplete(notebookUri: string): Promise { logger.debug(`[Snapshot] onExecutionComplete called for ${notebookUri}`); + // The run is over, but its metadata stays readable — the deferred save and any file save that + // follows still serialize it. The next run's first queue resets it. + this.endedExecutionSessions.add(notebookUri); + // Wait for any pending cell state change events to be processed. // This is needed because the queue completion event can fire before the // last cell's Idle state change event has been processed (race condition). @@ -836,8 +852,15 @@ export class SnapshotService implements ISnapshotMetadataService, IExtensionSync } catch (error) { // Fire-and-forget save: swallow so a failure never becomes an unhandled rejection. logger.error(`[Snapshot] Failed to save deferred snapshot for ${notebookUri}`, error); - } finally { - // Clear execution state so the next run starts fresh, even if the save above failed. + } + } + + /** + * Drops the previous run's metadata once a new run begins. Deferred until then so the save that + * follows a run — and any file save after it — still serializes what that run did. + */ + private retireFinishedExecutionSession(notebookUri: string): void { + if (this.endedExecutionSessions.delete(notebookUri)) { this.clearExecutionState(notebookUri); } } diff --git a/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts b/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts index b817d8b229..0a042032bc 100644 --- a/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts +++ b/src/notebooks/deepnote/snapshots/snapshotService.unit.test.ts @@ -1148,6 +1148,7 @@ project: suite('deferred snapshot save timing', () => { const notebookUri = activatedServiceNotebookUri; let clock: fakeTimers.InstalledClock; + let activatedService: SnapshotService; let changeEmitter: EventEmitter; let closeEmitter: EventEmitter; let flush: sinon.SinonStub; @@ -1158,6 +1159,7 @@ project: clock = fakeTimers.install(); const built = buildActivatedSnapshotService(); + activatedService = built.service; changeEmitter = built.changeEmitter; closeEmitter = built.closeEmitter; @@ -1278,6 +1280,90 @@ project: await clock.tickAsync(3000); assert.isFalse(flush.called, 'an output change with no pending save must not arm a deferred save'); }); + + test('keeps the run metadata readable after the deferred save flushes (catches wiping the state the next file save serializes)', async () => { + await arm(); + await clock.tickAsync(150); + + assert.isTrue(flush.calledOnce, 'the deferred save must have flushed'); + // deepnoteSerializer reads this on every save; wiping it at flush time meant a Ctrl+S a + // moment later wrote the .deepnote file with no execution metadata at all. + assert.isDefined(activatedService.getExecutionMetadata(notebookUri)); + }); + + test('a run that opens no kernel queue still clears the finished run (an agent run generating no cells)', async () => { + await arm(); + await clock.tickAsync(150); + + // Nothing calls captureEnvironmentBeforeExecution here: an agent cell runs off the kernel, + // so without the run-start signal the finished run's counters would be what gets saved. + notebookCellExecutions.notifyQueueStart(notebookUri); + + assert.isUndefined( + activatedService.getExecutionMetadata(notebookUri), + "a run that executes nothing must not report the previous run's counters" + ); + }); + + test('a run start with no completion before it keeps the run alive', async () => { + notebookCellExecutions.notifyQueueStart(notebookUri); + + assert.strictEqual( + activatedService.getExecutionMetadata(notebookUri)?.summary?.blocksExecuted, + 1, + 'only a finished run may be retired' + ); + }); + + test("the next run's first queue clears the finished run (catches counters accumulating across runs)", async () => { + await arm(); + await clock.tickAsync(150); + + // A queue opening after the completion belongs to the next run. + await activatedService.captureEnvironmentBeforeExecution(notebookUri); + + assert.isUndefined( + activatedService.getExecutionMetadata(notebookUri), + "a new run must not inherit the previous run's counters" + ); + }); + + test('a second queue inside the same run keeps the run alive (catches clearing per queue, which an agent run opens one of per generated cell)', async () => { + // No completion since the fixture recorded its executed cell: still the same run. + await activatedService.captureEnvironmentBeforeExecution(notebookUri); + + assert.strictEqual( + activatedService.getExecutionMetadata(notebookUri)?.summary?.blocksExecuted, + 1, + 'a mid-run queue must not reset the session' + ); + }); + + test("the run's own captured environment survives its first executing cell (catches clearing on Executing, which lands after capture)", async () => { + const capturedEnvironment: Environment = { + hash: 'sha256:abc', + packages: {}, + platform: 'linux-x64', + python: { environment: 'venv', version: '3.12.0' } + }; + when(mockEnvironmentCapture.captureEnvironment(anything())).thenResolve(capturedEnvironment); + + await arm(); + await clock.tickAsync(150); + + await activatedService.captureEnvironmentBeforeExecution(notebookUri); + + const cellNotebook = mock(); + when(cellNotebook.uri).thenReturn(Uri.parse(notebookUri)); + + const cell = mock(); + when(cell.notebook).thenReturn(instance(cellNotebook)); + when(cell.metadata).thenReturn({ id: 'cell-1' }); + + notebookCellExecutions.changeCellState(instance(cell), NotebookCellExecutionState.Executing); + + assert.deepStrictEqual(await activatedService.getEnvironmentMetadata(notebookUri), capturedEnvironment); + }); }); suite('createSnapshot', () => { @@ -1848,6 +1934,64 @@ project: ); }); + // An agent block is a Code-kind cell (agentBlockConverter) that runs off the kernel, so the + // tracker never sees it — this is the shape executeAgentCell left behind before it started + // reporting to the execution-state shim. + test('an untracked Code cell (an agent block) blocks Run-All even though every kernel cell executed', async () => { + const mockConfig = mock(); + when(mockConfig.get('snapshots.enabled', true)).thenReturn(true); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + + const projectId = 'test-project-id'; + const notebookId = 'test-notebook-id'; + + const mockNotebook = mockNotebookDoc({ + uri: Uri.parse(notebookUri), + projectId, + notebookId, + cells: [mockCell({ id: 'cell-1', source: 'print(1)' }), mockCell({ id: 'agent-cell' })] + }); + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([mockNotebook]); + + const originalProject: DeepnoteFile = { + metadata: { createdAt: '2025-01-01T00:00:00Z' }, + version: '1.0.0', + project: { + id: projectId, + name: 'Test Project', + notebooks: [{ id: notebookId, name: 'Test Notebook', blocks: [] }] + } + }; + const mockNotebookManager = mock(); + when(mockNotebookManager.getProjectForNotebook(projectId, notebookId)).thenReturn(originalProject); + + const testService = new SnapshotService( + instance(mockEnvironmentCapture), + mockDisposables, + instance(mockNotebookManager), + tracker + ); + + // Only cell-1 is tracked as executed — agent-cell never reaches recordCellExecutionStart/End. + const startTime = Date.now(); + tracker.recordCellExecutionStart(notebookUri, 'cell-1', startTime); + tracker.recordCellExecutionEnd(notebookUri, 'cell-1', startTime + 100, true); + + const writtenUris = captureSnapshotWrites(); + + testService.activate(); + await flushDeferredSave(notebookUri); + + assert.isFalse( + wroteTimestampedSnapshot(writtenUris), + 'a Code-kind cell the tracker never saw must keep the run off the Run-All branch' + ); + assert.isTrue( + wroteLatestSnapshot(writtenUris), + 'the run instead falls back to the partial-run (latest-only) path' + ); + }); + test('writes the snapshot next to the saved notebook, not a sibling that shares the project id', async () => { const mockConfig = mock(); when(mockConfig.get('snapshots.enabled', true)).thenReturn(true); diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index b429263154..6071f2f1c3 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -95,7 +95,10 @@ import { DeepnoteNotebookEnvironmentMapper } from '../kernels/deepnote/environme import { DeepnoteCellExecutionAnalytics } from './deepnote/deepnoteCellExecutionAnalytics'; import { DeepnoteNotebookCommandListener } from './deepnote/deepnoteNotebookCommandListener'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; +import { AgentCellStatusBarProvider } from './deepnote/agentCellStatusBarProvider'; +import { AgentOpenAiApiKeyCommandHandler } from './deepnote/agentOpenAiApiKeyCommandHandler'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; +import { EphemeralCellStatusBarProvider } from './deepnote/ephemeralCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIntegrationStartupCodeProvider'; import { DeepnoteCellCopyHandler } from './deepnote/deepnoteCellCopyHandler'; @@ -266,6 +269,18 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentOpenAiApiKeyCommandHandler + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellStatusBarProvider + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNewCellLanguageService diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index 28937d6b60..fe6a8dbe90 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 { AgentOpenAiApiKeyCommandHandler } from './deepnote/agentOpenAiApiKeyCommandHandler'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; +import { EphemeralCellStatusBarProvider } from './deepnote/ephemeralCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler'; @@ -127,6 +130,18 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentOpenAiApiKeyCommandHandler + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellStatusBarProvider + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNewCellLanguageService diff --git a/src/platform/analytics/types.ts b/src/platform/analytics/types.ts index 4ce4f15a39..8428043247 100644 --- a/src/platform/analytics/types.ts +++ b/src/platform/analytics/types.ts @@ -33,7 +33,7 @@ export type CommandOutcome = 'completed' | 'cancelled' | 'failed'; /** Caller-supplied properties per event; `undefined` means none beyond the common properties the service attaches. */ export interface TelemetryEventProperties { - add_block: { blockType: string }; + add_block: { blockType: string; isEphemeral: boolean }; authenticate_integration: { integrationType: string; outcome: CommandOutcome }; configure_integration: { integrationType: string }; copy_notebook_details: undefined; @@ -44,7 +44,7 @@ export interface TelemetryEventProperties { delete_integration: { integrationType: string }; delete_notebook: { outcome: CommandOutcome }; duplicate_notebook: { outcome: CommandOutcome }; - execute_cell: { cellType: 'sql' | 'markdown' | 'code'; integrationType?: string }; + execute_cell: { cellType: 'sql' | 'markdown' | 'code'; isEphemeral: boolean; integrationType?: string }; execute_notebook: undefined; export_notebook: { outcome: CommandOutcome; format?: string }; import_notebook: { outcome: CommandOutcome; source: 'deepnote' | 'jupyter' }; diff --git a/src/platform/common/constants.ts b/src/platform/common/constants.ts index f78efcb475..f840ff1aa5 100644 --- a/src/platform/common/constants.ts +++ b/src/platform/common/constants.ts @@ -227,6 +227,7 @@ export namespace Commands { export const DisableSnapshots = 'deepnote.disableSnapshots'; export const AuthenticateIntegration = 'deepnote.authenticateIntegration'; export const ManageIntegrations = 'deepnote.manageIntegrations'; + export const AddAgentBlock = 'deepnote.addAgentBlock'; export const AddSqlBlock = 'deepnote.addSqlBlock'; export const AddBigNumberChartBlock = 'deepnote.addBigNumberChartBlock'; export const AddChartBlock = 'deepnote.addChartBlock'; diff --git a/src/platform/deepnote/pocket.ts b/src/platform/deepnote/pocket.ts index 1bf7b6286c..fd2d4bbeb7 100644 --- a/src/platform/deepnote/pocket.ts +++ b/src/platform/deepnote/pocket.ts @@ -2,7 +2,7 @@ import type { DeepnoteBlock, ExecutableBlock } from '@deepnote/blocks'; import { isExecutableBlockType } from '@deepnote/blocks'; import { NotebookCellKind, type NotebookCellData } from 'vscode'; -import { generateBlockId, generateSortingKey } from '../../notebooks/deepnote/dataConversionUtils'; +import { generateBlockId, generateSortingKey, getBlockId } from '../../notebooks/deepnote/dataConversionUtils'; import { logger } from '../logging'; import { generateUuid } from '../common/uuid'; @@ -74,9 +74,8 @@ export function createBlockFromPocket(cell: NotebookCellData, index: number): De const pocket = extractPocketFromCellMetadata(cell); const metadata = cell.metadata ? { ...cell.metadata } : undefined; - // Get id from top-level metadata before cleaning it up - // Check both 'id' and backup '__deepnoteBlockId' in case VS Code modifies 'id' - const cellId = (metadata?.__deepnoteBlockId as string | undefined) || (metadata?.id as string | undefined); + // Read the id before the copy below is stripped of it + const cellId = getBlockId(cell); logger.debug( `[Pocket] createBlockFromPocket index=${index}: cell.metadata.id=${metadata?.id}, __deepnoteBlockId=${metadata?.__deepnoteBlockId}, using cellId=${cellId}, metadata keys=${ diff --git a/src/platform/deepnote/pocket.unit.test.ts b/src/platform/deepnote/pocket.unit.test.ts index b8cff6033f..7671412e38 100644 --- a/src/platform/deepnote/pocket.unit.test.ts +++ b/src/platform/deepnote/pocket.unit.test.ts @@ -121,7 +121,8 @@ suite('Pocket', () => { sortingKey: 'a0', executionCount: 5 }, - id: 'block-123', + __deepnoteBlockId: 'block-123', + id: 'rewritten-by-vscode', custom: 'value' }; diff --git a/src/platform/notebooks/cellExecutionStateService.ts b/src/platform/notebooks/cellExecutionStateService.ts index 25b25dd376..c0800b72e2 100644 --- a/src/platform/notebooks/cellExecutionStateService.ts +++ b/src/platform/notebooks/cellExecutionStateService.ts @@ -57,6 +57,7 @@ const STATE_NAMES: Record = { export namespace notebookCellExecutions { const eventEmitter = trackDisposable(new EventEmitter()); const queueCompletionEmitter = trackDisposable(new EventEmitter()); + const queueStartEmitter = trackDisposable(new EventEmitter()); /** * An {@link Event} which fires when the execution state of a cell has changed. @@ -71,6 +72,13 @@ export namespace notebookCellExecutions { */ export const onDidCompleteQueueExecution = queueCompletionEmitter.event; + /** + * An {@link Event} which fires when a user-initiated run begins, before any cell executes. + * A run that executes no cells at all still fires it, which is what separates it from the + * first cell going Executing. + */ + export const onDidStartQueueExecution = queueStartEmitter.event; + /** * Notify listeners that a notebook's cell execution queue has completed. * @param notebookUri The URI of the notebook whose queue completed @@ -80,6 +88,15 @@ export namespace notebookCellExecutions { queueCompletionEmitter.fire({ notebookUri }); } + /** + * Notify listeners that a user-initiated run is starting. + * @param notebookUri The URI of the notebook whose run is starting + */ + export function notifyQueueStart(notebookUri: string) { + logger.debug(`[CellExecState] Queue execution starting for ${notebookUri}`); + queueStartEmitter.fire({ notebookUri }); + } + export function changeCellState(cell: NotebookCell, state: NotebookCellExecutionState, executionOrder?: number) { const cellId = cell.metadata?.id as string | undefined; const stateName = STATE_NAMES[state] || String(state); diff --git a/src/test/datascience/editor-integration/helpers.ts b/src/test/datascience/editor-integration/helpers.ts index bfd90523ad..2a01e61f2d 100644 --- a/src/test/datascience/editor-integration/helpers.ts +++ b/src/test/datascience/editor-integration/helpers.ts @@ -191,6 +191,7 @@ export function createMockedNotebookDocument( when(cell.document).thenReturn(mockedDocument); when(cell.index).thenReturn(index); when(cell.kind).thenReturn(data.kind); + when(cell.metadata).thenReturn(data.metadata ?? {}); const cellOutput: NotebookCellOutput[] = []; when(cell.outputs).thenReturn(cellOutput); when(cell.notebook).thenReturn(instance(notebook)); diff --git a/src/test/mocks/deepnoteRuntimeCore.ts b/src/test/mocks/deepnoteRuntimeCore.ts index 1032f55c9e..4478994237 100644 --- a/src/test/mocks/deepnoteRuntimeCore.ts +++ b/src/test/mocks/deepnoteRuntimeCore.ts @@ -1,17 +1,16 @@ -import type { ServerInfo, ServerOptions } from '@deepnote/runtime-core'; +import type { AgentBlock } from '@deepnote/blocks'; +import type { AgentBlockContext, ServerInfo, ServerOptions } from '@deepnote/runtime-core'; import type { ChildProcess } from 'child_process'; /** - * Mock of @deepnote/runtime-core for unit tests: the real startServer/stopServer spawn and - * kill Python processes, so this records calls and returns fake server info instead. - * - * build/mocha-esm-loader.js resolves the '@deepnote/runtime-core' specifier to this module, - * so code under test and tests importing the __ helpers below share one module instance. - * The exports are typed against the real package so the mock cannot drift from its API. + * Mock @deepnote/runtime-core: no Python spawns or live agent API; records calls, returns stubs. + * build/mocha-esm-loader.js aliases the package here for one shared instance; exports match real types. */ type RuntimeCore = typeof import('@deepnote/runtime-core'); +const executeAgentBlockCalls: { block: AgentBlock; context: AgentBlockContext }[] = []; +const serializeNotebookContextFromBlocksCalls: { blockCount: number; notebookName: string }[] = []; const startServerCalls: ServerOptions[] = []; const stopServerCalls: ServerInfo[] = []; let nextServerId = 0; @@ -27,6 +26,21 @@ function makeFakeProcess(id: number): ChildProcess { } as unknown as ChildProcess; } +export const executeAgentBlock: RuntimeCore['executeAgentBlock'] = async (block, context) => { + executeAgentBlockCalls.push({ block, context }); + + return { finalOutput: '' }; +}; + +export const serializeNotebookContextFromBlocks: RuntimeCore['serializeNotebookContextFromBlocks'] = ({ + blocks, + notebookName +}) => { + serializeNotebookContextFromBlocksCalls.push({ blockCount: blocks.length, notebookName }); + + return `notebook:${notebookName} blocks:${blocks.length}`; +}; + /** The real one probes the filesystem; unit tests have no venv on disk, so it always fails there. */ export const resolvePythonExecutable: RuntimeCore['resolvePythonExecutable'] = async (pythonPath) => { throw new Error(`No Python executable found under ${pythonPath}`); @@ -54,6 +68,14 @@ export const stopServer: RuntimeCore['stopServer'] = async (info) => { }; // Test-only helpers (prefixed with __ to signal they are not part of the real API). +export function __getExecuteAgentBlockCalls(): { block: AgentBlock; context: AgentBlockContext }[] { + return executeAgentBlockCalls; +} + +export function __getSerializeNotebookContextFromBlocksCalls(): { blockCount: number; notebookName: string }[] { + return serializeNotebookContextFromBlocksCalls; +} + export function __getStartServerCalls(): ServerOptions[] { return startServerCalls; } @@ -67,6 +89,8 @@ export function __setStartServerImpl(impl: RuntimeCore['startServer'] | null): v } export function __resetRuntimeCoreMock(): void { + executeAgentBlockCalls.length = 0; + serializeNotebookContextFromBlocksCalls.length = 0; startServerCalls.length = 0; stopServerCalls.length = 0; nextServerId = 0; diff --git a/test/e2e/.mocharc.js b/test/e2e/.mocharc.js index 14040de246..caa8289f15 100644 --- a/test/e2e/.mocharc.js +++ b/test/e2e/.mocharc.js @@ -3,12 +3,13 @@ // tests are the real guard rails; this is a generous suite-level safety net. const path = require('path'); +// ExTester ignores Mocha `require`; wire rootHooks from compiled output. +const { mochaHooks } = require(path.resolve(__dirname, '..', '..', 'out', 'e2e', 'rootHooks.js')); + module.exports = { timeout: 1500000, // 25 min — env creation + first kernel start (venv + toolkit) can be slow retries: 1, // absorb transient UI flakiness with a single retry reporter: 'spec', color: true, - // Dismiss notification toasts between tests (rootHooks) so they don't accumulate across the one - // shared VS Code instance. Points at compiled output, so compile-e2e must run first. - require: [path.resolve(__dirname, '..', '..', 'out', 'e2e', 'rootHooks.js')] + rootHooks: mochaHooks }; diff --git a/test/e2e/fixtures/agent-block-batch.deepnote b/test/e2e/fixtures/agent-block-batch.deepnote new file mode 100644 index 0000000000..5d426fdc42 --- /dev/null +++ b/test/e2e/fixtures/agent-block-batch.deepnote @@ -0,0 +1,36 @@ +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00.000Z' + modifiedAt: '2025-01-01T00:00:00.000Z' +project: + id: e2e-agent-batch-project + name: E2E Agent Batch + notebooks: + - id: e2e-agent-batch-notebook + name: Agent Batch + blocks: + - id: e2e-agent-batch-failing-block + blockGroup: e2e-agent-batch-group + type: code + content: |- + raise ValueError("e2e-batch-boom") + sortingKey: a0 + metadata: {} + - id: e2e-agent-batch-agent-block + blockGroup: e2e-agent-batch-group + type: agent + content: |- + Run some Python, then add a markdown block summarizing this notebook. + sortingKey: a1 + metadata: + deepnote_agent_model: 'gpt-5.6-sol' + - id: e2e-agent-batch-trailing-block + blockGroup: e2e-agent-batch-group + type: code + content: |- + print("e2e-batch-trailing-ran") + sortingKey: a2 + metadata: {} + executionMode: block + isModule: false + settings: {} diff --git a/test/e2e/fixtures/agent-block-stop.deepnote b/test/e2e/fixtures/agent-block-stop.deepnote new file mode 100644 index 0000000000..a8a2177794 --- /dev/null +++ b/test/e2e/fixtures/agent-block-stop.deepnote @@ -0,0 +1,29 @@ +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00.000Z' + modifiedAt: '2025-01-01T00:00:00.000Z' +project: + id: e2e-agent-stop-project + name: E2E Agent Stop + notebooks: + - id: e2e-agent-stop-notebook + name: Agent Stop + blocks: + - id: e2e-agent-stop-agent-block + blockGroup: e2e-agent-stop-group + type: agent + content: |- + Run some Python, then add a markdown block summarizing this notebook. + sortingKey: a0 + metadata: + deepnote_agent_model: 'gpt-5.6-sol' + - id: e2e-agent-stop-trailing-block + blockGroup: e2e-agent-stop-group + type: code + content: |- + print("e2e-stop-trailing-ran") + sortingKey: a1 + metadata: {} + executionMode: block + isModule: false + settings: {} diff --git a/test/e2e/fixtures/agent-block.deepnote b/test/e2e/fixtures/agent-block.deepnote new file mode 100644 index 0000000000..c5fed3056c --- /dev/null +++ b/test/e2e/fixtures/agent-block.deepnote @@ -0,0 +1,22 @@ +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00.000Z' + modifiedAt: '2025-01-01T00:00:00.000Z' +project: + id: e2e-agent-block-project + name: E2E Agent Block + notebooks: + - id: e2e-agent-block-notebook + name: Agent Block + blocks: + - id: e2e-agent-block + blockGroup: e2e-agent-group + type: agent + content: |- + Run some Python, then add a markdown block summarising this notebook. + sortingKey: a0 + metadata: + deepnote_agent_model: 'gpt-5.6-sol' + executionMode: block + isModule: false + settings: {} diff --git a/test/e2e/helpers/index.ts b/test/e2e/helpers/index.ts index 3ed57bb00d..45a939edc4 100644 --- a/test/e2e/helpers/index.ts +++ b/test/e2e/helpers/index.ts @@ -4,6 +4,7 @@ export * from './constants'; export * from './deepnoteEnvironment'; export * from './deepnoteTree'; export * from './fixtures'; +export * from './mockOpenAiServer'; export * from './modals'; export * from './notebook'; export * from './notifications'; diff --git a/test/e2e/helpers/mockOpenAiServer.ts b/test/e2e/helpers/mockOpenAiServer.ts new file mode 100644 index 0000000000..f9fb50cf6e --- /dev/null +++ b/test/e2e/helpers/mockOpenAiServer.ts @@ -0,0 +1,219 @@ +import { spawn } from 'child_process'; +import * as fs from 'fs'; +import { connect } from 'net'; +import * as os from 'os'; +import * as path from 'path'; +import { setTimeout as delay } from 'timers/promises'; +import { InputBox, Workbench } from 'vscode-extension-tester'; + +import { QUICK_PICK_TIMEOUT } from './constants'; +import { waitForNotification } from './notifications'; + +// npx aimock — keep jest/vitest peers out of the lockfile. +const AIMOCK_VERSION = '1.37.4'; +const AIMOCK_BIN = 'llmock'; + +// Fixed port below ephemeral range (connect pre-flight). +const MOCK_OPENAI_PORT = 18_937; + +/** Set OPENAI_BASE_URL at module scope — ExTester spawns the host before `before` hooks. */ +export function pointExtensionHostAtMockServer(): void { + process.env.OPENAI_BASE_URL = `http://127.0.0.1:${MOCK_OPENAI_PORT}/v1`; +} + +const MOCK_API_KEY = 'sk-e2e-mock-key'; +const SET_API_KEY_COMMAND = 'Deepnote: Set OpenAI API Key'; +const CLEAR_API_KEY_COMMAND = 'Deepnote: Clear OpenAI API Key'; +const API_KEY_SAVED_NOTIFICATION = /OpenAI API key has been saved/; + +/** + * Stores a throwaway key so the agent has credentials to send. The request goes to the mock, which + * never checks it — `pointExtensionHostAtMockServer` is what keeps it off the real API. + */ +export async function storeMockOpenAiApiKey(): Promise { + await new Workbench().executeCommand(SET_API_KEY_COMMAND); + + const input = await InputBox.create(QUICK_PICK_TIMEOUT); + await input.setText(MOCK_API_KEY); + await input.confirm(); + + await waitForNotification(API_KEY_SAVED_NOTIFICATION, QUICK_PICK_TIMEOUT, true); +} + +/** Removes the stored key. The key outlives a suite, so every suite that stores one clears it. */ +export async function clearStoredOpenAiApiKey(): Promise { + await new Workbench().executeCommand(CLEAR_API_KEY_COMMAND); +} + +const START_TIMEOUT = 90_000; +const POLL_INTERVAL = 200; +const STOP_TIMEOUT = 2_000; + +export interface MockOpenAiServer { + stop: () => Promise; +} + +export interface MockToolCall { + arguments: string; + id: string; + name: string; +} + +/** Per-leg match predicate (not call order); Mocha-retry-safe. */ +export type MockAgentMatch = { hasToolResult: false } | { toolResultContains: string }; + +export type MockAgentResponse = { content: string } | { toolCall: MockToolCall }; + +export interface MockAgentTurn { + match: MockAgentMatch; + response: MockAgentResponse; +} + +function canConnect(port: number): Promise { + return new Promise((resolve) => { + const socket = connect({ host: '127.0.0.1', port }); + const settle = (reachable: boolean) => { + socket.destroy(); + resolve(reachable); + }; + + socket.once('connect', () => settle(true)); + socket.once('error', () => settle(false)); + }); +} + +function assertBaseUrlPointsAtMock(): void { + if (!process.env.OPENAI_BASE_URL?.includes(`:${MOCK_OPENAI_PORT}`)) { + throw new Error( + `OPENAI_BASE_URL must point at 127.0.0.1:${MOCK_OPENAI_PORT}; run via "npm run test:e2e". ` + + `Without it the agent would call the real OpenAI API. Current value: ` + + `${JSON.stringify(process.env.OPENAI_BASE_URL)}` + ); + } +} + +function writeFixtures(turns: MockAgentTurn[]): string { + const fixtures = turns.map(({ match, response }) => ({ + match, + response: 'toolCall' in response ? { toolCalls: [response.toolCall] } : { content: response.content } + })); + + const directory = fs.mkdtempSync(path.join(os.tmpdir(), 'deepnote-e2e-aimock-')); + fs.writeFileSync(path.join(directory, 'fixtures.json'), JSON.stringify({ fixtures }, undefined, 4)); + + return directory; +} + +export async function startMockOpenAiServer(turns: MockAgentTurn[]): Promise { + assertBaseUrlPointsAtMock(); + + if (await canConnect(MOCK_OPENAI_PORT)) { + throw new Error( + `Port ${MOCK_OPENAI_PORT} is already in use — most likely a mock server left behind by an ` + + `interrupted run. Kill it before running the suite.` + ); + } + + const fixturesDirectory = writeFixtures(turns); + + const child = spawn( + 'npx', + [ + '--prefer-offline', + '-y', + '-p', + `@copilotkit/aimock@${AIMOCK_VERSION}`, + AIMOCK_BIN, + '-f', + fixturesDirectory, + '-p', + String(MOCK_OPENAI_PORT), + '--strict', + '--log-level', + 'warn' + ], + { + detached: true, + stdio: ['ignore', 'inherit', 'inherit'] + } + ); + + let exitReason: string | undefined; + child.once('exit', (code, signal) => { + exitReason = `code ${code}, signal ${signal}`; + }); + child.once('error', (error) => { + exitReason = `spawn failed: ${error.message}`; + }); + + const signalTree = (signal: NodeJS.Signals) => { + try { + if (child.pid === undefined) { + return; + } + + process.kill(-child.pid, signal); + } catch { + // process already exited + } + }; + + const killChild = () => signalTree('SIGKILL'); + process.once('exit', killChild); + + const hasShutDown = async () => + (child.exitCode !== null || child.signalCode !== null) && !(await canConnect(MOCK_OPENAI_PORT)); + + const waitForShutdown = async (): Promise => { + const deadline = Date.now() + STOP_TIMEOUT; + + while (Date.now() < deadline) { + if (await hasShutDown()) { + return true; + } + + await delay(POLL_INTERVAL); + } + + return false; + }; + + const stop = async () => { + fs.rmSync(fixturesDirectory, { force: true, recursive: true }); + + signalTree('SIGTERM'); + await waitForShutdown(); + signalTree('SIGKILL'); + + if (!(await waitForShutdown())) { + throw new Error( + `aimock did not shut down after SIGKILL: port ${MOCK_OPENAI_PORT} still accepts ` + + `connections, or the npx process has not exited.` + ); + } + + process.removeListener('exit', killChild); + }; + + const deadline = Date.now() + START_TIMEOUT; + while (Date.now() < deadline) { + if (exitReason) { + await stop(); + + throw new Error(`aimock exited before it started listening (${exitReason}); see its output above`); + } + + if (await canConnect(MOCK_OPENAI_PORT)) { + return { stop }; + } + + await delay(POLL_INTERVAL); + } + + await stop(); + + throw new Error( + `aimock did not listen on port ${MOCK_OPENAI_PORT} within ${START_TIMEOUT}ms. A stale server from an ` + + `earlier run may still hold the port.` + ); +} diff --git a/test/e2e/helpers/notebook.ts b/test/e2e/helpers/notebook.ts index 6fc561a7b3..36784deb5c 100644 --- a/test/e2e/helpers/notebook.ts +++ b/test/e2e/helpers/notebook.ts @@ -4,12 +4,12 @@ import { OUTPUT_FRAME_SWITCH_TIMEOUT, OUTPUT_POLL_INTERVAL, OUTPUT_SELECTOR, WOR import { dismissAllNotifications } from './notifications'; /** - * Focuses the given notebook editor and clicks its toolbar "Run All" button. The command-palette - * entry for `deepnote.runallcells` ("Jupyter: Run All Cells") is gated behind context keys - * (`deepnote.ispythonornativeactive`, …) that are not reliably set under automation, so driving it + * Focuses the given notebook editor and clicks the toolbar button carrying `ariaLabel`. The + * command-palette entries for these actions are gated behind context keys + * (`deepnote.ispythonornativeactive`, …) that are not reliably set under automation, so driving them * through `Workbench.executeCommand` can silently miss and trigger the wrong command. */ -export async function clickRunAll(notebookFileName: string): Promise { +async function clickNotebookToolbarButton(notebookFileName: string, ariaLabel: string): Promise { const driver = VSBrowser.instance.driver; await new EditorView().openEditor(notebookFileName); @@ -17,12 +17,11 @@ export async function clickRunAll(notebookFileName: string): Promise { // Locate AND click inside the same wait loop. The notebook toolbar can re-render between finding // the button and clicking it (the editor re-focuses, kernel status / notifications change), which // would otherwise surface as a StaleElementReferenceError. Re-finding and clicking on the next - // tick is still a SINGLE "Run All" — the run is only issued once the click actually lands, so - // this does not re-run a notebook whose first execution was accepted. + // tick still issues the action only ONCE — it is only issued when the click actually lands. await driver.wait( async () => { try { - const [button] = await driver.findElements(By.css('a.action-label[aria-label="Run All"]')); + const [button] = await driver.findElements(By.css(`a.action-label[aria-label="${ariaLabel}"]`)); if (!button) { return false; } @@ -31,38 +30,114 @@ export async function clickRunAll(notebookFileName: string): Promise { return true; } catch (error) { - console.warn('[deepnote-e2e] locate/click notebook Run All (retrying):', error); + console.warn(`[deepnote-e2e] locate/click notebook "${ariaLabel}" (retrying):`, error); return false; } }, WORKBENCH_TIMEOUT, - 'notebook "Run All" button did not appear or could not be clicked' + `notebook "${ariaLabel}" button did not appear or could not be clicked` ); } +export async function clickRunAll(notebookFileName: string): Promise { + return clickNotebookToolbarButton(notebookFileName, 'Run All'); +} + /** - * Reads the notebook cell output once. + * Clicks the toolbar's "Interrupt" button — VS Code's `notebook.interruptExecution`, shown while + * `notebookHasSomethingRunning && notebookInterruptibleKernel`. * - * Output lives two iframes deep (iframe.webview.ready -> #active-frame). We only attempt to switch - * when an output webview iframe actually exists (`getViewToSwitchTo`), and we read output-specific - * elements inside the frame — so we never match the cell's source code that is visible in the editor - * of the main document. Returns '' when no output is present yet. + * It is the only toolbar action that reaches the controller's `interruptHandler`, and therefore the + * only one that signals a running agent to stop. "Stop Execution" (`notebook.cancelExecution`, + * which VS Code shows in its place for a kernel that declares no interrupt handler) cancels the + * cells without ever telling the agent, so it must not stand in as a fallback here. */ -export async function readRenderedOutput(): Promise { +export async function clickInterrupt(notebookFileName: string): Promise { + return clickNotebookToolbarButton(notebookFileName, 'Interrupt'); +} + +/** + * Clicks the notebook cell status bar item whose text contains `label`. Cell chrome lives in the + * main window DOM (not the output iframe), so this switches out of the webview first and matches on + * `textContent` — Selenium's `getText()` is empty for items scrolled out of view. + */ +export async function clickCellStatusBarItem(label: string): Promise { + const driver = VSBrowser.instance.driver; + + await new WebView().switchBack().catch((error) => { + console.warn('[deepnote-e2e] switch back before clicking a cell status bar item:', error); + }); + + // Locate AND click in the same wait loop: the status bar re-renders as cells execute, which + // would otherwise surface as a StaleElementReferenceError between finding and clicking. + await driver.wait( + async () => { + try { + for (const item of await driver.findElements(By.css('.cell-statusbar-container .cell-status-item'))) { + const text = (await item.getAttribute('textContent')) ?? ''; + if (!text.includes(label)) { + continue; + } + + await driver.executeScript('arguments[0].scrollIntoView({block: "center"})', item); + await item.click(); + + return true; + } + + return false; + } catch (error) { + console.warn('[deepnote-e2e] locate/click cell status bar item (retrying):', error); + + return false; + } + }, + WORKBENCH_TIMEOUT, + `notebook cell status bar item "${label}" did not appear or could not be clicked` + ); +} + +/** Run `read` in the notebook output webview; '' if the frame is missing. */ +async function readInsideNotebookWebview(read: (webView: WebView) => Promise): Promise { + const driver = VSBrowser.instance.driver; const webView = new WebView(); - const outputFrame = await webView.getViewToSwitchTo().catch((error) => { - console.warn('[deepnote-e2e] locate notebook output webview:', error); + const frame = await webView.getViewToSwitchTo().catch((error) => { + console.warn('[deepnote-e2e] locate notebook webview:', error); return undefined; }); - if (!outputFrame) { + if (!frame) { return ''; } - let text = ''; try { await webView.switchToFrame(OUTPUT_FRAME_SWITCH_TIMEOUT); + + if (await driver.executeScript('return window.self === window.top')) { + return ''; + } + + return (await read(webView)).trim(); + } catch (error) { + console.warn('[deepnote-e2e] read inside notebook webview:', error); + + return ''; + } finally { + await webView.switchBack().catch((error) => { + console.warn('[deepnote-e2e] switch back from notebook webview:', error); + }); + } +} + +/** Notebook webview body (markdown previews and outputs). */ +export async function readNotebookWebviewText(): Promise { + return readInsideNotebookWebview(async (webView) => (await webView.findWebElement(By.css('body'))).getText()); +} + +/** Cell output once; falls back to frame body if output selectors miss. */ +export async function readRenderedOutput(): Promise { + return readInsideNotebookWebview(async (webView) => { const elements = await webView.findWebElements(By.css(OUTPUT_SELECTOR)); const texts = await Promise.all( elements.map((element) => @@ -73,36 +148,72 @@ export async function readRenderedOutput(): Promise { }) ) ); - text = texts.join('\n').trim(); - - // Fallback: if the renderer used unexpected classes, read the frame body — safe here because - // we have confirmed we are inside the output iframe, not the editor. - if (!text) { - const body = await webView.findWebElement(By.css('body')).catch((error) => { - console.warn('[deepnote-e2e] read output frame body:', error); - - return undefined; - }); - text = body - ? ( - await body.getText().catch((error) => { - console.warn('[deepnote-e2e] read output frame body text:', error); - - return ''; - }) - ).trim() - : ''; + const text = texts.join('\n').trim(); + + return text || (await webView.findWebElement(By.css('body'))).getText(); + }); +} + +/** + * Polls the notebook webview until every marker in `markers` is rendered and none of `absentMarkers` + * is, then returns the text it settled on. `context` names the state being waited for; it is only + * used to make the timeout message say what did not happen. + */ +export async function awaitWebviewMarkers( + markers: string[], + timeout: number, + context: string, + absentMarkers: string[] = [] +): Promise { + const driver = VSBrowser.instance.driver; + const deadline = Date.now() + timeout; + let text = ''; + + while (Date.now() < deadline) { + text = await readNotebookWebviewText(); + const missing = markers.filter((marker) => !text.includes(marker)); + const lingering = absentMarkers.filter((marker) => text.includes(marker)); + if (missing.length === 0 && lingering.length === 0) { + return text; } - } catch (error) { - // Frame went stale or output not painted yet — treat as no output this tick. - console.warn('[deepnote-e2e] read rendered notebook output:', error); - } finally { - await webView.switchBack().catch((error) => { - console.warn('[deepnote-e2e] switch back from notebook output webview:', error); - }); + + await driver.sleep(OUTPUT_POLL_INTERVAL); } - return text; + const missing = markers.filter((marker) => !text.includes(marker)); + const lingering = absentMarkers.filter((marker) => text.includes(marker)); + throw new Error( + `Timed out after ${timeout}ms waiting for notebook webview (${context}). Missing: ${JSON.stringify( + missing + )}. ` + `Lingering: ${JSON.stringify(lingering)}. Last text: ${JSON.stringify(text)}` + ); +} + +/** + * Fails if any of `markers` renders in the notebook webview during the next `windowMs`. + * + * Non-occurrence needs a window rather than one read: the regressions this guards render the + * forbidden text a beat *after* the state the test waited for — a batch that should have stopped + * carries on into the agent's round trip to the local mock and then the trailing cell. Size the + * window well above that round trip, since the whole window is spent on every passing run. + */ +export async function assertMarkersStayAbsent(markers: string[], windowMs: number, context: string): Promise { + const driver = VSBrowser.instance.driver; + const deadline = Date.now() + windowMs; + + while (Date.now() < deadline) { + const text = await readNotebookWebviewText(); + const rendered = markers.filter((marker) => text.includes(marker)); + + if (rendered.length > 0) { + throw new Error( + `Notebook webview rendered ${JSON.stringify(rendered)}, which must not appear (${context}). ` + + `Full text: ${JSON.stringify(text)}` + ); + } + + await driver.sleep(OUTPUT_POLL_INTERVAL); + } } /** diff --git a/test/e2e/helpers/yaml.ts b/test/e2e/helpers/yaml.ts index d179bd8b1f..e09f89038e 100644 --- a/test/e2e/helpers/yaml.ts +++ b/test/e2e/helpers/yaml.ts @@ -1,6 +1,34 @@ -import { deserializeDeepnoteFile } from '@deepnote/blocks'; +import { deserializeDeepnoteFile, isExecutableBlock } from '@deepnote/blocks'; /** Counts notebooks in a serialized `.deepnote` file by parsing it with the canonical schema. */ export function notebookCount(yaml: string): number { return deserializeDeepnoteFile(yaml).project.notebooks.length; } + +/** + * The text a block's stream outputs carry in a serialized `.deepnote`, concatenated in order. + * + * Parse rather than search the raw YAML: `serializeDeepnoteFile` folds at 120 columns, so a marker + * that is one unbroken string in the block can sit across two lines in the file. + */ +export function blockStreamOutputText(yaml: string, blockId: string): string { + const block = deserializeDeepnoteFile(yaml) + .project.notebooks.flatMap((notebook) => notebook.blocks ?? []) + .find((candidate) => candidate.id === blockId); + + if (!block) { + throw new Error(`No block ${JSON.stringify(blockId)} in the serialized project.`); + } + + if (!isExecutableBlock(block)) { + throw new Error(`Block ${JSON.stringify(blockId)} is a ${block.type} block, which carries no outputs.`); + } + + return (block.outputs ?? []) + .map((output: { text?: string | string[] }) => { + const text = output.text ?? ''; + + return Array.isArray(text) ? text.join('') : text; + }) + .join(''); +} diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts new file mode 100644 index 0000000000..87fa34be6e --- /dev/null +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -0,0 +1,668 @@ +/** + * Agent block E2E vs local aimock; legs 2–3 advance on real tool results (no live OpenAI). + * + * Three fixtures share one workspace and one environment — provisioning a second environment costs + * ~90s of CI and every test here wants the same kernel. Only that setup is shared: each group below + * opens and binds the notebook it runs, so the groups are order-independent and either can run on + * its own. + */ + +import * as fs from 'fs'; +import * as path from 'path'; +import { EditorView, VSBrowser, WebView, Workbench } from 'vscode-extension-tester'; + +import { + FIRST_RUN_OUTPUT_TIMEOUT, + MockOpenAiServer, + SUITE_TIMEOUT, + WORKBENCH_TIMEOUT, + assertMarkersStayAbsent, + awaitWebviewMarkers, + blockStreamOutputText, + clearStoredOpenAiApiKey, + clickCellStatusBarItem, + clickInterrupt, + clickRunAll, + confirmModalDialog, + copyFixtureToTempDir, + createEnvironment, + createScreenshotter, + dismissAllNotifications, + openFolderViaDialog, + openWorkspaceFile, + pointExtensionHostAtMockServer, + selectEnvironmentForNotebook, + startMockOpenAiServer, + storeMockOpenAiApiKey +} from '../helpers'; + +pointExtensionHostAtMockServer(); + +const AGENT_FILE = 'agent-block.deepnote'; +const AGENT_BLOCK_ID = 'e2e-agent-block'; +// Mixed-batch fixtures: a failing cell before the agent, and a trailing cell after it. +const BATCH_FILE = 'agent-block-batch.deepnote'; +const STOP_FILE = 'agent-block-stop.deepnote'; +const CODE_TOOL_NAME = 'add_code_block'; +const MARKDOWN_TOOL_NAME = 'add_markdown_block'; +// Leg 3 match: agentCellExecutionHandler add_markdown_block tool result. +const MARKDOWN_BLOCK_ADDED_TEXT = 'Markdown block added.'; +const ENVIRONMENT_NAME = 'E2E Agent Env'; +const AGENT_RUN_TIMEOUT = 60_000; +const PYTHON_OUTPUT_MARKER = 'agent-generated-python-ran'; +const GENERATED_PYTHON = `print("${PYTHON_OUTPUT_MARKER}")`; +const EPHEMERAL_MARKDOWN_TEXT = 'Ephemeral markdown written by the E2E agent run'; +// aimock emits 20-char chunks (multiple text_delta). +const FINAL_AGENT_TEXT = 'Summary added as a markdown block, streamed across several deltas.'; +// Disjoint from first-run markers (assertOccurrences below). +const RERUN_PYTHON_OUTPUT_MARKER = 'rerun-python-ran'; +const RERUN_GENERATED_PYTHON = `print("${RERUN_PYTHON_OUTPUT_MARKER}")`; +const RERUN_MARKDOWN_TEXT = 'Second-run markdown from the E2E agent'; +const RERUN_FINAL_AGENT_TEXT = 'Re-run summary added as a markdown block.'; +// executeAgentCell stale-run error substring. +const STALE_CELLS_ERROR_TEXT = 'from its previous run'; +// Third run, disjoint from both prior runs so the clear test stands on its own. +const CLEAR_RUN_PYTHON_OUTPUT_MARKER = 'clear-run-python-ran'; +const CLEAR_RUN_GENERATED_PYTHON = `print("${CLEAR_RUN_PYTHON_OUTPUT_MARKER}")`; +const CLEAR_RUN_MARKDOWN_TEXT = 'Third-run markdown from the E2E agent'; +const CLEAR_RUN_FINAL_AGENT_TEXT = 'Clear-run summary added as a markdown block.'; +// AgentCellStatusBarProvider button and its confirmation. +const CLEAR_EPHEMERAL_BUTTON = 'Clear ephemeral blocks'; +const CLEAR_EPHEMERAL_CONFIRM = 'Clear'; +const CLEAR_EPHEMERAL_CONFIRM_TEXT = 'ephemeral block'; +const CLEAR_EPHEMERAL_TIMEOUT = 30_000; +const REVERT_FILE_COMMAND = 'File: Revert File'; +const DISCARD_CHANGES_BUTTON = "Don't Save"; +// Fourth run: the generated cell raises, so the agent must be handed the failure and carry on. +const FAILING_PYTHON_MARKER = 'e2e-agent-code-boom'; +const FAILING_GENERATED_PYTHON = `raise ValueError("${FAILING_PYTHON_MARKER}")`; +// addAndExecuteCodeBlock's prefix for a cell that ran and failed (as opposed to one that never ran). +const EXECUTION_FAILED_TEXT = 'Execution failed:'; +const FAILURE_RECOVERY_MARKDOWN = 'Recovered from the failed cell and carried on'; +const FAILURE_FINAL_AGENT_TEXT = 'Reported the failure as a markdown block.'; +// Fifth run: read back from the snapshot sidecar rather than the webview. +const PERSISTED_PYTHON_OUTPUT_MARKER = 'persisted-python-ran'; +const PERSISTED_GENERATED_PYTHON = `print("${PERSISTED_PYTHON_OUTPUT_MARKER}")`; +const PERSISTED_MARKDOWN_TEXT = 'Fifth-run markdown from the E2E agent'; +const PERSISTED_FINAL_AGENT_TEXT = 'Transcript that must survive the save in full.'; +// executeAgentCell's first output item — all that survived the save before the streamed-item fix. +const AGENT_FIRST_OUTPUT_ITEM = '[Agent] Planning next steps...'; +const SNAPSHOT_WRITE_TIMEOUT = 60_000; +const SNAPSHOT_POLL_INTERVAL = 1_500; + +// Mixed batch, failing first cell: the two markers after it must never render. +const BATCH_FAILURE_MARKER = 'e2e-batch-boom'; +const BATCH_TRAILING_MARKER = 'e2e-batch-trailing-ran'; +const BATCH_AGENT_PYTHON_MARKER = 'e2e-batch-agent-python-ran'; +const BATCH_AGENT_MARKDOWN_TEXT = 'Markdown the agent must never get to write'; +const BATCH_AGENT_FINAL_TEXT = 'Summary the agent must never get to write.'; +/** + * How long the forbidden markers must stay away. A batch that carried on renders them within one + * agent round trip to the local mock (~1–2s) plus one cell execution (~1–3s); this keeps a wide + * margin over that without being open-ended, since a passing run spends the whole window. + */ +const BATCH_SETTLE_WINDOW = 10_000; + +// The generated cell prints, then sleeps: a bounded window in which the notebook is demonstrably +// running and Stop has something to interrupt. Long enough that a slow click still lands inside it, +// and never actually waited out on a passing run. +const STOP_SLEEP_MARKER = 'e2e-stop-sleeping'; +const STOP_GENERATED_PYTHON = `print("${STOP_SLEEP_MARKER}", flush=True)\nimport time\ntime.sleep(60)`; +const STOP_TRAILING_MARKER = 'e2e-stop-trailing-ran'; +const STOP_MARKDOWN_TEXT = 'Markdown the stopped agent must never write'; +const STOP_FINAL_TEXT = 'Summary the stopped agent must never write.'; +const AGENT_STOPPED_TEXT = '[Agent] Stopped'; +const STOP_ACKNOWLEDGED_TIMEOUT = 30_000; +// Shorter than the batch window: `[Agent] Stopped` already proves the run ended, so this only has +// to outlast the trailing cell that a batch which ignored the stop would dispatch next. +const STOP_SETTLE_WINDOW = 8_000; + +/** + * Keeps exactly one editor open: `clickRunAll` takes the first toolbar in DOM order. + * + * A run leaves the notebook dirty, so closing it raises the save prompt. Revert first so the close + * is clean, and answer the prompt if one appears anyway: an unanswered modal dims the workbench and + * intercepts every later click, which surfaces as an unrelated "element is visible" timeout in + * whatever runs next rather than here. + */ +async function openOnly(fileName: string): Promise { + await new WebView().switchBack().catch((error) => { + console.warn('[agent-block] switch back from webview before opening an editor:', error); + }); + + const alreadyOpen = await new EditorView().getOpenEditorTitles().catch(() => [] as string[]); + if (alreadyOpen.length > 0) { + await new Workbench().executeCommand(REVERT_FILE_COMMAND).catch((error) => { + console.warn('[agent-block] revert notebook before closing it:', error); + }); + } + + await new EditorView().closeAllEditors().catch((error) => { + console.warn('[agent-block] close editors before opening the next notebook:', error); + }); + + // Only when editors survived the close, since `confirmModalDialog` waits out its full timeout + // and then throws when no dialog is up. + const stillOpen = await new EditorView().getOpenEditorTitles().catch(() => [] as string[]); + if (stillOpen.length > 0) { + await confirmModalDialog(DISCARD_CHANGES_BUTTON).catch((error) => { + console.warn('[agent-block] discard unsaved changes before opening the next notebook:', error); + }); + await new EditorView().closeAllEditors().catch((error) => { + console.warn('[agent-block] close editors after discarding unsaved changes:', error); + }); + } + + await openWorkspaceFile(fileName); + await VSBrowser.instance.driver.wait( + async () => (await new EditorView().getOpenEditorTitles()).some((title) => title.includes(fileName)), + WORKBENCH_TIMEOUT, + `${fileName} did not open` + ); +} + +function assertRenderedContiguously(transcript: string, expected: string): void { + if (transcript.includes(expected)) { + return; + } + + throw new Error( + `Agent transcript does not contain ${JSON.stringify(expected)} as one unbroken run — appended stdout ` + + `items are not rendering as a single block. Full transcript: ${JSON.stringify(transcript)}` + ); +} + +function assertOccurrences(rendered: string, needle: string, expected: number): void { + const actual = rendered.split(needle).length - 1; + + if (actual === expected) { + return; + } + + throw new Error( + `Expected ${expected} occurrence(s) of ${JSON.stringify(needle)} in the notebook, found ${actual}. ` + + `Full text: ${JSON.stringify(rendered)}` + ); +} + +describe('Deepnote — running an agent block against a stand-in OpenAI API', function () { + this.timeout(SUITE_TIMEOUT); + + let cleanupTempDir: (() => void) | undefined; + let mockServer: MockOpenAiServer | undefined; + let screenshot: (label: string) => Promise; + let workspaceDir = ''; + + before(async function () { + screenshot = createScreenshotter(this); + + const copy = copyFixtureToTempDir(AGENT_FILE); + cleanupTempDir = copy.cleanup; + workspaceDir = copy.tempDir; + + for (const fixture of [BATCH_FILE, STOP_FILE]) { + fs.copyFileSync( + path.resolve(process.cwd(), 'test', 'e2e', 'fixtures', fixture), + path.join(copy.tempDir, fixture) + ); + } + + await VSBrowser.instance.waitForWorkbench(WORKBENCH_TIMEOUT); + await openFolderViaDialog(copy.tempDir); + await VSBrowser.instance.waitForWorkbench(WORKBENCH_TIMEOUT); + + // createEnvironment needs an active deepnote notebook; which one does not matter, and each + // group below binds the kernel for the notebook it actually runs. + await openOnly(AGENT_FILE); + await createEnvironment(ENVIRONMENT_NAME); + + await dismissAllNotifications(); + await storeMockOpenAiApiKey(); + await screenshot('environment-created'); + }); + + async function releaseMockServer(): Promise { + await mockServer?.stop().catch((error) => { + console.warn('[agent-block] stop the mock OpenAI server:', error); + }); + mockServer = undefined; + } + + beforeEach(releaseMockServer); + afterEach(releaseMockServer); + + after(async function () { + try { + cleanupTempDir?.(); + } catch (error) { + console.warn('[agent-block] remove temp workspace dir during cleanup:', error); + } + + await new WebView().switchBack().catch((error) => { + console.warn('[agent-block] switch back from webview during cleanup:', error); + }); + await new Workbench().executeCommand(REVERT_FILE_COMMAND).catch((error) => { + console.warn('[agent-block] revert notebook during cleanup:', error); + }); + await new EditorView().closeAllEditors().catch((error) => { + console.warn('[agent-block] close all editors during cleanup:', error); + }); + + const openEditors = await new EditorView().getOpenEditorTitles().catch(() => [] as string[]); + if (openEditors.length > 0) { + await confirmModalDialog(DISCARD_CHANGES_BUTTON).catch((error) => { + console.warn('[agent-block] discard unsaved changes during cleanup:', error); + }); + } + await clearStoredOpenAiApiKey().catch((error) => { + console.warn('[agent-block] clear the stored OpenAI API key during cleanup:', error); + }); + }); + + describe('one agent block on its own', function () { + // Opens and binds the notebook this group runs, so the group does not care what ran before + // it. Closing and reopening drops the block's generated cells, which is why it happens here + // once and never between the tests below. + before(async function () { + await openOnly(AGENT_FILE); + await selectEnvironmentForNotebook(ENVIRONMENT_NAME, AGENT_FILE); + await dismissAllNotifications(); + }); + + it('executes the generated code block, inserts its markdown block, and streams one transcript', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: GENERATED_PYTHON }), + id: 'call_e2e_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: PYTHON_OUTPUT_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: EPHEMERAL_MARKDOWN_TEXT }), + id: 'call_e2e_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: FINAL_AGENT_TEXT } + } + ]); + + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); + + await awaitWebviewMarkers([PYTHON_OUTPUT_MARKER], FIRST_RUN_OUTPUT_TIMEOUT, 'generated code cell stdout'); + + const transcript = await awaitWebviewMarkers( + [ + `[Agent] Tool called: ${CODE_TOOL_NAME}`, + `[Agent] Tool called: ${MARKDOWN_TOOL_NAME}`, + EPHEMERAL_MARKDOWN_TEXT, + FINAL_AGENT_TEXT + ], + AGENT_RUN_TIMEOUT, + 'agent tool loop and ephemeral markdown' + ); + + await screenshot('agent-run'); + + assertRenderedContiguously( + transcript, + `[Agent] Tool called: ${MARKDOWN_TOOL_NAME}\n\n[Agent] Tool output: ${MARKDOWN_TOOL_NAME}` + ); + assertRenderedContiguously(transcript, `[Agent] Text:\n${FINAL_AGENT_TEXT}`); + }); + + // Serial with prior it — block still owns first-run cells. + it('clears the cells its previous run generated instead of stacking a second copy', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: RERUN_GENERATED_PYTHON }), + id: 'call_e2e_rerun_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: RERUN_PYTHON_OUTPUT_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: RERUN_MARKDOWN_TEXT }), + id: 'call_e2e_rerun_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: RERUN_FINAL_AGENT_TEXT } + } + ]); + + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); + + const rendered = await awaitWebviewMarkers( + [RERUN_PYTHON_OUTPUT_MARKER, RERUN_MARKDOWN_TEXT, RERUN_FINAL_AGENT_TEXT], + AGENT_RUN_TIMEOUT, + 'second agent run' + ); + + await screenshot('agent-rerun'); + + // assertOccurrences — retries would duplicate markers. + assertOccurrences(rendered, PYTHON_OUTPUT_MARKER, 0); + assertOccurrences(rendered, EPHEMERAL_MARKDOWN_TEXT, 0); + assertOccurrences(rendered, RERUN_PYTHON_OUTPUT_MARKER, 1); + assertOccurrences(rendered, RERUN_MARKDOWN_TEXT, 1); + assertOccurrences(rendered, STALE_CELLS_ERROR_TEXT, 0); + }); + + // Self-contained: generates the run it clears, so it survives --grep and a Mocha retry (Run All + // drops any stale generated cells first). + it('clears the whole generated run from the agent block status bar button', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: CLEAR_RUN_GENERATED_PYTHON }), + id: 'call_e2e_clear_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: CLEAR_RUN_PYTHON_OUTPUT_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: CLEAR_RUN_MARKDOWN_TEXT }), + id: 'call_e2e_clear_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: CLEAR_RUN_FINAL_AGENT_TEXT } + } + ]); + + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); + + await awaitWebviewMarkers( + [CLEAR_RUN_PYTHON_OUTPUT_MARKER, CLEAR_RUN_MARKDOWN_TEXT, CLEAR_RUN_FINAL_AGENT_TEXT], + AGENT_RUN_TIMEOUT, + 'agent run whose cells the button clears' + ); + + await clickCellStatusBarItem(CLEAR_EPHEMERAL_BUTTON); + await confirmModalDialog(CLEAR_EPHEMERAL_CONFIRM, { messageIncludes: CLEAR_EPHEMERAL_CONFIRM_TEXT }); + + // The button lives on the agent block and takes both cells its run generated. Requiring the + // agent's own transcript to survive keeps an unreadable webview (which reads as '') from + // passing this as "the generated cells are gone". + await awaitWebviewMarkers( + [CLEAR_RUN_FINAL_AGENT_TEXT], + CLEAR_EPHEMERAL_TIMEOUT, + 'ephemeral cells cleared', + [CLEAR_RUN_PYTHON_OUTPUT_MARKER, CLEAR_RUN_MARKDOWN_TEXT] + ); + + await screenshot('agent-ephemeral-cleared'); + }); + + it('hands the agent a failed generated cell and lets the run carry on', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: FAILING_GENERATED_PYTHON }), + id: 'call_e2e_failing_code', + name: CODE_TOOL_NAME + } + } + }, + // Only a tool result carrying "Execution failed:" reaches this leg. The mock runs with + // --strict, so if the handler swallowed the failure and reported success instead, this + // request goes unmatched, the agent errors out, and the markers below never render — + // which is what makes the assertion about the tool result and not just about the cell. + { + match: { toolResultContains: EXECUTION_FAILED_TEXT }, + response: { + toolCall: { + arguments: JSON.stringify({ content: FAILURE_RECOVERY_MARKDOWN }), + id: 'call_e2e_failing_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: FAILURE_FINAL_AGENT_TEXT } + } + ]); + + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); + + await awaitWebviewMarkers( + [FAILING_PYTHON_MARKER, FAILURE_RECOVERY_MARKDOWN, FAILURE_FINAL_AGENT_TEXT], + AGENT_RUN_TIMEOUT, + 'agent run whose generated cell raises' + ); + + await screenshot('agent-generated-cell-failed'); + }); + + it('persists every streamed output item to the snapshot, not just the first', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: PERSISTED_GENERATED_PYTHON }), + id: 'call_e2e_persisted_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: PERSISTED_PYTHON_OUTPUT_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: PERSISTED_MARKDOWN_TEXT }), + id: 'call_e2e_persisted_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: PERSISTED_FINAL_AGENT_TEXT } + } + ]); + + await dismissAllNotifications(); + await clickRunAll(AGENT_FILE); + + await awaitWebviewMarkers( + [PERSISTED_PYTHON_OUTPUT_MARKER, PERSISTED_FINAL_AGENT_TEXT], + AGENT_RUN_TIMEOUT, + 'agent run whose transcript must survive the save' + ); + + // With snapshots on (the default) the main .deepnote has outputs stripped, so the agent's + // transcript only exists in the sidecar. The save is deferred off queue completion, hence + // the poll rather than a single read. + const snapshotsDir = path.join(workspaceDir, 'snapshots'); + const driver = VSBrowser.instance.driver; + const deadline = Date.now() + SNAPSHOT_WRITE_TIMEOUT; + let transcript = ''; + + while (Date.now() < deadline) { + const files = fs.existsSync(snapshotsDir) + ? fs.readdirSync(snapshotsDir).filter((file) => file.endsWith('_latest.snapshot.deepnote')) + : []; + transcript = + files.length > 0 + ? blockStreamOutputText( + fs.readFileSync(path.join(snapshotsDir, files[0]), 'utf8'), + AGENT_BLOCK_ID + ) + : ''; + + if (transcript.includes(PERSISTED_FINAL_AGENT_TEXT)) { + break; + } + + await driver.sleep(SNAPSHOT_POLL_INTERVAL); + } + + // The first item alone is what a truncating converter leaves behind, so requiring it AND the + // later ones is the whole assertion: the run is on disk and it is not just its opening line. + for (const expected of [ + AGENT_FIRST_OUTPUT_ITEM, + `[Agent] Tool called: ${CODE_TOOL_NAME}`, + `[Agent] Tool called: ${MARKDOWN_TOOL_NAME}`, + PERSISTED_FINAL_AGENT_TEXT + ]) { + if (!transcript.includes(expected)) { + throw new Error( + `Saved snapshot is missing ${JSON.stringify(expected)} — the agent's streamed output items ` + + `did not all survive the save. Saved transcript: ${JSON.stringify(transcript)}` + ); + } + } + }); + }); + + describe('a mixed batch of kernel cells and an agent block', function () { + // Each test here opens and binds the notebook it needs, so both stand alone: run either by + // itself, in either order, or after the group above, and it sets up the same state. + it('stops at a failing cell instead of running the agent block and the cell after it', async function () { + // Scripted so a batch that carried on has something to render. Nothing should reach the mock: + // if these legs are never requested the agent never started, which is the point. + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: `print("${BATCH_AGENT_PYTHON_MARKER}")` }), + id: 'call_e2e_batch_code', + name: CODE_TOOL_NAME + } + } + }, + { + match: { toolResultContains: BATCH_AGENT_PYTHON_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: BATCH_AGENT_MARKDOWN_TEXT }), + id: 'call_e2e_batch_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: BATCH_AGENT_FINAL_TEXT } + } + ]); + + await openOnly(BATCH_FILE); + await selectEnvironmentForNotebook(ENVIRONMENT_NAME, BATCH_FILE); + await dismissAllNotifications(); + await clickRunAll(BATCH_FILE); + + await awaitWebviewMarkers( + [BATCH_FAILURE_MARKER], + FIRST_RUN_OUTPUT_TIMEOUT, + 'traceback from the failing cell' + ); + + await assertMarkersStayAbsent( + [BATCH_AGENT_PYTHON_MARKER, BATCH_AGENT_MARKDOWN_TEXT, BATCH_AGENT_FINAL_TEXT, BATCH_TRAILING_MARKER], + BATCH_SETTLE_WINDOW, + 'a failing cell must end the batch, so neither the agent block nor the cell after it runs' + ); + + await screenshot('batch-stopped-at-failure'); + }); + + it('stops the agent and the cell after it when Interrupt is clicked mid-run', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: STOP_GENERATED_PYTHON }), + id: 'call_e2e_stop_code', + name: CODE_TOOL_NAME + } + } + }, + // Reached only if the agent survived the stop and asked for its next turn. + { + match: { toolResultContains: STOP_SLEEP_MARKER }, + response: { + toolCall: { + arguments: JSON.stringify({ content: STOP_MARKDOWN_TEXT }), + id: 'call_e2e_stop_markdown', + name: MARKDOWN_TOOL_NAME + } + } + }, + { + match: { toolResultContains: MARKDOWN_BLOCK_ADDED_TEXT }, + response: { content: STOP_FINAL_TEXT } + } + ]); + + await openOnly(STOP_FILE); + await selectEnvironmentForNotebook(ENVIRONMENT_NAME, STOP_FILE); + await dismissAllNotifications(); + await clickRunAll(STOP_FILE); + + // Both markers, not just the sleep one: a Mocha retry starts with the previous attempt's + // generated cell still on screen, and the agent's own transcript is cleared at run start, so + // the tool-call line is what proves we are looking at this attempt. + await awaitWebviewMarkers( + [`[Agent] Tool called: ${CODE_TOOL_NAME}`, STOP_SLEEP_MARKER], + AGENT_RUN_TIMEOUT, + 'the generated cell reached its sleep, so the notebook is running and Stop has something to interrupt' + ); + + await clickInterrupt(STOP_FILE); + + await awaitWebviewMarkers( + [AGENT_STOPPED_TEXT], + STOP_ACKNOWLEDGED_TIMEOUT, + 'the agent reports the stop rather than treating the interrupted cell as a retryable failure', + [STOP_MARKDOWN_TEXT, STOP_FINAL_TEXT] + ); + + await assertMarkersStayAbsent( + [STOP_MARKDOWN_TEXT, STOP_FINAL_TEXT, STOP_TRAILING_MARKER], + STOP_SETTLE_WINDOW, + 'a stopped agent ends the batch, so nothing after it runs' + ); + + await screenshot('batch-stopped-by-interrupt'); + }); + }); +});