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 d9bf27dd73..2c3955037a 100644 --- a/build/esbuild/build.ts +++ b/build/esbuild/build.ts @@ -72,7 +72,8 @@ const commonExternals = [ const webExternals = [ ...commonExternals, 'canvas', // Native module used by vega for server-side rendering, not needed in browser - 'mathjax-electron' // Uses Node.js path module, MathJax rendering handled differently in browser + 'mathjax-electron', // Uses Node.js path module, MathJax rendering handled differently in browser + '@deepnote/runtime-core' // Node built-ins (net, child_process); agent blocks run on desktop only ]; const desktopExternals = [...commonExternals, ...deskTopNodeModulesToExternalize]; const bundleConfig = getBundleConfiguration(); diff --git a/cspell.json b/cspell.json index 8aa54f8759..91b895eca3 100644 --- a/cspell.json +++ b/cspell.json @@ -49,6 +49,7 @@ "evalue", "findstr", "getsitepackages", + "hubot", "IMAGENAME", "ipykernel", "ipynb", diff --git a/package.json b/package.json index d2adcbfaff..8b89c5523d 100644 --- a/package.json +++ b/package.json @@ -340,6 +340,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%", @@ -2660,7 +2670,8 @@ "compile-e2e-watch": "tsc -p ./test/e2e/tsconfig.json --watch", "setup:e2e:vscode": "extest get-vscode -c max && extest get-chromedriver -c max", "setup:e2e:deps": "extest install-from-marketplace ms-python.python -e .test-extensions", - "setup:e2e": "npm run setup:e2e:vscode && npm run setup:e2e:deps", + "setup:e2e:mock": "npx -y -p @copilotkit/aimock@1.37.4 llmock --help", + "setup:e2e": "npm run setup:e2e:vscode && npm run setup:e2e:deps && npm run setup:e2e:mock", "test:e2e": "extest setup-and-run \"./out/e2e/suite/*.e2e.test.js\" -c max -o ./test/e2e/settings.json -e .test-extensions -m ./test/e2e/.mocharc.js -i", "test:e2e:prebuilt": "extest run-tests \"./out/e2e/suite/*.e2e.test.js\" -c max -o ./test/e2e/settings.json -e .test-extensions -m ./test/e2e/.mocharc.js", "test:unittests": "mocha --config ./build/.mocha.unittests.js.json ./out/**/*.unit.test.js", diff --git a/package.nls.json b/package.nls.json index f3b21dc525..7d98f46777 100644 --- a/package.nls.json +++ b/package.nls.json @@ -116,6 +116,8 @@ "deepnote.command.deepnote.openOutlineView.title": "Show Table Of Contents (Outline View)", "deepnote.command.deepnote.openOutlineView.shorttitle": "Outline", "deepnote.command.manageAccessToKernels": "Manage Access To Jupyter Kernels", + "deepnote.command.setOpenAiApiKey": "Set OpenAI API Key", + "deepnote.command.clearOpenAiApiKey": "Clear OpenAI API Key", "deepnote.commandPalette.deepnote.replayPylanceLog.title": "Replay Pylance Log", "deepnote.notebookRenderer.IPyWidget.displayName": "Jupyter IPyWidget Renderer", "deepnote.notebookRenderer.Error.displayName": "Jupyter Error Renderer", diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index 9ab88bd119..7f61d0ab59 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -91,6 +91,8 @@ import { RemoteKernelReconnectBusyIndicator } from './remoteKernelReconnectBusyI import { IConnectionDisplayData, IConnectionDisplayDataProvider, IVSCodeNotebookController } from './types'; import { notebookPathToDeepnoteProjectFilePath } from '../../platform/deepnote/deepnoteProjectUtils'; import { DEEPNOTE_NOTEBOOK_TYPE, IDeepnoteKernelAutoSelector } from '../../kernels/deepnote/types'; +import { executeAgentCell, 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 @@ -624,18 +626,48 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont if (!this.cellQueue.has(doc)) { return; } + const queuedCells = this.cellQueue.get(doc) || []; + // Clear before await so agent-driven re-entrant runs start with an empty queue. + this.cellQueue.delete(doc); + + 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`); + await executeAgentCell(cell, this.controller).catch(noop); + } + + await this.executeKernelCells(doc, pendingKernelCells); + } + + private async executeKernelCells(doc: NotebookDocument, cells: NotebookCell[]) { // Start execution now (from the user's point of view) // Creating these execution objects marks the cell as queued for execution (vscode will update cell UI). type CellExec = { cell: NotebookCell; exec: NotebookCellExecution }; - const cellExecs: CellExec[] = (this.cellQueue.get(doc) || []).map((cell) => { + + // Stale cell handles report index -1; createNotebookCellExecution would abort the batch. + 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`); diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts new file mode 100644 index 0000000000..a91c58c348 --- /dev/null +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -0,0 +1,531 @@ +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 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 { 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-level MCP servers and database integrations declared in the `.deepnote` file, matching what + * the CLI's ExecutionEngine passes. `executeAgentBlock` merges the servers with any block-level + * `deepnote_mcp_servers` (block wins on name), and only names the integrations — along with the + * `dntk.execute_sql` instructions — in its system prompt when that list is non-empty, so leaving + * either empty silently drops the project-level half of that contract. + * + * Spawning MCP servers is arbitrary local command execution declared by a workspace file, so every + * caller must already be behind a `workspace.isTrusted` check. + */ +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 }; +} + +// Tool results reported back to the agent. These mirror the wording @deepnote/runtime-core uses in +// its own ExecutionEngine implementation of the same tools, so the agent sees identical phrasing +// whether a block runs in the extension or on the backend. +const MARKDOWN_BLOCK_ADDED_TEXT = 'Markdown block added.'; +const NO_OUTPUT_TEXT = '(no output)'; + +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; +} + +/** + * `translateCellDisplayOutput` follows nbformat's multiline convention and emits text as an array of + * lines — both for stream `text` and for the `text/*` entries of `execute_result`/`display_data` + * `data`. `extractOutputsText` reads stream text only when it is a string, and stringifies + * `data['text/plain']` with `String(...)`, which joins an array with commas. Join the lines first so + * `print()` output isn't dropped and a `df.head()` string representation doesn't reach the agent with a + * comma glued to the start of every line. + */ +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; +} + +/** + * Runs an agent block, streaming its progress into the cell's output and inserting the cells it + * generates below itself. + * + * Requires the cell's previous run to have been cleared first — call + * `removeEphemeralCellsForAgentBlocks` on the batch. Never rejects: failures, including an uncleared + * previous run, are reported on the cell as stderr output and end the execution unsuccessfully. + */ +export async function executeAgentCell( + cell: NotebookCell, + controller: NotebookController, + options?: ExecuteAgentCellOptions +): Promise { + const executeAgentBlockFn = options?.executeAgentBlockFn ?? executeAgentBlock; + const execution = controller.createNotebookCellExecution(cell); + execution.start(Date.now()); + + try { + await execution.clearOutput(); + + const prompt = cell.document.getText(); + + // Streamed as stdout items so each event can be appended rather than re-sending the whole + // transcript: `NotebookCellOutputItem.text` re-encodes the full buffer on every token, which + // is O(n²) bytes across the extension-host boundary — and since runtime-core awaits + // `onAgentEvent` inside its stream loop, that cost is added to the run's wall clock. + // Appended stdout items render as one continuous block, the same way kernel output streams; + // agentBlock.e2e.test.ts pins that, since only a real renderer can show it. + 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'); + } + + // Verify rather than assume the caller cleared them: nothing enforces the precondition across + // the three call sites, and running dirty fails silently — the agent would be handed its own + // previous output as context, and insertEphemeralCell appends below the stale cells rather + // than replacing them, so every run would leave another copy behind. + const staleCellCount = cell.notebook + .getCells() + .filter((c) => getEphemeralCellAgentSourceBlockId(c) === agentBlock.id).length; + + if (staleCellCount > 0) { + throw new Error( + `Agent block ${agentBlock.id} still has ${staleCellCount} generated cell(s) from its previous run` + ); + } + + const openAiToken = await getOrPromptOpenAiApiKey(); + + let lastAgentEventType: AgentStreamEvent['type'] | undefined; + + // serializeNotebookContextFromBlocks does no ephemeral filtering, so this is safe only + // because of the precondition checked above. + const notebookContext = serializeNotebookContext({ + cells: cell.notebook.getCells().filter((c) => c.index !== cell.index), + notebookName: (cell.notebook.metadata?.deepnoteNotebookName as string | undefined) ?? '' + }); + + const context: AgentBlockContext = { + openAiToken, + ...getProjectAgentContext(cell.notebook), + notebookContext, + addMarkdownBlock: async ({ content }: { content: string }) => { + 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 }) => { + try { + const insertedCell = await insertEphemeralCell( + cell.notebook, + cell.index, + agentBlock.id, + 'code', + code + ); + + const { success, outputs, error } = await executeEphemeralCell(insertedCell, execution.token); + const outputText = error ?? describeExecutionOutputs(outputs); + + return success ? `Output:\n${outputText}` : `Execution failed:\n${outputText}`; + } catch (error) { + return `Execution error: ${toError(error).message}`; + } + }, + onAgentEvent: async (event: AgentStreamEvent) => { + 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}`); + + execution.end(true, Date.now()); + } catch (error) { + // `logger.error(msg, error)` only renders `Error.prototype.toString()` unless the error is + // branded with `isJupyterError`, so the stack has to be logged explicitly. + logger.error('Agent cell execution failed', error); + if (error instanceof Error) { + if (error.cause) { + logger.error('Agent error cause:', error.cause); + } + if (error.stack) { + logger.error('Agent error stack:', error.stack); + } + } + + const message = error instanceof Error ? error.message : String(error); + const stderrOutput = new NotebookCellOutput([NotebookCellOutputItem.stderr(message)]); + await execution.appendOutput([stderrOutput]).then(undefined, () => undefined); + execution.end(false, Date.now()); + } +} + +function getInsertIndexAfterAgentCell( + notebook: NotebookDocument, + agentCellIndex: number, + agentBlockId: string +): number { + let index = agentCellIndex + 1; + + while (index < notebook.cellCount) { + if (getEphemeralCellAgentSourceBlockId(notebook.cellAt(index)) === agentBlockId) { + index++; + } else { + break; + } + } + + return index; +} + +/** + * Inserts an ephemeral cell after the agent cell and returns the cell that was actually created. + * + * Resolving by block id rather than by index matters: `cellAt` clamps out-of-range indices instead of + * throwing, so a rejected edit or a concurrent structural change would otherwise hand the caller a + * pre-existing user cell — which `addAndExecuteCodeBlock` would then run. + */ +async function insertEphemeralCell( + notebook: NotebookDocument, + agentCellIndex: number, + agentBlockId: string, + blockType: 'code' | 'markdown', + content: string +): Promise { + 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 { + // Bail before dispatching: rejecting the deferred alone would abandon the wait but still hand the + // generated code to the kernel. + if (token?.isCancellationRequested) { + throw new CancellationError(); + } + + const completionDeferred = createDeferred(); + const disposables: IDisposable[] = []; + + disposables.push( + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell === cell && e.state === NotebookCellExecutionState.Idle) { + completionDeferred.resolve(); + } + }) + ); + + if (token) { + 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; + + // The dispatch settles independently of the cell reaching Idle, so both waits have to start + // together — otherwise the timeout cannot end a run whose command never resolves, and a + // rejection arriving before the second await is reported as unhandled. + await Promise.all([ + commands.executeCommand('notebook.cell.execute', { + ranges: [{ start: cellIndex, end: cellIndex + 1 }], + document: cell.notebook.uri + }), + completionDeferred.promise + ]); + + return { + success: cell.executionSummary?.success === true, + outputs: cell.outputs.map(translateCellDisplayOutput), + executionCount: cell.executionSummary?.executionOrder ?? null + }; + } catch (error) { + if (error instanceof CancellationError) { + throw error; + } + + // Report the reason rather than collapsing everything into "(no output)" — a timed-out cell + // is still running, and telling the agent it produced nothing invites an immediate retry. + return { + success: false, + outputs: [], + executionCount: null, + error: error instanceof Error ? error.message : String(error) + }; + } finally { + dispose(disposables); + clearTimeout(timeout); + } +} + +/** + * Deletes the scratch cells the agent cells in `cells` generated on their previous run, and returns + * the batch without them. Call this before executing a batch of cells; `executeAgentCell` requires it. + * + * Ephemeral cells are agent-owned: the agent regenerates them on every run and the serializer never + * persists them. Left in the batch they would run the previous run's generated code against the + * kernel — so they are dropped up front, whether or not the agent that owns them gets far enough to + * replace them. + * + * Scoped to the agents present in the batch, because two callers legitimately run an ephemeral cell + * on its own: the user selecting one, and the agent executing the code cell it just generated via + * `notebook.cell.execute`. Neither carries its agent, so neither is touched. + * + * A rejected edit is logged rather than thrown: the agent cells re-check the notebook themselves and + * report it on the cell, and a failed edit is no reason to hold back the batch's ordinary code cells. + */ +export async function removeEphemeralCellsForAgentBlocks( + notebook: NotebookDocument, + cells: NotebookCell[] +): Promise { + const agentBlockIds = new Set( + cells + .filter(isAgentCell) + .map(getBlockId) + .filter((id): id is string => typeof id === 'string') + ); + + if (agentBlockIds.size === 0) { + return cells; + } + + const isOwnedScratch = (cell: NotebookCell) => { + const owner = getEphemeralCellAgentSourceBlockId(cell); + + return owner !== undefined && agentBlockIds.has(owner); + }; + + const remainingCells = cells.filter((cell) => !isOwnedScratch(cell)); + const deletions: NotebookEdit[] = []; + + for (let i = notebook.cellCount - 1; i >= 0; i--) { + 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..f8e70c30df --- /dev/null +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -0,0 +1,698 @@ +import { expect } from 'chai'; +import * as sinon from 'sinon'; +import { anything, capture, instance, mock, reset, verify, when } from 'ts-mockito'; +import { + CancellationError, + CancellationTokenSource, + Disposable, + EventEmitter, + ExtensionMode, + NotebookCell, + NotebookCellData, + NotebookCellOutput, + NotebookCellOutputItem, + NotebookController, + NotebookDocument, + SecretStorage, + SecretStorageChangeEvent, + WorkspaceEdit +} from 'vscode'; + +import type { AgentBlock } from '@deepnote/blocks'; +import type { AgentBlockContext, AgentBlockResult } from '@deepnote/runtime-core'; + +import type { IDisposable } from '../../platform/common/types'; +import { IExtensionContext } from '../../platform/common/types'; +import { dispose } from '../../platform/common/utils/lifecycle'; +import { ServiceContainer } from '../../platform/ioc/container'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; +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'; + +/** + * Wires up a ServiceContainer whose IExtensionContext exposes an in-memory SecretStorage, so the + * secret-store helpers take their real code paths instead of the ExtensionMode.Test no-op branch. + */ +function stubSecretStorage(secretStorage: Map): ServiceContainer { + const context = mock(); + const secrets = mock(); + const onDidChangeSecrets = new EventEmitter(); + const serviceContainer = mock(); + + sinon.stub(ServiceContainer, 'instance').get(() => instance(serviceContainer)); + when(serviceContainer.get(IExtensionContext)).thenReturn(instance(context)); + when(context.extensionMode).thenReturn(ExtensionMode.Production); + when(context.secrets).thenReturn(instance(secrets)); + when(secrets.onDidChange).thenReturn(onDidChangeSecrets.event); + when(secrets.get(anything())).thenCall((key: string) => Promise.resolve(secretStorage.get(key))); + when(secrets.store(anything(), anything())).thenCall((key: string, value: string) => { + secretStorage.set(key, value); + + return Promise.resolve(); + }); + + return serviceContainer; +} + +/** + * Makes `workspace.applyEdit` apply the notebook edits it is given to `cells`, so the code under test + * observes its own inserts and deletes. + * + * The mocked `WorkspaceEdit.set` discards the edits it is given, so record them off the prototype + * rather than reading them back off the edit object. + * + * Returns the number of edits applied so far — the shared `workspace` mock is never reset between + * tests, so its own call counts are useless here. + */ +function applyNotebookEditsTo(cells: NotebookCell[], notebook: NotebookDocument) { + type RecordedEdit = { range: { start: number; end: number }; newCells?: NotebookCellData[] }; + let recordedEdits: RecordedEdit[] = []; + let appliedEdits = 0; + + sinon.stub(WorkspaceEdit.prototype, 'set').callsFake((_uri, edits) => { + recordedEdits = edits as unknown as RecordedEdit[]; + }); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { + appliedEdits++; + + for (const notebookEdit of recordedEdits) { + const { start, end } = notebookEdit.range; + const 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'); + }); + + // translateCellDisplayOutput splits `text/plain` into a line array, and @deepnote/blocks + // stringifies it with String(...) — which joins with commas. Without the fix the agent reads + // its own DataFrame output with a comma glued to the start of every line but the first. + test('joins nbformat line arrays in execute_result text/plain', () => { + const output = { + output_type: 'execute_result', + data: { 'text/plain': [' a b\n', '0 1 4\n', '1 2 5'] }, + metadata: {}, + execution_count: 1 + }; + + expect(describeExecutionOutputs([output])).to.equal(' a b\n0 1 4\n1 2 5'); + }); + + test('joins nbformat line arrays in display_data text/plain', () => { + const output = { + output_type: 'display_data', + data: { 'text/plain': ['line one\n', 'line two'] }, + metadata: {} + }; + + expect(describeExecutionOutputs([output])).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; + + setup(() => { + secretStorage.clear(); + secretStorage.set('openAiApiKey', 'test-key'); + mockServiceContainer = stubSecretStorage(secretStorage); + 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' } as AgentBlockResult); + }); + + teardown(() => { + disposables = dispose(disposables); + reset(mockedVSCodeNamespaces.commands); + // Restore the default from vscode-mock rather than reset()ing the whole workspace + // namespace, which other suites rely on. + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(true)); + }); + + function createAgentCell(text: string = 'Test prompt') { + return createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text + }); + } + + /** + * Builds an agent cell inside a notebook whose cell list the test can mutate, and applies + * insert/delete notebook edits to that list so the handler observes its own mutations. + */ + function createAgentCellInMutableNotebook(cells: NotebookCell[] = [], agentBlockId = 'agent-block-1') { + const notebook = createMockNotebook({ cells }); + const agentCell = createMockCell({ + 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 and starts it', async () => { + const cell = createAgentCell('Analyze data'); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect((mockController.createNotebookCellExecution as sinon.SinonStub).calledOnceWith(cell)).to.be.true; + expect(mockExecution.start.calledOnce).to.be.true; + }); + + test('clears output before streaming', async () => { + const cell = createAgentCell('Analyze data'); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.clearOutput.calledOnce).to.be.true; + expect(mockExecution.clearOutput.calledBefore(mockExecution.replaceOutput)).to.be.true; + }); + + test('sets initial output via replaceOutput', async () => { + const cell = createAgentCell('Hello world'); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.replaceOutput.calledOnce).to.be.true; + + const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; + expect(outputs).to.have.lengthOf(1); + expect(outputs[0].items).to.have.lengthOf(1); + + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('[Agent] Planning next steps...'); + }); + + test('streams events via appendOutputItems using onAgentEvent callback', async () => { + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.onAgentEvent?.({ type: 'text_delta', text: 'Hello ' }); + await context.onAgentEvent?.({ type: 'text_delta', text: 'world' }); + + return { finalOutput: 'Hello world' } as AgentBlockResult; + }); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { 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'); + }); + + // Each event must ship only its own delta: re-sending the whole transcript per token is + // O(n²) bytes over the extension-host boundary, and runtime-core awaits this callback. + test('streaming sends only the incremental text per event', async () => { + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.onAgentEvent?.({ type: 'text_delta', text: 'first' }); + await context.onAgentEvent?.({ type: 'text_delta', text: ' second' }); + + return { finalOutput: 'first second' } as AgentBlockResult; + }); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(getStdoutChunkText(0)).to.equal('[Agent] Text:\nfirst'); + expect(getStdoutChunkText(1)).to.equal(' second'); + }); + + test('separates different event types with blank lines', async () => { + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.onAgentEvent?.({ type: 'text_delta', text: 'thinking...' }); + await context.onAgentEvent?.({ type: 'tool_called', toolName: 'search' }); + + return { finalOutput: '' } as AgentBlockResult; + }); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + const chunk2 = getStdoutChunkText(1); + expect(chunk2).to.include('\n\n'); + expect(chunk2).to.include('[Agent] Tool called: search'); + }); + + test('ends execution with success', async () => { + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.true; + }); + + test('ends execution with failure when error occurs', async () => { + mockExecution.clearOutput.rejects(new Error('Test error')); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.false; + }); + + test('writes error message to stderr output on failure', async () => { + mockExecution.clearOutput.rejects(new Error('Something went wrong')); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.appendOutput.calledOnce).to.be.true; + + const outputs = mockExecution.appendOutput.firstCall.args[0] as NotebookCellOutput[]; + expect(outputs).to.have.lengthOf(1); + + const item = outputs[0].items[0]; + expect(item.mime).to.equal('application/vnd.code.notebook.stderr'); + + const text = Buffer.from(item.data).toString('utf-8'); + expect(text).to.equal('Something went wrong'); + }); + + test('handles empty prompt', async () => { + const cell = createAgentCell(''); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.true; + + const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('[Agent] Planning next steps...'); + }); + + test('ends with failure and writes error when API key is not set', async () => { + secretStorage.clear(); + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.false; + expect(mockExecution.appendOutput.calledOnce).to.be.true; + + const outputs = mockExecution.appendOutput.firstCall.args[0] as NotebookCellOutput[]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('OpenAI API key is not set'); + }); + + // Clearing the previous run belongs to the caller. Running against a dirty notebook fails + // silently — the agent gets its own old output as context and appends a second copy below it + // — so the precondition is checked rather than assumed. + test('refuses to run rather than clearing the previous run itself', async () => { + const previousResult = createMockCell({ + text: 'print("previous run")', + metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }, + index: 1 + }); + const { agentCell, cells } = createAgentCellInMutableNotebook([previousResult]); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(executeAgentBlockStub.called).to.be.false; + expect(mockExecution.end.firstCall.args[0]).to.be.false; + expect(cells).to.include(previousResult); + + const [outputs] = mockExecution.appendOutput.firstCall.args as [NotebookCellOutput[]]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('previous run'); + }); + + test('inserts a markdown cell after the agent cell with ephemeral metadata', async () => { + const { agentCell, cells } = createAgentCellInMutableNotebook(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.addMarkdownBlock({ content: '## Findings' }); + + return { finalOutput: '' } as AgentBlockResult; + }); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(cells).to.have.lengthOf(2); + expect(cells[1].document.getText()).to.equal('## Findings'); + expect(cells[1].metadata?.is_ephemeral).to.be.true; + expect(cells[1].metadata?.agent_source_block_id).to.equal(agentCell.metadata?.id); + }); + + test('inserts successive cells after the ones it already added', async () => { + const { agentCell, cells } = createAgentCellInMutableNotebook(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.addMarkdownBlock({ content: 'first' }); + await context.addMarkdownBlock({ content: 'second' }); + + return { finalOutput: '' } as AgentBlockResult; + }); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(cells.map((cell) => cell.document.getText())).to.deep.equal(['Test prompt', 'first', 'second']); + }); + + // cellAt clamps rather than throwing, so resolving the inserted cell by index would hand the + // agent a pre-existing user cell and execute it. + test('fails the tool call without executing anything when the insert edit is rejected', async () => { + const { agentCell } = createAgentCellInMutableNotebook(); + let toolResult: string | undefined; + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(false)); + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + toolResult = await context.addAndExecuteCodeBlock({ code: 'print(1)' }); + + return { finalOutput: '' } as AgentBlockResult; + }); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(toolResult).to.include('Execution error'); + expect(toolResult).to.include('Failed to insert ephemeral code cell'); + verify(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).never(); + }); + + test('passes project MCP servers and integrations to the agent', async () => { + const integrations = [{ id: 'warehouse', name: 'Warehouse', type: 'postgres' }]; + const mcpServers = [{ name: 'files', command: 'mcp-files', args: [] }]; + const notebookManager = mock(); + + when(mockServiceContainer.tryGet(IDeepnoteNotebookManager)).thenReturn( + instance(notebookManager) + ); + when(notebookManager.getProjectForNotebook('project-1', 'notebook-1')).thenReturn( + createDeepnoteFile({ project: createDeepnoteProject({ integrations, settings: { mcpServers } }) }) + ); + + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test prompt', + notebookMetadata: { deepnoteProjectId: 'project-1', deepnoteNotebookId: 'notebook-1' } + }); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + const context = executeAgentBlockStub.firstCall.args[1] as AgentBlockContext; + expect(context.mcpServers).to.deep.equal(mcpServers); + expect(context.integrations).to.deep.equal(integrations); + }); + }); + + suite('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 + }); + } + + /** Wires the cells into a notebook whose list the applied edits actually mutate. */ + function createMutableNotebook(cells: NotebookCell[]) { + const notebook = createMockNotebook({ cells }); + + cells.forEach((cell, index) => { + (cell as { notebook: NotebookDocument }).notebook = notebook; + (cell as { index: number }).index = index; + }); + + return { notebook, ...applyNotebookEditsTo(cells, notebook) }; + } + + test('drops the previous run from the batch and deletes it from the notebook', async () => { + const agentCell = createAgentCell('agent-block-1'); + const previousResult = createEphemeralCell('agent-block-1', 'print("previous run")'); + const cells = [agentCell, previousResult]; + const { notebook } = createMutableNotebook(cells); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [...cells]); + + expect(batch).to.deep.equal([agentCell]); + expect(cells).to.deep.equal([agentCell]); + }); + + test('keeps another agent and ordinary cells', async () => { + const agentCell = createAgentCell('agent-block-1'); + const ownResult = createEphemeralCell('agent-block-1', 'own'); + const otherAgentResult = createEphemeralCell('agent-block-2', 'other agent'); + const userCell = createMockCell({ text: 'user code', metadata: {} }); + const cells = [agentCell, ownResult, otherAgentResult, userCell]; + const { notebook } = createMutableNotebook(cells); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [...cells]); + + expect(batch).to.deep.equal([agentCell, otherAgentResult, userCell]); + expect(cells).to.deep.equal([agentCell, otherAgentResult, userCell]); + }); + + // An agent runs the code cell it just generated through `notebook.cell.execute`, which arrives + // here as a batch of that cell alone. Dropping it would hang the agent until its timeout. + test('leaves an ephemeral cell whose agent is not in the batch', async () => { + const agentCell = createAgentCell('agent-block-1'); + const generatedCell = createEphemeralCell('agent-block-1', 'print("just generated")'); + const cells = [agentCell, generatedCell]; + const { notebook, appliedEdits } = createMutableNotebook(cells); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [generatedCell]); + + expect(batch).to.deep.equal([generatedCell]); + expect(cells).to.deep.equal([agentCell, generatedCell]); + expect(appliedEdits()).to.equal(0); + }); + + test('applies no edit when the batch has no agent cell', async () => { + const userCell = createMockCell({ text: 'user code', metadata: {} }); + const cells = [userCell]; + const { notebook, appliedEdits } = createMutableNotebook(cells); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [...cells]); + + expect(batch).to.deep.equal([userCell]); + expect(appliedEdits()).to.equal(0); + }); + + // The agent cells re-check the notebook and report it on the cell, so a rejected edit must not + // hold back the batch's ordinary code cells. + test('still drops the previous run from the batch when the edit is rejected', async () => { + const agentCell = createAgentCell('agent-block-1'); + const previousResult = createEphemeralCell('agent-block-1', 'print("previous run")'); + const userCell = createMockCell({ text: 'user code', metadata: {} }); + const cells = [agentCell, previousResult, userCell]; + const { notebook } = createMutableNotebook(cells); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(false)); + + const batch = await removeEphemeralCellsForAgentBlocks(notebook, [...cells]); + + expect(batch).to.deep.equal([agentCell, userCell]); + expect(cells).to.deep.equal([agentCell, previousResult, userCell]); + }); + }); + + suite('executeEphemeralCell', () => { + teardown(() => { + reset(mockedVSCodeNamespaces.commands); + }); + + test('uses current cell index, not stale index from insertion time', async () => { + const staleIndex = 5; + const currentIndex = 6; + + const cell = createMockCell({ index: staleIndex }); + + // Simulate a concurrent insertion shifting the cell's index + (cell as { index: number }).index = currentIndex; + + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenCall(async () => { + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Idle); + }); + + await executeEphemeralCell(cell); + + const [commandName, commandArg] = capture( + mockedVSCodeNamespaces.commands.executeCommand as (cmd: string, arg: unknown) => Thenable + ).last(); + + expect(commandName).to.equal('notebook.cell.execute'); + expect(commandArg).to.deep.equal({ + ranges: [{ start: currentIndex, end: currentIndex + 1 }], + document: cell.notebook.uri + }); + }); + + // Rejecting the deferred alone abandons only the wait — the generated code would still reach + // the kernel after the user cancelled. + test('throws without dispatching to the kernel when the token is pre-cancelled', async () => { + const cell = createMockCell({ index: 0 }); + const tokenSource = new CancellationTokenSource(); + tokenSource.cancel(); + + 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); + } finally { + tokenSource.dispose(); + } + + verify(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).never(); + }); + + test('reports the failure reason instead of swallowing it', async () => { + const cell = createMockCell({ index: 0 }); + + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenReject( + new Error('kernel is dead') + ); + + const result = await executeEphemeralCell(cell); + + expect(result.success).to.be.false; + expect(result.error).to.equal('kernel is dead'); + }); + + // The dispatch settles independently of the cell reaching Idle, so waiting on it first would + // leave the timeout unable to end a run whose command never resolves. + test('times out while the dispatch is still pending', async () => { + const cell = createMockCell({ index: 0 }); + const clock = sinon.useFakeTimers({ toFake: ['setTimeout', 'clearTimeout'] }); + + try { + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenCall( + () => new Promise(() => undefined) + ); + + const resultPromise = executeEphemeralCell(cell); + await clock.tickAsync(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'); + } finally { + clock.restore(); + } + }); + }); +}); diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts new file mode 100644 index 0000000000..6cbb2797f5 --- /dev/null +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -0,0 +1,169 @@ +import { + CancellationToken, + Disposable, + EventEmitter, + NotebookCell, + NotebookCellStatusBarItem, + NotebookCellStatusBarItemProvider, + NotebookEdit, + WorkspaceEdit, + commands, + l10n, + notebooks, + window, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { isAgentCell } from './dataConversionUtils'; + +/** The key `agentBlockSchema` defines and `executeAgentBlock` reads. */ +const AGENT_MODEL_METADATA_KEY = 'deepnote_agent_model'; + +/** Must be stored explicitly; a missing key becomes `undefined` in runtime-core and is passed to openai() as the model name. */ +const AGENT_MODEL_AUTO = 'auto'; + +const AGENT_MODEL_OPTIONS = [AGENT_MODEL_AUTO, 'gpt-4o', 'gpt-5']; + +const AGENT_INDICATOR_PRIORITY = 100; +const MODEL_PICKER_PRIORITY = 90; + +@injectable() +export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService { + private readonly disposables: Disposable[] = []; + private readonly _onDidChangeCellStatusBarItems = new EventEmitter(); + + public readonly onDidChangeCellStatusBarItems = this._onDidChangeCellStatusBarItems.event; + + public activate(): void { + this.disposables.push(notebooks.registerNotebookCellStatusBarItemProvider('deepnote', this)); + + this.disposables.push( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook.notebookType === 'deepnote') { + this._onDidChangeCellStatusBarItems.fire(); + } + }) + ); + + this.disposables.push( + commands.registerCommand('deepnote.switchAgentModel', async (cell?: NotebookCell) => { + const activeCell = cell || this.getActiveCell(); + if (activeCell) { + await this.switchModel(activeCell); + } + }) + ); + + this.disposables.push(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); + + return [this.createAgentIndicatorItem(), this.createModelPickerItem(cell, model)]; + } + + 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 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: 'deepnote.switchAgentModel', + arguments: [cell] + } + }; + } + + private getActiveCell(): NotebookCell | undefined { + const activeEditor = window.activeNotebookEditor; + if (activeEditor && activeEditor.selection) { + return activeEditor.notebook.cellAt(activeEditor.selection.start); + } + + return undefined; + } + + private 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..c8463c4a34 --- /dev/null +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -0,0 +1,176 @@ +import { expect } from 'chai'; +import { CancellationToken } from 'vscode'; + +import { AgentCellStatusBarProvider } from './agentCellStatusBarProvider'; +import { createMockCell } from './deepnoteTestHelpers'; + +suite('AgentCellStatusBarProvider', () => { + let provider: AgentCellStatusBarProvider; + let mockToken: CancellationToken; + + setup(() => { + mockToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + provider = new AgentCellStatusBarProvider(); + }); + + teardown(() => { + provider.dispose(); + }); + + suite('Agent Cell Detection', () => { + test('Should return status bar items for agent cell', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.not.be.undefined; + expect(items).to.have.lengthOf(2); + }); + + test('Should return undefined for code cell', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined for sql cell', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'sql' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined for markdown cell', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'markdown' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined for cell without pocket', () => { + const cell = createMockCell({ metadata: {} }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined for cell without metadata', () => { + const cell = createMockCell({ metadata: undefined }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined when cancellation is requested', () => { + const cancelledToken: CancellationToken = { + isCancellationRequested: true, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, cancelledToken); + + expect(items).to.be.undefined; + }); + }); + + suite('Agent Block Indicator', () => { + test('Should display agent block label with icon', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[0].text).to.include('$(hubot)'); + expect(items[0].text).to.include('Agent Block'); + expect(items[0].alignment).to.equal(1); + expect(items[0].priority).to.equal(100); + }); + + test('Should not have a command on the indicator', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[0].command).to.be.undefined; + }); + }); + + suite('Model Picker', () => { + test('Should display "auto" when no model is set', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: auto'); + expect(items[1].text).to.include('$(symbol-enum)'); + }); + + test('Should display configured model from metadata', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_agent_model: 'gpt-4o' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: gpt-4o'); + }); + + test('Should display gpt-5 model', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_agent_model: 'gpt-5' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: gpt-5'); + }); + + test('Should display "auto" when model is empty string', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_agent_model: '' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: auto'); + }); + + test('Should have switch model command', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].command).to.not.be.undefined; + const cmd = items[1].command as any; + expect(cmd.command).to.equal('deepnote.switchAgentModel'); + }); + + test('Should have priority 90', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].priority).to.equal(90); + }); + }); + + suite('Combined metadata', () => { + test('Should ignore metadata keys the runtime does not consume', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_agent_model: 'gpt-4o', + deepnote_max_iterations: 50 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items).to.have.lengthOf(2); + expect(items[0].text).to.include('Agent Block'); + expect(items[1].text).to.include('Model: gpt-4o'); + }); + }); +}); diff --git a/src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts b/src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts new file mode 100644 index 0000000000..09ada758df --- /dev/null +++ b/src/notebooks/deepnote/agentOpenAiApiKeyCommandHandler.ts @@ -0,0 +1,30 @@ +import { inject, injectable } from 'inversify'; +import { commands, l10n, window } from 'vscode'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { IExtensionContext } from '../../platform/common/types'; +import { clearOpenAiApiKey, promptForOpenAiApiKey } from './deepnoteSecretStore'; + +@injectable() +export class AgentOpenAiApiKeyCommandHandler implements IExtensionSyncActivationService { + constructor(@inject(IExtensionContext) private readonly extensionContext: IExtensionContext) {} + + public activate(): void { + this.extensionContext.subscriptions.push( + commands.registerCommand('deepnote.setOpenAiApiKey', () => this.setApiKey()), + commands.registerCommand('deepnote.clearOpenAiApiKey', () => this.clearApiKey()) + ); + } + + private async setApiKey(): Promise { + const key = await promptForOpenAiApiKey(); + if (key) { + void window.showInformationMessage(l10n.t('OpenAI API key has been saved.')); + } + } + + private async clearApiKey(): Promise { + await clearOpenAiApiKey(); + void window.showInformationMessage(l10n.t('OpenAI API key has been cleared.')); + } +} diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.ts b/src/notebooks/deepnote/converters/agentBlockConverter.ts new file mode 100644 index 0000000000..ebf6cdff2c --- /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 prompts render as plaintext code cells; metadata passes through in DeepnoteDataConverter. */ +export class AgentBlockConverter implements BlockConverter { + applyChangesToBlock(block: DeepnoteBlock, cell: NotebookCellData): void { + block.content = cell.value; + } + + canConvert(blockType: string): boolean { + return blockType.toLowerCase() === 'agent'; + } + + convertToCell(block: DeepnoteBlock): NotebookCellData { + const cell = new NotebookCellData(NotebookCellKind.Code, block.content || '', 'plaintext'); + + return cell; + } + + getSupportedTypes(): string[] { + return ['agent']; + } +} diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts new file mode 100644 index 0000000000..a3ce26acf8 --- /dev/null +++ b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts @@ -0,0 +1,198 @@ +import type { DeepnoteBlock } from '@deepnote/blocks'; +import { assert } from 'chai'; +import { NotebookCellData, NotebookCellKind } from 'vscode'; +import { AgentBlockConverter } from './agentBlockConverter'; +import dedent from 'dedent'; + +suite('AgentBlockConverter', () => { + let converter: AgentBlockConverter; + + setup(() => { + converter = new AgentBlockConverter(); + }); + + suite('canConvert', () => { + test('returns true for "agent" type', () => { + assert.strictEqual(converter.canConvert('agent'), true); + }); + + test('returns true for "Agent" type (case insensitive)', () => { + assert.strictEqual(converter.canConvert('Agent'), true); + }); + + test('returns false for other types', () => { + assert.strictEqual(converter.canConvert('code'), false); + assert.strictEqual(converter.canConvert('markdown'), false); + assert.strictEqual(converter.canConvert('sql'), false); + }); + }); + + suite('getSupportedTypes', () => { + test('returns array with "agent"', () => { + const types = converter.getSupportedTypes(); + + assert.deepStrictEqual(types, ['agent']); + }); + }); + + suite('convertToCell', () => { + test('converts agent block to code cell with plaintext language', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Analyze the dataset and create a summary report', + id: 'agent-block-123', + sortingKey: 'a0', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, 'Analyze the dataset and create a summary report'); + assert.strictEqual(cell.languageId, 'plaintext'); + }); + + test('handles empty content', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: '', + id: 'agent-block-456', + sortingKey: 'a1', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, ''); + assert.strictEqual(cell.languageId, 'plaintext'); + }); + + test('handles undefined content', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + id: 'agent-block-789', + sortingKey: 'a2', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, ''); + assert.strictEqual(cell.languageId, 'plaintext'); + }); + + test('preserves multiline prompt', () => { + const prompt = dedent` + You are a senior data analyst. + + Perform a thorough exploratory analysis: + 1. Create a grouped bar chart of revenue by quarter + 2. Create a line chart showing churn rate trends + 3. Compute a pivot table of average revenue + `; + + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: prompt, + id: 'agent-block-multiline', + sortingKey: 'a3', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, prompt); + assert.strictEqual(cell.languageId, 'plaintext'); + }); + + test('preserves agent block with metadata', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Analyze the data', + id: 'agent-block-with-metadata', + metadata: { + deepnote_agent_model: 'gpt-4o' + }, + sortingKey: 'a4', + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, 'Analyze the data'); + assert.strictEqual(cell.languageId, 'plaintext'); + }); + }); + + suite('applyChangesToBlock', () => { + test('updates block content from cell value', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Old prompt', + id: 'agent-block-123', + sortingKey: 'a0', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + const cell = new NotebookCellData( + NotebookCellKind.Code, + 'New prompt with updated instructions', + 'plaintext' + ); + + converter.applyChangesToBlock(block, cell); + + assert.strictEqual(block.content, 'New prompt with updated instructions'); + }); + + test('handles empty cell value', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Some prompt', + id: 'agent-block-456', + sortingKey: 'a1', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + const cell = new NotebookCellData(NotebookCellKind.Code, '', 'plaintext'); + + converter.applyChangesToBlock(block, cell); + + assert.strictEqual(block.content, ''); + }); + + test('does not modify other block properties', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Old prompt', + id: 'agent-block-789', + metadata: { + deepnote_agent_model: 'gpt-4o', + custom: 'value' + }, + sortingKey: 'a2', + type: 'agent' + }; + const cell = new NotebookCellData(NotebookCellKind.Code, 'New prompt', 'plaintext'); + + converter.applyChangesToBlock(block, cell); + + assert.strictEqual(block.content, 'New prompt'); + assert.strictEqual(block.id, 'agent-block-789'); + assert.strictEqual(block.type, 'agent'); + assert.strictEqual(block.sortingKey, 'a2'); + assert.deepStrictEqual(block.metadata, { + deepnote_agent_model: 'gpt-4o', + custom: 'value' + }); + }); + }); +}); diff --git a/src/notebooks/deepnote/dataConversionUtils.ts b/src/notebooks/deepnote/dataConversionUtils.ts index 1b30484770..69fc7a8d86 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,51 @@ export function generateBlockId(): string { return id; } +/** + * Returns true if the cell is backed by an agent block. + * + * Lives here rather than next to the execution handler so callers that only need the predicate + * don't pull `@deepnote/runtime-core` into their module graph. + */ +export function isAgentCell(cell: NotebookCell): boolean { + const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; + + return pocket?.type === 'agent'; +} + +/** + * Returns true if the cell metadata indicates an ephemeral cell (auto-generated by agent). + */ +export function isEphemeralCell(cell: NotebookCell | NotebookCellData): boolean { + return cell.metadata?.is_ephemeral === true; +} + +/** + * Returns the id of the block a cell is backed by, or undefined for a cell that has never been + * serialized. + * + * `__deepnoteBlockId` wins because VS Code may rewrite `id`, which is why the converter mirrors it. + * `deepnoteBlockId` is a third name the fallback-cell path writes; it is only ever set alongside the + * other two, so it resolves nothing new in practice and exists to tolerate metadata that has lost + * them. Losing an id here is worse than reading a redundant one: callers mint a fresh one, which + * reassigns the block on save. + */ +export function getBlockId(cell: NotebookCell | NotebookCellData): string | undefined { + return ( + (cell.metadata?.__deepnoteBlockId as string | undefined) || + (cell.metadata?.id as string | undefined) || + (cell.metadata?.deepnoteBlockId as string | undefined) + ); +} + +/** + * Returns the id of the agent block that generated this ephemeral cell, or undefined if the cell + * isn't agent-generated scratch. + */ +export function getEphemeralCellAgentSourceBlockId(cell: NotebookCell): string | undefined { + return isEphemeralCell(cell) ? (cell.metadata?.agent_source_block_id as string | undefined) : undefined; +} + /** * Generate sorting key based on index (format: a0, a1, ..., a99, b0, b1, ...) */ diff --git a/src/notebooks/deepnote/dataConversionUtils.unit.test.ts b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts new file mode 100644 index 0000000000..2b39574c45 --- /dev/null +++ b/src/notebooks/deepnote/dataConversionUtils.unit.test.ts @@ -0,0 +1,93 @@ +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 cell with code pocket type', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } }); + + expect(isAgentCell(cell)).to.be.false; + }); + + test('returns false for cell with markdown pocket type', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'markdown' } } }); + + expect(isAgentCell(cell)).to.be.false; + }); + + test('returns false for cell without pocket', () => { + const cell = createMockCell({ metadata: {} }); + + expect(isAgentCell(cell)).to.be.false; + }); + + test('returns false for cell without metadata', () => { + const cell = createMockCell({ metadata: undefined }); + + expect(isAgentCell(cell)).to.be.false; + }); + }); + + suite('getBlockId', () => { + test('prefers the backup id VS Code cannot rewrite', () => { + const cell = createMockCell({ metadata: { __deepnoteBlockId: 'backup-id', id: 'rewritten-id' } }); + + expect(getBlockId(cell)).to.equal('backup-id'); + }); + + test('falls back to id when the backup is absent', () => { + const cell = createMockCell({ metadata: { id: 'block-id' } }); + + expect(getBlockId(cell)).to.equal('block-id'); + }); + + // The fallback-cell path writes this third name. Reading it beats minting a fresh id, which + // would reassign the block on save. + test('falls back to the legacy deepnoteBlockId when both are absent', () => { + const cell = createMockCell({ metadata: { deepnoteBlockId: 'legacy-id' } }); + + expect(getBlockId(cell)).to.equal('legacy-id'); + }); + + test('ranks the legacy name below both current ones', () => { + const cell = createMockCell({ metadata: { id: 'block-id', deepnoteBlockId: 'legacy-id' } }); + + expect(getBlockId(cell)).to.equal('block-id'); + }); + + test('returns undefined for a cell that was never serialized', () => { + const cell = createMockCell({ metadata: {} }); + + expect(getBlockId(cell)).to.be.undefined; + }); + }); + + suite('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'); + }); + + // An ordinary cell that happens to carry the metadata is not the agent's to delete. + test('returns undefined when the cell is not marked ephemeral', () => { + const cell = createMockCell({ metadata: { agent_source_block_id: 'agent-block-1' } }); + + expect(getEphemeralCellAgentSourceBlockId(cell)).to.be.undefined; + }); + + test('returns undefined for an ordinary cell', () => { + const cell = createMockCell({ metadata: {} }); + + expect(getEphemeralCellAgentSourceBlockId(cell)).to.be.undefined; + }); + }); +}); diff --git a/src/notebooks/deepnote/deepnoteDataConverter.ts b/src/notebooks/deepnote/deepnoteDataConverter.ts index 9f71700a71..09b06cc2be 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; diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts index c5abf5a442..159b4b4b7c 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.ts @@ -18,6 +18,7 @@ import { IExtensionSyncActivationService } from '../../platform/activation/types import { IDisposableRegistry } from '../../platform/common/types'; import { logger } from '../../platform/logging'; import { IDeepnoteNotebookManager } from '../types'; +import { getBlockId, isEphemeralCell } from './dataConversionUtils'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; import { DeepnoteNotebookSerializer } from './deepnoteSerializer'; @@ -163,7 +164,9 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic * has fewer/no outputs), it's an auto-save of stripped content — skip reload. */ private contentActuallyChanged(notebook: NotebookDocument, newCells: NotebookCellData[]): boolean { - const liveCells = notebook.getCells(); + // Compare against what the serializer persists: ephemeral cells are never written, so + // counting them reads our own save back as an external edit that deletes them. + const liveCells = notebook.getCells().filter((cell) => !isEphemeralCell(cell)); if (liveCells.length !== newCells.length) { return true; } @@ -279,14 +282,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 +303,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic edits.push(NotebookEdit.replaceCells(new NotebookRange(0, notebook.cellCount), newCells)); for (let i = 0; i < newCells.length; i++) { - const blockId = this.getBlockIdFromMetadata(newCells[i].metadata); + const blockId = getBlockId(newCells[i]); if (blockId) { edits.push( NotebookEdit.updateCellMetadata(i, { @@ -376,7 +379,7 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic for (let i = 0; i < liveCells.length; i++) { try { const cell = liveCells[i]; - let blockId = this.getBlockIdFromMetadata(cell.metadata); + let blockId = getBlockId(cell); let blockIdFromFallback = false; // Fallback to original project blocks when metadata was lost @@ -502,10 +505,6 @@ export class DeepnoteFileChangeWatcher implements IExtensionSyncActivationServic logger.info(`[FileChangeWatcher] Updated notebook outputs from external snapshot: ${notebook.uri.path}`); } - private getBlockIdFromMetadata(metadata: Record | undefined): string | undefined { - return (metadata?.__deepnoteBlockId ?? metadata?.id) as string | undefined; - } - private handleFileChange(uri: Uri): void { // Deterministic self-write check — no timers involved if (this.consumeSelfWrite(uri)) { diff --git a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts index cd9eb4d818..2a95a82051 100644 --- a/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteFileChangeWatcher.unit.test.ts @@ -178,6 +178,40 @@ project: assert.strictEqual(applyEditCount, 0, 'applyEdit should not be called when cells match'); }); + test('should skip reload when the live notebook only adds ephemeral cells', async () => { + const uri = Uri.file('/workspace/test.deepnote'); + // The serializer never persists ephemeral cells, so a plain save produces a file with + // fewer cells than the live document. That difference must not read as an external edit. + const notebook = createMockNotebook({ + uri, + cells: [ + { + metadata: { id: 'block-1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("hello")', languageId: 'python' } + }, + { + metadata: { id: 'eph-1', is_ephemeral: true, agent_source_block_id: 'agent-1' }, + outputs: [], + kind: NotebookCellKind.Code, + document: { getText: () => 'print("agent generated")', languageId: 'python' } + } + ] + }); + + when(mockedVSCodeNamespaces.workspace.notebookDocuments).thenReturn([notebook]); + setupMockFs(validYaml); + + onDidChangeFile.fire(uri); + + await waitFor(() => readFileCalls > 0); + await new Promise((resolve) => setTimeout(resolve, autoSaveGraceMs)); + + assert.strictEqual(applyEditCount, 0, 'ephemeral-only difference should not trigger a reload'); + assert.strictEqual(saveCount, 0, 'ephemeral-only difference should not trigger a save'); + }); + test('should reload on external change', async () => { const uri = Uri.file('/workspace/test.deepnote'); const notebook = createMockNotebook({ uri, cellCount: 0 }); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 8c0b93cc4a..eecb22ec2c 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -1089,9 +1089,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 + // Run here only prompts for an environment; kernel execution uses the real controller afterward. controller.executeHandler = async (cells, doc) => { logger.info( `Placeholder controller execute handler called for ${getDisplayPath(doc.uri)} with ${ @@ -1099,6 +1099,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) => { @@ -1117,38 +1123,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 d5f71bc77c..159124e7a0 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -20,7 +20,7 @@ import { IConfigurationService } from '../../platform/common/types'; import { IDeepnoteNotebookManager } from '../types'; import { IKernelProvider, IKernel, IJupyterKernelSpec } from '../../kernels/types'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; -import { NotebookDocument, Uri, NotebookController, CancellationToken } from 'vscode'; +import { EventEmitter, NotebookCell, NotebookDocument, Uri, NotebookController, CancellationToken } from 'vscode'; import { DeepnoteEnvironment } from '../../kernels/deepnote/environments/deepnoteEnvironment'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; @@ -1034,6 +1034,70 @@ 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 = { + index: 0, + metadata: { __deepnotePocket: { type: 'agent' } } + } as unknown as NotebookCell; + const codeCell = { index: 1, metadata: {} } as unknown as NotebookCell; + + 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'); + }); + }); }); /** diff --git a/src/notebooks/deepnote/deepnoteSecretStore.ts b/src/notebooks/deepnote/deepnoteSecretStore.ts new file mode 100644 index 0000000000..9e940557f6 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteSecretStore.ts @@ -0,0 +1,120 @@ +import { ExtensionMode, l10n, window } from 'vscode'; + +import { ServiceContainer } from '../../platform/ioc/container'; +import { IExtensionContext } from '../../platform/common/types'; + +export interface SecretPromptOptions { + prompt: string; + placeHolder?: string; + password?: boolean; +} + +function getContext(): IExtensionContext | null { + const context = ServiceContainer.instance.get(IExtensionContext); + + if (context.extensionMode === ExtensionMode.Test) { + return null; + } + + return context; +} + +export async function getSecret(key: string): Promise { + const context = getContext(); + + if (!context) { + return undefined; + } + + const value = await context.secrets.get(key); + + return value && value.length > 0 ? value : undefined; +} + +export async function setSecret(key: string, value: string): Promise { + const context = getContext(); + + if (!context) { + return; + } + + await context.secrets.store(key, value); +} + +export async function clearSecret(key: string): Promise { + const context = getContext(); + + if (!context) { + return; + } + + await context.secrets.delete(key); +} + +export async function promptForSecret(key: string, options: SecretPromptOptions): Promise { + const input = await window.showInputBox({ + prompt: options.prompt, + placeHolder: options.placeHolder, + password: options.password ?? true, + ignoreFocusOut: true + }); + + if (!input || input.trim().length === 0) { + return undefined; + } + + const trimmed = input.trim(); + await setSecret(key, trimmed); + + return trimmed; +} + +export async function getOrPromptSecret( + key: string, + options: SecretPromptOptions, + errorMessage: string +): Promise { + let value = await getSecret(key); + + if (!value) { + value = await promptForSecret(key, options); + } + + if (!value) { + throw new Error(errorMessage); + } + + return value; +} + +const OPENAI_API_KEY = 'openAiApiKey'; + +const OPENAI_PROMPT_OPTIONS: SecretPromptOptions = { + prompt: l10n.t('Enter your OpenAI API key'), + placeHolder: l10n.t('sk-...'), + password: true +}; + +export async function getOpenAiApiKey(): Promise { + return getSecret(OPENAI_API_KEY); +} + +export async function setOpenAiApiKey(key: string): Promise { + return setSecret(OPENAI_API_KEY, key); +} + +export async function clearOpenAiApiKey(): Promise { + return clearSecret(OPENAI_API_KEY); +} + +export async function promptForOpenAiApiKey(): Promise { + return promptForSecret(OPENAI_API_KEY, OPENAI_PROMPT_OPTIONS); +} + +export async function getOrPromptOpenAiApiKey(): Promise { + return getOrPromptSecret( + OPENAI_API_KEY, + OPENAI_PROMPT_OPTIONS, + l10n.t('OpenAI API key is not set. Use the command "Deepnote: Set OpenAI API Key" to configure it.') + ); +} diff --git a/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts new file mode 100644 index 0000000000..3ba41edd7f --- /dev/null +++ b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts @@ -0,0 +1,241 @@ +import { assert } from 'chai'; +import * as sinon from 'sinon'; +import { anything, instance, mock, when } from 'ts-mockito'; +import { EventEmitter, ExtensionMode, SecretStorage, SecretStorageChangeEvent } from 'vscode'; + +import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; +import { IExtensionContext } from '../../platform/common/types'; +import { ServiceContainer } from '../../platform/ioc/container'; +import { + clearOpenAiApiKey, + clearSecret, + getOpenAiApiKey, + getOrPromptOpenAiApiKey, + getOrPromptSecret, + getSecret, + promptForOpenAiApiKey, + promptForSecret, + setOpenAiApiKey, + setSecret +} from './deepnoteSecretStore'; + +suite('deepnoteSecretStore', () => { + const secretStorage = new Map(); + let context: IExtensionContext; + let secrets: SecretStorage; + let onDidChangeSecrets: EventEmitter; + + setup(() => { + secretStorage.clear(); + context = mock(); + secrets = mock(); + onDidChangeSecrets = new EventEmitter(); + + const serviceContainer = mock(); + sinon.stub(ServiceContainer, 'instance').get(() => instance(serviceContainer)); + when(serviceContainer.get(IExtensionContext)).thenReturn(instance(context)); + when(context.extensionMode).thenReturn(ExtensionMode.Production); + when(context.secrets).thenReturn(instance(secrets)); + when(secrets.onDidChange).thenReturn(onDidChangeSecrets.event); + when(secrets.get(anything())).thenCall((key: string) => Promise.resolve(secretStorage.get(key))); + when(secrets.store(anything(), anything())).thenCall((key: string, value: string) => { + secretStorage.set(key, value); + onDidChangeSecrets.fire({ key }); + + return Promise.resolve(); + }); + when(secrets.delete(anything())).thenCall((key: string) => { + secretStorage.delete(key); + + return Promise.resolve(); + }); + }); + + teardown(() => { + sinon.restore(); + }); + + suite('generic getSecret', () => { + test('returns value when stored', async () => { + secretStorage.set('customKey', 'custom-value'); + + const value = await getSecret('customKey'); + + assert.strictEqual(value, 'custom-value'); + }); + + test('returns undefined when not set', async () => { + const value = await getSecret('customKey'); + + assert.isUndefined(value); + }); + + test('returns undefined when value is empty string', async () => { + secretStorage.set('customKey', ''); + + const value = await getSecret('customKey'); + + assert.isUndefined(value); + }); + }); + + suite('generic setSecret', () => { + test('stores value in secrets', async () => { + await setSecret('customKey', 'custom-value'); + + assert.strictEqual(secretStorage.get('customKey'), 'custom-value'); + }); + }); + + suite('generic clearSecret', () => { + test('deletes value from secrets', async () => { + secretStorage.set('customKey', 'custom-value'); + + await clearSecret('customKey'); + + assert.isFalse(secretStorage.has('customKey')); + }); + }); + + suite('generic promptForSecret', () => { + test('stores and returns value when user enters input', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('user-input')); + + const value = await promptForSecret('customKey', { + prompt: 'Enter value', + placeHolder: 'placeholder', + password: false + }); + + assert.strictEqual(value, 'user-input'); + assert.strictEqual(secretStorage.get('customKey'), 'user-input'); + }); + + test('returns undefined when user cancels', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const value = await promptForSecret('customKey', { prompt: 'Enter value' }); + + assert.isUndefined(value); + }); + }); + + suite('generic getOrPromptSecret', () => { + test('returns value when present in store', async () => { + secretStorage.set('customKey', 'stored-value'); + + const value = await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); + + assert.strictEqual(value, 'stored-value'); + }); + + test('throws when value missing and user cancels prompt', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + try { + await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); + assert.fail('Should have thrown'); + } catch (e) { + assert.strictEqual((e as Error).message, 'Value is required'); + } + }); + }); + + suite('getOpenAiApiKey', () => { + test('returns key when stored', async () => { + secretStorage.set('openAiApiKey', 'test-key'); + + const key = await getOpenAiApiKey(); + + assert.strictEqual(key, 'test-key'); + }); + + test('returns undefined when not set', async () => { + const key = await getOpenAiApiKey(); + + assert.isUndefined(key); + }); + + test('returns undefined when key is empty string', async () => { + secretStorage.set('openAiApiKey', ''); + + const key = await getOpenAiApiKey(); + + assert.isUndefined(key); + }); + }); + + suite('setOpenAiApiKey', () => { + test('stores key in secrets', async () => { + await setOpenAiApiKey('my-api-key'); + + assert.strictEqual(secretStorage.get('openAiApiKey'), 'my-api-key'); + }); + }); + + suite('clearOpenAiApiKey', () => { + test('deletes key from secrets', async () => { + secretStorage.set('openAiApiKey', 'test-key'); + + await clearOpenAiApiKey(); + + assert.isFalse(secretStorage.has('openAiApiKey')); + }); + }); + + suite('promptForOpenAiApiKey', () => { + test('stores and returns key when user enters value', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('sk-abc123')); + + const key = await promptForOpenAiApiKey(); + + assert.strictEqual(key, 'sk-abc123'); + assert.strictEqual(secretStorage.get('openAiApiKey'), 'sk-abc123'); + }); + + test('returns undefined when user cancels', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const key = await promptForOpenAiApiKey(); + + assert.isUndefined(key); + }); + + test('returns undefined when user enters empty string', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(' ')); + + const key = await promptForOpenAiApiKey(); + + assert.isUndefined(key); + }); + }); + + suite('getOrPromptOpenAiApiKey', () => { + test('returns key when present in store', async () => { + secretStorage.set('openAiApiKey', 'stored-key'); + + const key = await getOrPromptOpenAiApiKey(); + + assert.strictEqual(key, 'stored-key'); + }); + + test('prompts and returns key when missing', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('prompted-key')); + + const key = await getOrPromptOpenAiApiKey(); + + assert.strictEqual(key, 'prompted-key'); + }); + + test('throws when key missing and user cancels prompt', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + try { + await getOrPromptOpenAiApiKey(); + assert.fail('Should have thrown'); + } catch (e) { + assert.include((e as Error).message, 'OpenAI API key is not set'); + } + }); + }); +}); diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 5d029bcf12..2c6df36a15 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 { isEphemeralCell } from './dataConversionUtils'; import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; import { SnapshotService } from './snapshots/snapshotService'; import { computeHash } from '../../platform/common/crypto'; @@ -230,11 +231,17 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { throw new Error(`Notebook with ID ${notebookId} not found in project`); } - logger.debug(`SerializeNotebook: Found notebook, converting ${data.cells.length} cells to blocks`); + // Exclude ephemeral cells (agent-generated) from persistence + const nonEphemeralCells = data.cells.filter((cell) => !isEphemeralCell(cell)); + + logger.debug( + `SerializeNotebook: Found notebook, converting ${nonEphemeralCells.length} cells to blocks ` + + `(${data.cells.length - nonEphemeralCells.length} ephemeral excluded)` + ); // Log cell metadata IDs before conversion - for (let i = 0; i < data.cells.length; i++) { - const cell = data.cells[i]; + for (let i = 0; i < nonEphemeralCells.length; i++) { + const cell = nonEphemeralCells[i]; logger.trace( `SerializeNotebook: cell[${i}] metadata.id=${cell.metadata?.id}, metadata keys=${ cell.metadata ? Object.keys(cell.metadata).join(',') : 'none' @@ -244,7 +251,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { // Clone blocks while removing circular references that may have been // introduced by VS Code's notebook cell/output handling - const blocks = this.converter.convertCellsToBlocks(data.cells); + const blocks = this.converter.convertCellsToBlocks(nonEphemeralCells); logger.debug(`SerializeNotebook: Converted to ${blocks.length} blocks`); @@ -258,7 +265,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } // Add snapshot metadata to blocks (contentHash and execution timing) - await this.addSnapshotMetadataToBlocks(blocks, data); + await this.addSnapshotMetadataToBlocks(blocks, { ...data, cells: nonEphemeralCells }); // Handle snapshot mode: strip outputs and execution metadata from main file if (this.snapshotService?.isSnapshotsEnabled()) { diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index 6af7e32ee7..43e0be82cd 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'; diff --git a/src/notebooks/deepnote/deepnoteTestHelpers.ts b/src/notebooks/deepnote/deepnoteTestHelpers.ts index c28ecc04bf..b93e0bae59 100644 --- a/src/notebooks/deepnote/deepnoteTestHelpers.ts +++ b/src/notebooks/deepnote/deepnoteTestHelpers.ts @@ -47,6 +47,7 @@ export interface CreateMockNotebookOptions { notebookType?: string; uri?: Uri; metadata?: Record; + cells?: NotebookCell[]; } /** @@ -56,13 +57,29 @@ 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; } /** @@ -135,7 +152,8 @@ export function createMockCell(options?: CreateMockCellOptions): NotebookCell { positionAt: () => ({}) as unknown, validateRange: () => ({}) as unknown, validatePosition: () => ({}) as unknown, - getWordRangeAtPosition: () => undefined + getWordRangeAtPosition: () => undefined, + encoding: 'utf-8' } as unknown as TextDocument; return { diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts new file mode 100644 index 0000000000..2e1108b52e --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts @@ -0,0 +1,121 @@ +import { + Disposable, + NotebookCell, + NotebookDocument, + OverviewRulerLane, + Range, + TextEditor, + TextEditorDecorationType, + ThemeColor, + window, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import { isEphemeralCell } from './dataConversionUtils'; + +const NOTEBOOK_CELL_SCHEME = 'vscode-notebook-cell'; + +/** Code cell editor decorations for `is_ephemeral` blocks; markup uses `src/renderers/client/markdown.ts`. */ +@injectable() +export class EphemeralCellDecorationProvider implements IExtensionSyncActivationService { + private readonly disposables: Disposable[] = []; + + private ephemeralDecorationType!: TextEditorDecorationType; + + public activate(): void { + this.ephemeralDecorationType = window.createTextEditorDecorationType({ + opacity: '0.8', + isWholeLine: true, + overviewRulerColor: new ThemeColor('charts.yellow'), + overviewRulerLane: OverviewRulerLane.Left, + before: { + contentText: '\u200B', + width: '3px', + backgroundColor: new ThemeColor('charts.yellow'), + margin: '0 8px 0 0' + } + }); + + this.disposables.push(this.ephemeralDecorationType); + + this.disposables.push( + window.onDidChangeVisibleTextEditors(() => { + this.updateDecorations(); + }) + ); + + this.disposables.push( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook.notebookType === 'deepnote') { + this.updateDecorations(); + } + }) + ); + + this.updateDecorations(); + } + + public dispose(): void { + for (const disposable of this.disposables) { + disposable.dispose(); + } + } + + private findCellForEditor(editor: TextEditor): NotebookCell | undefined { + const uri = editor.document.uri; + if (uri.scheme !== NOTEBOOK_CELL_SCHEME) { + return undefined; + } + + for (const notebook of workspace.notebookDocuments) { + if (notebook.notebookType !== 'deepnote') { + continue; + } + + const cell = this.findMatchingCell(notebook, editor); + if (cell) { + return cell; + } + } + + return undefined; + } + + private findMatchingCell(notebook: NotebookDocument, editor: TextEditor): NotebookCell | undefined { + for (const cell of notebook.getCells()) { + if (cell.document.uri.toString() === editor.document.uri.toString()) { + return cell; + } + } + + return undefined; + } + + private updateDecorations(): void { + for (const editor of window.visibleTextEditors) { + try { + if (editor.document.uri.scheme !== NOTEBOOK_CELL_SCHEME) { + continue; + } + + const cell = this.findCellForEditor(editor); + if (!cell || !isEphemeralCell(cell)) { + editor.setDecorations(this.ephemeralDecorationType, []); + continue; + } + + const lineRanges: Range[] = []; + for (let i = 0; i < editor.document.lineCount; i++) { + const line = editor.document.lineAt(i); + lineRanges.push(line.range); + } + + editor.setDecorations(this.ephemeralDecorationType, lineRanges); + } catch { + continue; + } + } + } +} diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts new file mode 100644 index 0000000000..12a7f08d0b --- /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); + } + + private createEphemeralIndicatorItem(agentSourceBlockId?: string): NotebookCellStatusBarItem { + const tooltipLines = [l10n.t('Auto-generated ephemeral block')]; + if (agentSourceBlockId) { + tooltipLines.push(l10n.t('Source agent block: {0}', agentSourceBlockId)); + } + + return { + text: `$(sparkle) ${l10n.t('Ephemeral')}`, + alignment: 1, + priority: EPHEMERAL_INDICATOR_PRIORITY, + tooltip: tooltipLines.join('\n') + }; + } +} diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts new file mode 100644 index 0000000000..27e53dcdf2 --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts @@ -0,0 +1,169 @@ +import { expect } from 'chai'; +import { CancellationToken } from 'vscode'; + +import { EphemeralCellStatusBarProvider } from './ephemeralCellStatusBarProvider'; +import { createMockCell } from './deepnoteTestHelpers'; + +suite('EphemeralCellStatusBarProvider', () => { + let provider: EphemeralCellStatusBarProvider; + let mockToken: CancellationToken; + + setup(() => { + mockToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + provider = new EphemeralCellStatusBarProvider(); + }); + + teardown(() => { + provider.dispose(); + }); + + suite('Ephemeral Cell Detection', () => { + test('Should return a status bar item for ephemeral cell', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + + test('Should return undefined when is_ephemeral is false', () => { + const cell = createMockCell({ metadata: { is_ephemeral: false } }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined when is_ephemeral is not set', () => { + const cell = createMockCell({ metadata: {} }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined for cell without metadata', () => { + const cell = createMockCell({ metadata: undefined }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined when is_ephemeral is a non-boolean truthy value', () => { + const cell = createMockCell({ metadata: { is_ephemeral: 'true' } }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined when cancellation is requested', () => { + const cancelledToken: CancellationToken = { + isCancellationRequested: true, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, cancelledToken); + + expect(item).to.be.undefined; + }); + }); + + suite('Status Bar Item Properties', () => { + test('Should display sparkle icon with Ephemeral label', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.text).to.include('$(sparkle)'); + expect(item.text).to.include('Ephemeral'); + }); + + test('Should have left alignment', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.alignment).to.equal(1); + }); + + test('Should have priority 1000 to appear before all other items', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.priority).to.equal(1000); + }); + + test('Should not have a command', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.command).to.be.undefined; + }); + }); + + suite('Tooltip', () => { + test('Should include auto-generated description in tooltip', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.tooltip).to.include('Auto-generated ephemeral block'); + }); + + test('Should include agent source block ID in tooltip when present', () => { + const cell = createMockCell({ + metadata: { + is_ephemeral: true, + agent_source_block_id: 'a0000000000000000000000000000004' + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.tooltip).to.include('a0000000000000000000000000000004'); + expect(item.tooltip).to.include('Source agent block'); + }); + + test('Should not include source block line in tooltip when agent_source_block_id is absent', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.tooltip).to.not.include('Source agent block'); + }); + }); + + suite('Coexistence with other cell types', () => { + test('Should return item for ephemeral agent cell', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + is_ephemeral: true, + agent_source_block_id: 'source-id' + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + + test('Should return item for ephemeral code cell', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'code' }, + is_ephemeral: true + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + + test('Should return item for ephemeral markdown cell', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'markdown' }, + is_ephemeral: true + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + }); +}); diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index eb87ce58a3..a70e84ac1c 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -94,7 +94,11 @@ import { DeepnoteExtensionSidecarWriter } from '../kernels/deepnote/environments import { DeepnoteNotebookEnvironmentMapper } from '../kernels/deepnote/environments/deepnoteNotebookEnvironmentMapper.node'; import { DeepnoteNotebookCommandListener } from './deepnote/deepnoteNotebookCommandListener'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; +import { AgentCellStatusBarProvider } from './deepnote/agentCellStatusBarProvider'; +import { AgentOpenAiApiKeyCommandHandler } from './deepnote/agentOpenAiApiKeyCommandHandler'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; +import { EphemeralCellDecorationProvider } from './deepnote/ephemeralCellDecorationProvider'; +import { EphemeralCellStatusBarProvider } from './deepnote/ephemeralCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIntegrationStartupCodeProvider'; import { DeepnoteCellCopyHandler } from './deepnote/deepnoteCellCopyHandler'; @@ -261,6 +265,22 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentOpenAiApiKeyCommandHandler + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellDecorationProvider + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNewCellLanguageService diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index 28937d6b60..8e7b5e665a 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/src/notebooks/serviceRegistry.web.ts @@ -50,7 +50,11 @@ 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 { EphemeralCellDecorationProvider } from './deepnote/ephemeralCellDecorationProvider'; +import { EphemeralCellStatusBarProvider } from './deepnote/ephemeralCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler'; @@ -127,6 +131,22 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentOpenAiApiKeyCommandHandler + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellDecorationProvider + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNewCellLanguageService 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..73b317561a 100644 --- a/src/platform/deepnote/pocket.unit.test.ts +++ b/src/platform/deepnote/pocket.unit.test.ts @@ -135,6 +135,20 @@ suite('Pocket', () => { assert.strictEqual((block as any).outputs, undefined); }); + test('takes the id from the backup rather than a rewritten id', () => { + const cell = new NotebookCellData(NotebookCellKind.Code, 'print("hello")', 'python'); + + cell.metadata = { + __deepnotePocket: { type: 'code', sortingKey: 'a0' }, + __deepnoteBlockId: 'block-123', + id: 'rewritten-by-vscode' + }; + + const block = createBlockFromPocket(cell, 0); + + assert.strictEqual(block.id, 'block-123'); + }); + test('creates block with generated ID and sortingKey when no pocket exists', () => { const cell = new NotebookCellData(NotebookCellKind.Code, 'print("hello")', 'python'); diff --git a/src/renderers/client/markdown.ts b/src/renderers/client/markdown.ts index b5399de2df..8537598b5e 100644 --- a/src/renderers/client/markdown.ts +++ b/src/renderers/client/markdown.ts @@ -1,3 +1,30 @@ +import type { ActivationFunction } from 'vscode-notebook-renderer'; + +// Minimal markdown-it surface (no package types; transitive dependency only). +interface MarkdownItToken { + content: string; +} + +interface MarkdownItRuleState { + Token: new (type: string, tag: string, nesting: number) => MarkdownItToken; + env?: { + outputItem?: { + metadata?: Record; + }; + }; + tokens: MarkdownItToken[]; +} + +interface MarkdownIt { + core: { + ruler: { + push(name: string, rule: (state: MarkdownItRuleState) => void): void; + }; + }; +} + +type ExtendMarkdownIt = (callback: (md: MarkdownIt) => void) => void; + const styleContent = ` .alert { width: auto; @@ -31,13 +58,59 @@ const styleContent = ` background-color: rgb(255,205,210); color: rgb(183,28,28); } + +.ephemeral-cell { + border-left: 3px solid var(--vscode-charts-yellow, #cca700); + padding-left: 8px; + opacity: 0.8; +} +.ephemeral-badge { + display: inline-block; + font-size: 0.75em; + padding: 1px 6px; + border-radius: 3px; + background: var(--vscode-charts-yellow, #cca700); + color: var(--vscode-editor-background, #1e1e1e); + margin-bottom: 4px; + font-weight: 600; + letter-spacing: 0.03em; +} `; -export async function activate() { +export const activate: ActivationFunction = async (ctx) => { const style = document.createElement('style'); style.textContent = styleContent; const template = document.createElement('template'); template.classList.add('markdown-style'); template.content.appendChild(style); document.head.appendChild(template); + + const markdownRenderer = await ctx.getRenderer('vscode.markdown-it-renderer'); + const extendMarkdownIt = markdownRenderer?.extendMarkdownIt as ExtendMarkdownIt | undefined; + + if (typeof extendMarkdownIt === 'function') { + extendMarkdownIt((md) => { + addEphemeralCellWrapper(md); + }); + } + + return undefined; +}; + +function addEphemeralCellWrapper(md: MarkdownIt): void { + md.core.ruler.push('ephemeral_wrapper', (state) => { + const metadata = state.env?.outputItem?.metadata; + if (!metadata || metadata.is_ephemeral !== true) { + return; + } + + const openToken = new state.Token('html_block', '', 0); + openToken.content = '
\u2728 Ephemeral\n'; + + const closeToken = new state.Token('html_block', '', 0); + closeToken.content = '
\n'; + + state.tokens.unshift(openToken); + state.tokens.push(closeToken); + }); } diff --git a/src/test/mocks/deepnoteRuntimeCore.ts b/src/test/mocks/deepnoteRuntimeCore.ts index 26e25d6797..4dccd7953f 100644 --- a/src/test/mocks/deepnoteRuntimeCore.ts +++ b/src/test/mocks/deepnoteRuntimeCore.ts @@ -1,9 +1,11 @@ -import type { ServerInfo, ServerOptions } from '@deepnote/runtime-core'; +import type { 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. + * kill Python processes, and the real executeAgentBlock calls the OpenAI API, so this records + * calls and returns fake results instead. * * build/mocha-esm-loader.js resolves the '@deepnote/runtime-core' specifier to this module, * so code under test and tests importing the __ helpers below share one module instance. @@ -12,6 +14,8 @@ import type { ChildProcess } from 'child_process'; type RuntimeCore = typeof import('@deepnote/runtime-core'); +const executeAgentBlockCalls: { block: AgentBlock; context: AgentBlockContext }[] = []; +const serializeNotebookContextFromBlocksCalls: { blockCount: number; notebookName: string }[] = []; const startServerCalls: ServerOptions[] = []; const stopServerCalls: ServerInfo[] = []; let nextServerId = 0; @@ -27,6 +31,21 @@ function makeFakeProcess(id: number): ChildProcess { } as unknown as ChildProcess; } +export const executeAgentBlock: RuntimeCore['executeAgentBlock'] = async (block, context) => { + executeAgentBlockCalls.push({ block, context }); + + return { finalOutput: '' }; +}; + +export const serializeNotebookContextFromBlocks: RuntimeCore['serializeNotebookContextFromBlocks'] = ({ + blocks, + notebookName +}) => { + serializeNotebookContextFromBlocksCalls.push({ blockCount: blocks.length, notebookName }); + + return `notebook:${notebookName} blocks:${blocks.length}`; +}; + export const startServer: RuntimeCore['startServer'] = async (options) => { startServerCalls.push(options); @@ -49,6 +68,14 @@ export const stopServer: RuntimeCore['stopServer'] = async (info) => { }; // Test-only helpers (prefixed with __ to signal they are not part of the real API). +export function __getExecuteAgentBlockCalls(): { block: AgentBlock; context: AgentBlockContext }[] { + return executeAgentBlockCalls; +} + +export function __getSerializeNotebookContextFromBlocksCalls(): { blockCount: number; notebookName: string }[] { + return serializeNotebookContextFromBlocksCalls; +} + export function __getStartServerCalls(): ServerOptions[] { return startServerCalls; } @@ -62,6 +89,8 @@ export function __setStartServerImpl(impl: RuntimeCore['startServer'] | null): v } export function __resetRuntimeCoreMock(): void { + executeAgentBlockCalls.length = 0; + serializeNotebookContextFromBlocksCalls.length = 0; startServerCalls.length = 0; stopServerCalls.length = 0; nextServerId = 0; diff --git a/test/e2e/.mocharc.js b/test/e2e/.mocharc.js index 14040de246..50d9e789fd 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 uses `new Mocha(config)` and ignores `require`; load rootHooks here as `rootHooks`. +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.deepnote b/test/e2e/fixtures/agent-block.deepnote new file mode 100644 index 0000000000..334562ffec --- /dev/null +++ b/test/e2e/fixtures/agent-block.deepnote @@ -0,0 +1,22 @@ +version: '1.0.0' +metadata: + createdAt: '2025-01-01T00:00:00.000Z' + modifiedAt: '2025-01-01T00:00:00.000Z' +project: + id: e2e-agent-block-project + name: E2E Agent Block + notebooks: + - id: e2e-agent-block-notebook + name: Agent Block + blocks: + - id: e2e-agent-block + blockGroup: e2e-agent-group + type: agent + content: |- + Run some Python, then add a markdown block summarising this notebook. + sortingKey: a0 + metadata: + deepnote_agent_model: gpt-5 + executionMode: block + isModule: false + settings: {} diff --git a/test/e2e/helpers/index.ts b/test/e2e/helpers/index.ts index 3ed57bb00d..45a939edc4 100644 --- a/test/e2e/helpers/index.ts +++ b/test/e2e/helpers/index.ts @@ -4,6 +4,7 @@ export * from './constants'; export * from './deepnoteEnvironment'; export * from './deepnoteTree'; export * from './fixtures'; +export * from './mockOpenAiServer'; export * from './modals'; export * from './notebook'; export * from './notifications'; diff --git a/test/e2e/helpers/mockOpenAiServer.ts b/test/e2e/helpers/mockOpenAiServer.ts new file mode 100644 index 0000000000..b998e99af5 --- /dev/null +++ b/test/e2e/helpers/mockOpenAiServer.ts @@ -0,0 +1,194 @@ +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'; + +// aimock is npx-only so its peer deps (jest/vitest) never enter this repo's lockfile. +const AIMOCK_VERSION = '1.37.4'; +const AIMOCK_BIN = 'llmock'; + +// Below typical ephemeral port range so a stray outbound source port cannot fake the pre-flight check. +const MOCK_OPENAI_PORT = 18_937; + +/** + * Set `OPENAI_BASE_URL` for the extension host. Call at spec module scope — ExTester spawns VS Code + * before Mocha `before` hooks, and the host inherits env at spawn time. + */ +export function pointExtensionHostAtMockServer(): void { + process.env.OPENAI_BASE_URL = `http://127.0.0.1:${MOCK_OPENAI_PORT}/v1`; +} + +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; +} + +/** Predicate per scripted leg (not sequence index — safe across Mocha retries). */ +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..83c912ec0b 100644 --- a/test/e2e/helpers/notebook.ts +++ b/test/e2e/helpers/notebook.ts @@ -41,28 +41,46 @@ export async function clickRunAll(notebookFileName: string): Promise { ); } -/** - * Reads the notebook cell output once. - * - * Output lives two iframes deep (iframe.webview.ready -> #active-frame). We only attempt to switch - * when an output webview iframe actually exists (`getViewToSwitchTo`), and we read output-specific - * elements inside the frame — so we never match the cell's source code that is visible in the editor - * of the main document. Returns '' when no output is present yet. - */ -export async function readRenderedOutput(): Promise { +/** Runs `read` inside the notebook output webview; returns '' when the frame is missing or not ready. */ +async function readInsideNotebookWebview(read: (webView: WebView) => Promise): Promise { + const driver = VSBrowser.instance.driver; const webView = new WebView(); - 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); + }); + } +} + +/** Full notebook webview body text (markdown previews and outputs). */ +export async function readNotebookWebviewText(): Promise { + return readInsideNotebookWebview(async (webView) => (await webView.findWebElement(By.css('body'))).getText()); +} + +/** Reads the notebook cell output once, falling back to the whole frame if the renderer used unexpected classes. */ +export async function readRenderedOutput(): Promise { + return readInsideNotebookWebview(async (webView) => { const elements = await webView.findWebElements(By.css(OUTPUT_SELECTOR)); const texts = await Promise.all( elements.map((element) => @@ -73,36 +91,10 @@ export async function readRenderedOutput(): Promise { }) ) ); - text = texts.join('\n').trim(); - - // Fallback: if the renderer used unexpected classes, read the frame body — safe here because - // we have confirmed we are inside the output iframe, not the editor. - if (!text) { - const body = await webView.findWebElement(By.css('body')).catch((error) => { - console.warn('[deepnote-e2e] read output frame body:', error); - - return undefined; - }); - text = body - ? ( - await body.getText().catch((error) => { - console.warn('[deepnote-e2e] read output frame body text:', error); - - return ''; - }) - ).trim() - : ''; - } - } catch (error) { - // Frame went stale or output not painted yet — treat as no output this tick. - console.warn('[deepnote-e2e] read rendered notebook output:', error); - } finally { - await webView.switchBack().catch((error) => { - console.warn('[deepnote-e2e] switch back from notebook output webview:', error); - }); - } + const text = texts.join('\n').trim(); - return text; + return text || (await webView.findWebElement(By.css('body'))).getText(); + }); } /** diff --git a/test/e2e/suite/agentBlock.e2e.test.ts b/test/e2e/suite/agentBlock.e2e.test.ts new file mode 100644 index 0000000000..6f60263260 --- /dev/null +++ b/test/e2e/suite/agentBlock.e2e.test.ts @@ -0,0 +1,222 @@ +/** + * Agent block E2E: three-leg tool loop against a local aimock server (no live OpenAI calls). + * Legs 2–3 match on tool results, so the mock only advances after real kernel stdout and + * markdown tool replies. First kernel run can take minutes (venv + toolkit). + */ + +import { EditorView, InputBox, VSBrowser, WebView, Workbench } from 'vscode-extension-tester'; + +import { + FIRST_RUN_OUTPUT_TIMEOUT, + MockOpenAiServer, + OUTPUT_POLL_INTERVAL, + QUICK_PICK_TIMEOUT, + SUITE_TIMEOUT, + WORKBENCH_TIMEOUT, + clickRunAll, + confirmModalDialog, + copyFixtureToTempDir, + createEnvironment, + createScreenshotter, + dismissAllNotifications, + openFolderViaDialog, + openWorkspaceFile, + pointExtensionHostAtMockServer, + readNotebookWebviewText, + selectEnvironmentForNotebook, + startMockOpenAiServer, + waitForNotification +} from '../helpers'; + +pointExtensionHostAtMockServer(); + +const AGENT_FILE = 'agent-block.deepnote'; +const CODE_TOOL_NAME = 'add_code_block'; +const MARKDOWN_TOOL_NAME = 'add_markdown_block'; +// Coupled to agentCellExecutionHandler tool result for add_markdown_block (leg 3 match). +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 streams content in 20-character chunks, so an answer this long reaches the agent cell as +// several text_delta events — one appended output item each. +const FINAL_AGENT_TEXT = 'Summary added as a markdown block, streamed across several deltas.'; +const MOCK_API_KEY = 'sk-e2e-mock-key'; +const SET_API_KEY_COMMAND = 'Deepnote: Set OpenAI API Key'; +const CLEAR_API_KEY_COMMAND = 'Deepnote: Clear OpenAI API Key'; +const REVERT_FILE_COMMAND = 'File: Revert File'; +const DISCARD_CHANGES_BUTTON = "Don't Save"; +const API_KEY_SAVED_NOTIFICATION = /OpenAI API key has been saved/; + +async function awaitWebviewMarkers(markers: string[], timeout: number, context: string): Promise { + 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)); + if (missing.length === 0) { + return text; + } + + await driver.sleep(OUTPUT_POLL_INTERVAL); + } + + const missing = markers.filter((marker) => !text.includes(marker)); + throw new Error( + `Timed out after ${timeout}ms waiting for notebook webview (${context}). Missing: ${JSON.stringify(missing)}. ` + + `Last text: ${JSON.stringify(text)}` + ); +} + +function assertRenderedContiguously(transcript: string, expected: string): void { + if (transcript.includes(expected)) { + return; + } + + throw new Error( + `Agent transcript does not contain ${JSON.stringify(expected)} as one unbroken run — appended stdout ` + + `items are not rendering as a single block. Full transcript: ${JSON.stringify(transcript)}` + ); +} + +async function storeMockOpenAiApiKey(): Promise { + await new Workbench().executeCommand(SET_API_KEY_COMMAND); + + const input = await InputBox.create(QUICK_PICK_TIMEOUT); + await input.setText(MOCK_API_KEY); + await input.confirm(); + + await waitForNotification(API_KEY_SAVED_NOTIFICATION, QUICK_PICK_TIMEOUT, true); +} + +describe('Deepnote — running an agent block against a stand-in OpenAI API', function () { + this.timeout(SUITE_TIMEOUT); + + let cleanupTempDir: (() => void) | undefined; + let mockServer: MockOpenAiServer | undefined; + let screenshot: (label: string) => Promise; + + before(async function () { + screenshot = createScreenshotter(this); + + const copy = copyFixtureToTempDir(AGENT_FILE); + cleanupTempDir = copy.cleanup; + + await VSBrowser.instance.waitForWorkbench(WORKBENCH_TIMEOUT); + await openFolderViaDialog(copy.tempDir); + await VSBrowser.instance.waitForWorkbench(WORKBENCH_TIMEOUT); + + await openWorkspaceFile(AGENT_FILE); + await VSBrowser.instance.driver.wait( + async () => (await new EditorView().getOpenEditorTitles()).some((title) => title.includes(AGENT_FILE)), + WORKBENCH_TIMEOUT, + `${AGENT_FILE} did not open` + ); + + await createEnvironment(ENVIRONMENT_NAME); + await selectEnvironmentForNotebook(ENVIRONMENT_NAME, AGENT_FILE); + + await dismissAllNotifications(); + await storeMockOpenAiApiKey(); + await screenshot('kernel-connected'); + }); + + 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 new Workbench().executeCommand(CLEAR_API_KEY_COMMAND).catch((error) => { + console.warn('[agent-block] clear the stored OpenAI API key during cleanup:', error); + }); + }); + + it('executes the generated code block, inserts its markdown block, and streams one transcript', async function () { + mockServer = await startMockOpenAiServer([ + { + match: { hasToolResult: false }, + response: { + toolCall: { + arguments: JSON.stringify({ code: GENERATED_PYTHON }), + id: 'call_e2e_code', + name: CODE_TOOL_NAME + } + } + }, + { + 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'); + + // agentCellExecutionHandler appends one stdout item per agent event instead of re-sending the + // transcript, which only pays off if the renderer joins those items back into one block — + // something only a real renderer can show: a section boundary keeps its blank line, and an + // answer split across several deltas arrives unbroken. + assertRenderedContiguously( + transcript, + `[Agent] Tool called: ${MARKDOWN_TOOL_NAME}\n\n[Agent] Tool output: ${MARKDOWN_TOOL_NAME}` + ); + assertRenderedContiguously(transcript, `[Agent] Text:\n${FINAL_AGENT_TEXT}`); + }); +});