diff --git a/build/esbuild/build.ts b/build/esbuild/build.ts index d9bf27dd73..77b7fa66ab 100644 --- a/build/esbuild/build.ts +++ b/build/esbuild/build.ts @@ -72,7 +72,8 @@ const commonExternals = [ const webExternals = [ ...commonExternals, 'canvas', // Native module used by vega for server-side rendering, not needed in browser - 'mathjax-electron' // Uses Node.js path module, MathJax rendering handled differently in browser + 'mathjax-electron', // Uses Node.js path module, MathJax rendering handled differently in browser + '@deepnote/runtime-core' // Uses tcp-port-used → net, only needed in desktop for agent block execution ]; const desktopExternals = [...commonExternals, ...deskTopNodeModulesToExternalize]; const bundleConfig = getBundleConfiguration(); diff --git a/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..dcbc0e6a77 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%", 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..b472167ded 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 } 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,52 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont if (!this.cellQueue.has(doc)) { return; } + const allCells = this.cellQueue.get(doc) || []; + // Cleared before any await so the re-entrant execute request an agent cell issues for its + // generated code starts from an empty queue. + this.cellQueue.delete(doc); + + // Walk in document order rather than running every agent cell first: an agent executes the + // code it generates against the kernel immediately, so it must not overtake the cells above + // it that set up the state it reads. + let pendingKernelCells: NotebookCell[] = []; + + for (const cell of allCells) { + 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) => { + + // An agent run deletes the ephemeral cells it produced last time, and those are ordinary code + // cells that Run All queues. createNotebookCellExecution throws for a cell that has since been + // removed, which would abort the rest of the batch. + const kernelCells = cells.filter((cell) => cell.index >= 0); + + if (kernelCells.length === 0) { + return; + } + + 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..c79742919c --- /dev/null +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -0,0 +1,483 @@ +import { + CancellationError, + CancellationToken, + NotebookCell, + 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, isEphemeralCell } 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: [] }; + } + + 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)'; + +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( + { + kind: cell.kind, + value: cell.document.getText(), + languageId: cell.document.languageId, + metadata: cell.metadata, + outputs: [...(cell.outputs || [])] + }, + cell.index + ); + acc.push(block); + } catch (error) { + logger.error(`Error converting cell to block: ${error}`); + } + return acc; + }, []); + + return serializeNotebookContextFromBlocks({ blocks, notebookName }); +} + +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; +} + +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. + // The stdout mime is the one the renderer concatenates, matching how kernel output streams. + const output = new NotebookCellOutput([NotebookCellOutputItem.stdout(`[Agent] Planning next steps...`)]); + await execution.replaceOutput([output]); + + const dataConverter = new DeepnoteDataConverter(); + const deepnoteBlock = dataConverter.convertCellToBlock( + { + kind: cell.kind, + value: cell.document.getText(), + languageId: cell.document.languageId, + metadata: cell.metadata, + outputs: [...(cell.outputs || [])] + }, + cell.index + ); + const agentBlock: AgentBlock | null = deepnoteBlock.type === 'agent' ? deepnoteBlock : null; + + if (agentBlock == null) { + // TODO: better DX error handling + throw new Error('Cell is not an agent cell'); + } + + // Acquire the key before the destructive cleanup below: it prompts, and throws when the user + // dismisses the prompt, which would otherwise leave the previous run's cells already deleted. + const openAiToken = await getOrPromptOpenAiApiKey(); + + await removeEphemeralCellsForAgent(cell.notebook, agentBlock.id); + + let lastAgentEventType: AgentStreamEvent['type'] | undefined; + + // Must run after the removal — serializeNotebookContextFromBlocks does no ephemeral + // filtering, so the agent would otherwise be handed its own previous scratch cells. + const notebookContext = serializeNotebookContext({ + cells: cell.notebook.getCells().filter((c) => c.index !== cell.index), + notebookName: (cell.notebook.metadata?.deepnoteNotebookName as string | undefined) ?? '' + }); + + const 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) { + const insertError = error instanceof Error ? error : new Error(String(error)); + + return `Failed to add markdown block: ${insertError.message}`; + } + }, + addAndExecuteCodeBlock: async ({ code }: { code: string }) => { + try { + const 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) { + const executionError = error instanceof Error ? error : new Error(String(error)); + + return `Execution error: ${executionError.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) { + const cell = notebook.cellAt(index); + if (isEphemeralCell(cell) && cell.metadata?.agent_source_block_id === 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}`); + } + + // The converter mirrors the block id into `__deepnoteBlockId` precisely because VS Code may + // rewrite `id`, so match on that. + const insertedCell = notebook.getCells().find((c) => c.metadata?.__deepnoteBlockId === block.id); + + if (!insertedCell) { + throw new Error(`Inserted ephemeral ${blockType} cell ${block.id} not found in notebook`); + } + + return insertedCell; +} + +const EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS = 5 * 60 * 1000; + +export interface EphemeralCellExecutionResult { + success: boolean; + outputs: unknown[]; + executionCount: number | null; + /** Why the run failed, when the failure wasn't the cell's own output (cancellation, timeout). */ + error?: string; +} + +export async function executeEphemeralCell( + cell: NotebookCell, + token?: CancellationToken +): Promise { + // 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); + } +} + +async function removeEphemeralCellsForAgent(notebook: NotebookDocument, agentBlockId: string): Promise { + const deletions: NotebookEdit[] = []; + + for (let i = notebook.cellCount - 1; i >= 0; i--) { + const cell = notebook.cellAt(i); + + if (isEphemeralCell(cell) && cell.metadata?.agent_source_block_id === agentBlockId) { + deletions.push(NotebookEdit.deleteCells(new NotebookRange(i, i + 1))); + } + } + + if (deletions.length === 0) { + return; + } + + const edit = new WorkspaceEdit(); + edit.set(notebook.uri, deletions); + + // Fatal rather than a warning: the notebook context the agent receives is read off the live + // document, and Run All keeps these cells out of its kernel batch only by their index going + // negative once they are deleted. + if (!(await workspace.applyEdit(edit))) { + throw new Error(`Failed to remove ephemeral cells for agent block ${agentBlockId}`); + } + + logger.info(`Removed ${deletions.length} ephemeral cell(s) for agent block ${agentBlockId}`); +} diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts new file mode 100644 index 0000000000..910ecd7d47 --- /dev/null +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -0,0 +1,655 @@ +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, + SecretStorage, + SecretStorageChangeEvent, + Uri, + 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 { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; +import { dispose } from '../../platform/common/utils/lifecycle'; +import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; +import { ServiceContainer } from '../../platform/ioc/container'; +import { describeExecutionOutputs, executeAgentCell, executeEphemeralCell } from './agentCellExecutionHandler'; +import { isAgentCell } from './dataConversionUtils'; +import { IDeepnoteNotebookManager } from '../types'; +import { createDeepnoteFile, createDeepnoteProject, createMockCell, createMockNotebook } from './deepnoteTestHelpers'; + +/** + * Wires up a ServiceContainer whose IExtensionContext exposes an in-memory SecretStorage, so the + * secret-store helpers take their real code paths instead of the ExtensionMode.Test no-op branch. + */ +function stubSecretStorage(secretStorage: Map): 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; +} + +suite('AgentCellExecutionHandler', () => { + const secretStorage = new Map(); + let disposables: IDisposable[] = []; + + suite('isAgentCell', () => { + test('returns true for cell with agent pocket type', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + + expect(isAgentCell(cell)).to.be.true; + }); + + test('returns false for cell with code pocket type', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } }); + + expect(isAgentCell(cell)).to.be.false; + }); + + test('returns false for cell with markdown pocket type', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'markdown' } } }); + + expect(isAgentCell(cell)).to.be.false; + }); + + test('returns false for cell without pocket', () => { + const cell = createMockCell({ metadata: {} }); + + expect(isAgentCell(cell)).to.be.false; + }); + + test('returns false for cell without metadata', () => { + const cell = createMockCell({ metadata: undefined }); + + expect(isAgentCell(cell)).to.be.false; + }); + }); + + suite('describeExecutionOutputs', () => { + test('joins nbformat line arrays in stream text', () => { + const output = { + 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. + * + * The mocked `WorkspaceEdit.set` discards the edits it is given, so record them off the + * prototype rather than reading them back off the edit object. + */ + function createAgentCellInMutableNotebook(cells: NotebookCell[] = [], agentBlockId = 'agent-block-1') { + const notebook = createMockNotebook({ cells }); + const agentCell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' }, id: agentBlockId }, + text: 'Test prompt' + }); + + (agentCell as { notebook: typeof notebook }).notebook = notebook; + (agentCell as { index: number }).index = 0; + cells.unshift(agentCell); + + type RecordedEdit = { range: { start: number; end: number }; newCells: NotebookCellData[] }; + let recordedEdits: RecordedEdit[] = []; + + sinon.stub(WorkspaceEdit.prototype, 'set').callsFake((_uri, edits) => { + recordedEdits = edits as unknown as RecordedEdit[]; + }); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => { + for (const notebookEdit of recordedEdits) { + const { start, end } = notebookEdit.range; + const inserted = notebookEdit.newCells.map((cellData) => { + const created = createMockCell({ + text: cellData.value, + metadata: cellData.metadata + }); + (created as { notebook: typeof notebook }).notebook = notebook; + + return created; + }); + + cells.splice(start, end - start, ...inserted); + } + cells.forEach((cell, index) => ((cell as { index: number }).index = index)); + recordedEdits = []; + + return Promise.resolve(true); + }); + + return { agentCell, cells, notebook }; + } + + test('creates execution and starts it', async () => { + const cell = createAgentCell('Analyze data'); + + 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 }); + + const getChunkText = (callIndex: number): string => { + const item = mockExecution.appendOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; + + return Buffer.from(item.data).toString('utf-8'); + }; + + expect(getChunkText(0)).to.equal('[Agent] Text:\nfirst'); + expect(getChunkText(1)).to.equal(' second'); + }); + + test('separates different event types with blank lines', async () => { + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.onAgentEvent?.({ type: 'text_delta', text: 'thinking...' }); + await context.onAgentEvent?.({ type: 'tool_called', toolName: 'search' }); + + return { finalOutput: '' } as AgentBlockResult; + }); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + const getChunkText = (callIndex: number): string => { + const item = mockExecution.appendOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; + + return Buffer.from(item.data).toString('utf-8'); + }; + + const chunk2 = getChunkText(1); + expect(chunk2).to.include('\n\n'); + expect(chunk2).to.include('[Agent] Tool called: search'); + }); + + test('ends execution with success', async () => { + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.true; + }); + + test('ends execution with failure when error occurs', async () => { + mockExecution.clearOutput.rejects(new Error('Test error')); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.false; + }); + + test('writes error message to stderr output on failure', async () => { + mockExecution.clearOutput.rejects(new Error('Something went wrong')); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.appendOutput.calledOnce).to.be.true; + + const outputs = mockExecution.appendOutput.firstCall.args[0] as NotebookCellOutput[]; + expect(outputs).to.have.lengthOf(1); + + const item = outputs[0].items[0]; + expect(item.mime).to.equal('application/vnd.code.notebook.stderr'); + + const text = Buffer.from(item.data).toString('utf-8'); + expect(text).to.equal('Something went wrong'); + }); + + test('handles empty prompt', async () => { + const cell = createAgentCell(''); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.true; + + const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('[Agent] Planning next steps...'); + }); + + test('ends with failure and writes error when API key is not set', async () => { + secretStorage.clear(); + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.false; + expect(mockExecution.appendOutput.calledOnce).to.be.true; + + const outputs = mockExecution.appendOutput.firstCall.args[0] as NotebookCellOutput[]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('OpenAI API key is not set'); + }); + + // The key prompt is the last fallible step before the run starts, so it has to come before + // the cleanup that throws away the previous run's generated cells. + test('keeps previous ephemeral cells when the API key prompt is cancelled', async () => { + secretStorage.clear(); + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const previousResult = createMockCell({ + text: 'print("previous run")', + metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }, + index: 1 + }); + const { agentCell, cells } = createAgentCellInMutableNotebook([previousResult]); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.end.firstCall.args[0]).to.be.false; + expect(cells).to.include(previousResult); + }); + + test('inserts a markdown cell after the agent cell with ephemeral metadata', async () => { + const { agentCell, cells } = createAgentCellInMutableNotebook(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.addMarkdownBlock({ content: '## Findings' }); + + return { finalOutput: '' } as AgentBlockResult; + }); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(cells).to.have.lengthOf(2); + expect(cells[1].document.getText()).to.equal('## Findings'); + expect(cells[1].metadata?.is_ephemeral).to.be.true; + expect(cells[1].metadata?.agent_source_block_id).to.equal(agentCell.metadata?.id); + }); + + test('inserts successive cells after the ones it already added', async () => { + const { agentCell, cells } = createAgentCellInMutableNotebook(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.addMarkdownBlock({ content: 'first' }); + await context.addMarkdownBlock({ content: 'second' }); + + return { finalOutput: '' } as AgentBlockResult; + }); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(cells.map((cell) => cell.document.getText())).to.deep.equal(['Test prompt', 'first', 'second']); + }); + + // cellAt clamps rather than throwing, so resolving the inserted cell by index would hand the + // agent a pre-existing user cell and execute it. + test('fails the tool call without executing anything when the insert edit is rejected', async () => { + const { agentCell } = createAgentCellInMutableNotebook(); + let toolResult: string | undefined; + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(false)); + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); + + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + toolResult = await context.addAndExecuteCodeBlock({ code: 'print(1)' }); + + return { finalOutput: '' } as AgentBlockResult; + }); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(toolResult).to.include('Execution error'); + expect(toolResult).to.include('Failed to insert ephemeral code cell'); + verify(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).never(); + }); + + test('removes only the ephemeral cells belonging to this agent', async () => { + const ownResult = createMockCell({ + text: 'own', + metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }, + index: 1 + }); + const otherAgentResult = createMockCell({ + text: 'other agent', + metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-2' }, + index: 2 + }); + const userCell = createMockCell({ text: 'user code', metadata: {}, index: 3 }); + + const { agentCell, cells } = createAgentCellInMutableNotebook([ownResult, otherAgentResult, userCell]); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(cells).to.not.include(ownResult); + expect(cells).to.include(otherAgentResult); + expect(cells).to.include(userCell); + }); + + // The notebook context is read off the live document, and Run All keeps stale ephemeral cells + // out of the kernel batch only by their index going negative on deletion. + test('fails the run without calling the agent when the cleanup edit is rejected', async () => { + const previousResult = createMockCell({ + text: 'print("previous run")', + metadata: { is_ephemeral: true, agent_source_block_id: 'agent-block-1' }, + index: 1 + }); + const { agentCell } = createAgentCellInMutableNotebook([previousResult]); + + when(mockedVSCodeNamespaces.workspace.applyEdit(anything())).thenCall(() => Promise.resolve(false)); + + await executeAgentCell(agentCell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(executeAgentBlockStub.called).to.be.false; + expect(mockExecution.end.firstCall.args[0]).to.be.false; + }); + + test('passes project MCP servers and integrations to the agent', async () => { + const integrations = [{ id: 'warehouse', name: 'Warehouse', type: 'postgres' }]; + const mcpServers = [{ name: 'files', command: 'mcp-files', args: [] }]; + const notebookManager = mock(); + + when(mockServiceContainer.tryGet(IDeepnoteNotebookManager)).thenReturn( + instance(notebookManager) + ); + when(notebookManager.getProjectForNotebook('project-1', 'notebook-1')).thenReturn( + createDeepnoteFile({ project: createDeepnoteProject({ integrations, settings: { mcpServers } }) }) + ); + + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test prompt', + notebookMetadata: { deepnoteProjectId: 'project-1', deepnoteNotebookId: 'notebook-1' } + }); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + const context = executeAgentBlockStub.firstCall.args[1] as AgentBlockContext; + expect(context.mcpServers).to.deep.equal(mcpServers); + expect(context.integrations).to.deep.equal(integrations); + }); + }); + + suite('executeEphemeralCell', () => { + 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(5 * 60 * 1000); + + const result = await resultPromise; + + expect(result.success).to.be.false; + expect(result.error).to.equal('Ephemeral cell execution timed out'); + } finally { + clock.restore(); + } + }); + }); +}); + +suite('createMockNotebook', () => { + test('reads through to the backing cell array', () => { + const cells: NotebookCell[] = [createMockCell({ text: 'first' })]; + const notebook = createMockNotebook({ cells, uri: Uri.file('/test/mutable.deepnote') }); + + expect(notebook.cellCount).to.equal(1); + + cells.push(createMockCell({ text: 'second', index: 1 })); + + expect(notebook.cellCount).to.equal(2); + expect(notebook.cellAt(1).document.getText()).to.equal('second'); + expect(notebook.getCells()).to.have.lengthOf(2); + }); +}); diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts new file mode 100644 index 0000000000..6ccd26fe6a --- /dev/null +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -0,0 +1,193 @@ +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'; +import { clearOpenAiApiKey, promptForOpenAiApiKey } from './deepnoteSecretStore'; + +/** The key `agentBlockSchema` defines and `executeAgentBlock` reads. */ +const AGENT_MODEL_METADATA_KEY = 'deepnote_agent_model'; + +/** The schema default, and the sentinel runtime-core compares against to fall back to its own choice. */ +const AGENT_MODEL_AUTO = 'auto'; + +const AGENT_MODEL_OPTIONS = [AGENT_MODEL_AUTO, 'gpt-4o', 'gpt-5']; + +/** + * Provides status bar items for agent cells showing the block type indicator + * and the AI model picker. + */ +@injectable() +export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService { + private readonly disposables: Disposable[] = []; + private readonly _onDidChangeCellStatusBarItems = new EventEmitter(); + + public readonly onDidChangeCellStatusBarItems = this._onDidChangeCellStatusBarItems.event; + + public activate(): void { + this.disposables.push(notebooks.registerNotebookCellStatusBarItemProvider('deepnote', this)); + + this.disposables.push( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook.notebookType === 'deepnote') { + this._onDidChangeCellStatusBarItems.fire(); + } + }) + ); + + this.disposables.push( + commands.registerCommand('deepnote.switchAgentModel', async (cell?: NotebookCell) => { + const activeCell = cell || this.getActiveCell(); + if (activeCell) { + await this.switchModel(activeCell); + } + }) + ); + + this.disposables.push( + commands.registerCommand('deepnote.setOpenAiApiKey', async () => { + const key = await promptForOpenAiApiKey(); + if (key) { + void window.showInformationMessage(l10n.t('OpenAI API key has been saved.')); + } + }) + ); + + this.disposables.push( + commands.registerCommand('deepnote.clearOpenAiApiKey', async () => { + await clearOpenAiApiKey(); + void window.showInformationMessage(l10n.t('OpenAI API key has been cleared.')); + }) + ); + + this.disposables.push(this._onDidChangeCellStatusBarItems); + } + + public dispose(): void { + for (const disposable of this.disposables) { + disposable.dispose(); + } + } + + public provideCellStatusBarItems( + cell: NotebookCell, + token: CancellationToken + ): NotebookCellStatusBarItem[] | undefined { + if (token.isCancellationRequested) { + return undefined; + } + + if (!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: 100, + 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: 90, + tooltip: l10n.t('AI Model: {0}\nClick to change', model), + command: { + title: l10n.t('Switch Model'), + command: 'deepnote.switchAgentModel', + arguments: [cell] + } + }; + } + + private getActiveCell(): NotebookCell | undefined { + const activeEditor = window.activeNotebookEditor; + if (activeEditor && activeEditor.selection) { + return activeEditor.notebook.cellAt(activeEditor.selection.start); + } + + return undefined; + } + + private 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; + } + + // Write 'auto' rather than deleting the key: `convertCellToBlock` doesn't re-run the zod + // schema, so a missing key reaches runtime-core as `undefined` — which fails its + // `!== "auto"` check and gets passed to `openai()` as the model name. + await this.updateCellMetadata(cell, { [AGENT_MODEL_METADATA_KEY]: selected.label }); + } + + private async updateCellMetadata(cell: NotebookCell, updates: Record): Promise { + const updatedMetadata = { ...cell.metadata, ...updates }; + + // Remove keys set to undefined so they don't persist + for (const [key, value] of Object.entries(updates)) { + if (value === undefined) { + delete updatedMetadata[key]; + } + } + + 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/converters/agentBlockConverter.ts b/src/notebooks/deepnote/converters/agentBlockConverter.ts new file mode 100644 index 0000000000..6f9ebbd31c --- /dev/null +++ b/src/notebooks/deepnote/converters/agentBlockConverter.ts @@ -0,0 +1,34 @@ +import type { DeepnoteBlock } from '@deepnote/blocks'; +import { NotebookCellData, NotebookCellKind } from 'vscode'; + +import type { BlockConverter } from './blockConverter'; + +/** + * Converter for agent blocks. + * + * Agent blocks are rendered as code cells with plaintext language so the + * natural-language prompt appears without syntax highlighting while remaining + * executable. The prompt text is stored in `block.content`. + * + * Agent-specific metadata (model, MCP servers, max iterations, etc.) is preserved + * through the generic metadata pass-through in DeepnoteDataConverter. + */ +export class AgentBlockConverter implements BlockConverter { + applyChangesToBlock(block: DeepnoteBlock, cell: NotebookCellData): void { + block.content = cell.value || ''; + } + + canConvert(blockType: string): boolean { + return blockType.toLowerCase() === 'agent'; + } + + convertToCell(block: DeepnoteBlock): NotebookCellData { + const cell = new NotebookCellData(NotebookCellKind.Code, block.content || '', 'plaintext'); + + return cell; + } + + getSupportedTypes(): string[] { + return ['agent']; + } +} diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts new file mode 100644 index 0000000000..a3ce26acf8 --- /dev/null +++ b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts @@ -0,0 +1,198 @@ +import type { DeepnoteBlock } from '@deepnote/blocks'; +import { assert } from 'chai'; +import { NotebookCellData, NotebookCellKind } from 'vscode'; +import { AgentBlockConverter } from './agentBlockConverter'; +import dedent from 'dedent'; + +suite('AgentBlockConverter', () => { + let converter: AgentBlockConverter; + + setup(() => { + converter = new AgentBlockConverter(); + }); + + suite('canConvert', () => { + test('returns true for "agent" type', () => { + assert.strictEqual(converter.canConvert('agent'), true); + }); + + test('returns true for "Agent" type (case insensitive)', () => { + assert.strictEqual(converter.canConvert('Agent'), true); + }); + + test('returns false for other types', () => { + assert.strictEqual(converter.canConvert('code'), false); + assert.strictEqual(converter.canConvert('markdown'), false); + assert.strictEqual(converter.canConvert('sql'), false); + }); + }); + + suite('getSupportedTypes', () => { + test('returns array with "agent"', () => { + const types = converter.getSupportedTypes(); + + assert.deepStrictEqual(types, ['agent']); + }); + }); + + suite('convertToCell', () => { + test('converts agent block to code cell with plaintext language', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Analyze the dataset and create a summary report', + id: 'agent-block-123', + sortingKey: 'a0', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, 'Analyze the dataset and create a summary report'); + assert.strictEqual(cell.languageId, 'plaintext'); + }); + + test('handles empty content', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: '', + id: 'agent-block-456', + sortingKey: 'a1', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, ''); + assert.strictEqual(cell.languageId, 'plaintext'); + }); + + test('handles undefined content', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + id: 'agent-block-789', + sortingKey: 'a2', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, ''); + assert.strictEqual(cell.languageId, 'plaintext'); + }); + + test('preserves multiline prompt', () => { + const prompt = dedent` + You are a senior data analyst. + + Perform a thorough exploratory analysis: + 1. Create a grouped bar chart of revenue by quarter + 2. Create a line chart showing churn rate trends + 3. Compute a pivot table of average revenue + `; + + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: prompt, + id: 'agent-block-multiline', + sortingKey: 'a3', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, prompt); + assert.strictEqual(cell.languageId, 'plaintext'); + }); + + test('preserves agent block with metadata', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Analyze the data', + id: 'agent-block-with-metadata', + metadata: { + deepnote_agent_model: 'gpt-4o' + }, + sortingKey: 'a4', + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, 'Analyze the data'); + assert.strictEqual(cell.languageId, 'plaintext'); + }); + }); + + suite('applyChangesToBlock', () => { + test('updates block content from cell value', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Old prompt', + id: 'agent-block-123', + sortingKey: 'a0', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + const cell = new NotebookCellData( + NotebookCellKind.Code, + 'New prompt with updated instructions', + 'plaintext' + ); + + converter.applyChangesToBlock(block, cell); + + assert.strictEqual(block.content, 'New prompt with updated instructions'); + }); + + test('handles empty cell value', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Some prompt', + id: 'agent-block-456', + sortingKey: 'a1', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + const cell = new NotebookCellData(NotebookCellKind.Code, '', 'plaintext'); + + converter.applyChangesToBlock(block, cell); + + assert.strictEqual(block.content, ''); + }); + + test('does not modify other block properties', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Old prompt', + id: 'agent-block-789', + metadata: { + deepnote_agent_model: 'gpt-4o', + custom: 'value' + }, + sortingKey: 'a2', + type: 'agent' + }; + const cell = new NotebookCellData(NotebookCellKind.Code, 'New prompt', 'plaintext'); + + converter.applyChangesToBlock(block, cell); + + assert.strictEqual(block.content, 'New prompt'); + assert.strictEqual(block.id, 'agent-block-789'); + assert.strictEqual(block.type, 'agent'); + assert.strictEqual(block.sortingKey, 'a2'); + assert.deepStrictEqual(block.metadata, { + deepnote_agent_model: 'gpt-4o', + custom: 'value' + }); + }); + }); +}); diff --git a/src/notebooks/deepnote/dataConversionUtils.ts b/src/notebooks/deepnote/dataConversionUtils.ts index 1b30484770..fc9d227ccd 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,25 @@ 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; +} + /** * Generate sorting key based on index (format: a0, a1, ..., a99, b0, b1, ...) */ diff --git a/src/notebooks/deepnote/deepnoteDataConverter.ts b/src/notebooks/deepnote/deepnoteDataConverter.ts index 9f71700a71..51007cae9d 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.ts @@ -11,6 +11,7 @@ import { MarkdownBlockConverter } from './converters/markdownBlockConverter'; import { VisualizationBlockConverter } from './converters/visualizationBlockConverter'; import { compile as convertVegaLiteSpecToVega, ensureVegaLiteLoaded } from './vegaLiteWrapper'; import { produce } from 'immer'; +import { AgentBlockConverter } from './converters/agentBlockConverter'; import { SqlBlockConverter } from './converters/sqlBlockConverter'; import { TextBlockConverter } from './converters/textBlockConverter'; // @ts-ignore - types_unstable subpath requires moduleResolution: "node16" which mandates module: "node16" and .js extensions on all imports @@ -38,6 +39,7 @@ export class DeepnoteDataConverter { private readonly registry = new ConverterRegistry(); constructor() { + this.registry.register(new AgentBlockConverter()); this.registry.register(new CodeBlockConverter()); this.registry.register(new MarkdownBlockConverter()); this.registry.register(new ChartBigNumberBlockConverter()); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 8c0b93cc4a..aed1e61bba 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -57,6 +57,8 @@ import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { IControllerRegistration, IVSCodeNotebookController } from '../controllers/types'; import { IDeepnoteNotebookManager } from '../types'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; +import { executeAgentCell } from './agentCellExecutionHandler'; +import { isAgentCell } from './dataConversionUtils'; import { computeRequirementsHash } from './deepnoteProjectUtils'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; @@ -1089,7 +1091,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, ); controller.supportsExecutionOrder = true; - controller.supportedLanguages = ['python', 'sql', 'markdown']; + controller.supportedLanguages = ['python', 'sql', 'markdown', 'plaintext']; // Execution handler that shows environment picker when user tries to run without an environment controller.executeHandler = async (cells, doc) => { @@ -1099,6 +1101,29 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } cells` ); + // Agent blocks can run arbitrary commands declared in the project file (MCP servers), so + // gate this path the same way VSCodeNotebookController gates its own execute handler. + if (!workspace.isTrusted) { + logger.info(`Workspace is not trusted, skipping execution for ${getDisplayPath(doc.uri)}`); + + return; + } + + const kernelCells = cells.filter((cell) => !isAgentCell(cell)); + + if (kernelCells.length === 0) { + // Nothing needs the kernel, so don't make the user configure an environment first. + for (const cell of cells) { + try { + await executeAgentCell(cell, controller); + } catch (cellError) { + logger.error(`Error executing agent cell ${cell.index}`, cellError); + } + } + + return; + } + // Create a cancellation token that cancels when the notebook is closed const cts = new CancellationTokenSource(); const closeListener = workspace.onDidCloseNotebookDocument((closedDoc) => { @@ -1127,7 +1152,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return; } - logger.info(`Executing ${cells.length} cells through kernel after environment configuration`); + logger.info(`Executing ${cells.length} cells after environment configuration`); // Get or create a kernel for this notebook with the new connection const kernel = this.kernelProvider.getOrCreate(doc, { @@ -1139,12 +1164,19 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Execute cells through the kernel const kernelExecution = this.kernelProvider.getKernelExecution(kernel); + // Document order matters: an agent executes the code it generates immediately, so it + // must not overtake the cells above it that set up the state it reads. for (const cell of cells) { try { - await kernelExecution.executeCell(cell); + if (isAgentCell(cell)) { + // Configuring the environment disposed this placeholder and handed the + // notebook to the real controller, which now owns its executions. + await executeAgentCell(cell, realController.controller); + } else { + await kernelExecution.executeCell(cell); + } } catch (cellError) { logger.error(`Error executing cell ${cell.index}`, cellError); - // Continue with remaining cells } } diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts index d5f71bc77c..2f6c4cae23 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.unit.test.ts @@ -20,7 +20,7 @@ import { IConfigurationService } from '../../platform/common/types'; import { IDeepnoteNotebookManager } from '../types'; import { IKernelProvider, IKernel, IJupyterKernelSpec } from '../../kernels/types'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; -import { NotebookDocument, Uri, NotebookController, CancellationToken } from 'vscode'; +import { EventEmitter, NotebookCell, NotebookDocument, Uri, NotebookController, CancellationToken } from 'vscode'; import { DeepnoteEnvironment } from '../../kernels/deepnote/environments/deepnoteEnvironment'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { getNotebookKey } from '../../platform/deepnote/deepnoteProjectUtils'; @@ -1034,6 +1034,84 @@ suite('DeepnoteKernelAutoSelector - rebuildController', () => { }); }); }); + + suite('Placeholder controller execution', () => { + function createExecutionStub() { + return { + start: sandbox.stub(), + end: sandbox.stub(), + clearOutput: sandbox.stub().resolves(), + replaceOutput: sandbox.stub().resolves(), + appendOutput: sandbox.stub().resolves(), + appendOutputItems: sandbox.stub().resolves() + }; + } + + // Configuring the environment disposes and deselects the placeholder, and VS Code throws for + // executions created on either a disposed or an unassociated controller. + test('runs agent cells on the real controller after configuring the environment', async () => { + const placeholderExecution = createExecutionStub(); + const placeholder = { + supportsExecutionOrder: false, + supportedLanguages: [] as string[], + updateNotebookAffinity: sandbox.stub(), + dispose: sandbox.stub(), + createNotebookCellExecution: sandbox.stub().returns(placeholderExecution) + } as unknown as NotebookController; + + when( + mockedVSCodeNamespaces.notebooks!.createNotebookController(anything(), anything(), anything()) + ).thenReturn(placeholder); + when(mockedVSCodeNamespaces.workspace.isTrusted).thenReturn(true); + + const onDidCloseNotebookDocument = new EventEmitter(); + when(mockedVSCodeNamespaces.workspace.onDidCloseNotebookDocument).thenReturn( + onDidCloseNotebookDocument.event + ); + + const realExecution = createExecutionStub(); + const realNotebookController = { + createNotebookCellExecution: sandbox.stub().returns(realExecution) + } as unknown as NotebookController; + const realController = mock(); + when(realController.controller).thenReturn(realNotebookController); + + const internals = selector as unknown as { + createPlaceholderController(notebook: NotebookDocument): NotebookController; + notebookControllers: Map; + }; + + internals.createPlaceholderController(mockNotebook); + internals.notebookControllers.set(getNotebookKey(mockNotebook.uri), instance(realController)); + + sandbox.stub(selector, 'ensureEnvironmentConfiguredBeforeExecution').resolves(true); + + const kernelExecution = { executeCell: sandbox.stub().resolves() }; + when(mockKernelProvider.getOrCreate(anything(), anything())).thenReturn(instance(mock())); + when(mockKernelProvider.getKernelExecution(anything())).thenReturn( + kernelExecution as unknown as ReturnType + ); + + const agentCell = { + index: 0, + metadata: { __deepnotePocket: { type: 'agent' } } + } as unknown as NotebookCell; + const codeCell = { index: 1, metadata: {} } as unknown as NotebookCell; + + await placeholder.executeHandler!([agentCell, codeCell], mockNotebook, placeholder); + + assert.isTrue( + (realNotebookController.createNotebookCellExecution as sinon.SinonStub).calledOnceWithExactly( + agentCell + ), + 'agent cell execution should be created on the real controller' + ); + assert.isTrue( + (placeholder.createNotebookCellExecution as sinon.SinonStub).notCalled, + 'no execution should be created on the disposed placeholder' + ); + }); + }); }); /** diff --git a/src/notebooks/deepnote/deepnoteSecretStore.ts b/src/notebooks/deepnote/deepnoteSecretStore.ts new file mode 100644 index 0000000000..fadf11bd42 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteSecretStore.ts @@ -0,0 +1,122 @@ +import { ExtensionMode, l10n, window } from 'vscode'; + +import { ServiceContainer } from '../../platform/ioc/container'; +import { IExtensionContext } from '../../platform/common/types'; + +export interface SecretPromptOptions { + prompt: string; + placeHolder?: string; + password?: boolean; +} + +function getContext(): IExtensionContext | null { + const context = ServiceContainer.instance.get(IExtensionContext); + + if (context.extensionMode === ExtensionMode.Test) { + return null; + } + + return context; +} + +export async function getSecret(key: string): Promise { + const context = getContext(); + + if (!context) { + return undefined; + } + + const value = await context.secrets.get(key); + + return value && value.length > 0 ? value : undefined; +} + +export async function setSecret(key: string, value: string): Promise { + const context = getContext(); + + if (!context) { + return; + } + + await context.secrets.store(key, value); +} + +export async function clearSecret(key: string): Promise { + const context = getContext(); + + if (!context) { + return; + } + + await context.secrets.delete(key); +} + +export async function promptForSecret(key: string, options: SecretPromptOptions): Promise { + const input = await window.showInputBox({ + prompt: options.prompt, + placeHolder: options.placeHolder, + password: options.password ?? true, + ignoreFocusOut: true + }); + + if (!input || input.trim().length === 0) { + return undefined; + } + + const trimmed = input.trim(); + await setSecret(key, trimmed); + + return trimmed; +} + +export async function getOrPromptSecret( + key: string, + options: SecretPromptOptions, + errorMessage: string +): Promise { + let value = await getSecret(key); + + if (!value) { + value = await promptForSecret(key, options); + } + + if (!value) { + throw new Error(errorMessage); + } + + return value; +} + +// OpenAI API key - specific wrappers + +const OPENAI_API_KEY = 'openAiApiKey'; + +const OPENAI_PROMPT_OPTIONS: SecretPromptOptions = { + prompt: l10n.t('Enter your OpenAI API key'), + placeHolder: l10n.t('sk-...'), + password: true +}; + +export async function getOpenAiApiKey(): Promise { + return getSecret(OPENAI_API_KEY); +} + +export async function setOpenAiApiKey(key: string): Promise { + return setSecret(OPENAI_API_KEY, key); +} + +export async function clearOpenAiApiKey(): Promise { + return clearSecret(OPENAI_API_KEY); +} + +export async function promptForOpenAiApiKey(): Promise { + return promptForSecret(OPENAI_API_KEY, OPENAI_PROMPT_OPTIONS); +} + +export async function getOrPromptOpenAiApiKey(): Promise { + return getOrPromptSecret( + OPENAI_API_KEY, + OPENAI_PROMPT_OPTIONS, + l10n.t('OpenAI API key is not set. Use the command "Deepnote: Set OpenAI API Key" to configure it.') + ); +} diff --git a/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts new file mode 100644 index 0000000000..e99bafc638 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts @@ -0,0 +1,241 @@ +import { expect } from 'chai'; +import * as sinon from 'sinon'; +import { anything, instance, mock, when } from 'ts-mockito'; +import { EventEmitter, ExtensionMode, SecretStorage, SecretStorageChangeEvent } from 'vscode'; + +import { IExtensionContext } from '../../platform/common/types'; +import { ServiceContainer } from '../../platform/ioc/container'; +import { + clearOpenAiApiKey, + clearSecret, + getOpenAiApiKey, + getOrPromptOpenAiApiKey, + getOrPromptSecret, + getSecret, + promptForOpenAiApiKey, + promptForSecret, + setOpenAiApiKey, + setSecret +} from './deepnoteSecretStore'; +import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; + +suite('deepnoteSecretStore', () => { + const secretStorage = new Map(); + let context: IExtensionContext; + let secrets: SecretStorage; + let onDidChangeSecrets: EventEmitter; + + setup(() => { + secretStorage.clear(); + context = mock(); + secrets = mock(); + onDidChangeSecrets = new EventEmitter(); + + const serviceContainer = mock(); + sinon.stub(ServiceContainer, 'instance').get(() => instance(serviceContainer)); + when(serviceContainer.get(IExtensionContext)).thenReturn(instance(context)); + when(context.extensionMode).thenReturn(ExtensionMode.Production); + when(context.secrets).thenReturn(instance(secrets)); + when(secrets.onDidChange).thenReturn(onDidChangeSecrets.event); + when(secrets.get(anything())).thenCall((key: string) => Promise.resolve(secretStorage.get(key))); + when(secrets.store(anything(), anything())).thenCall((key: string, value: string) => { + secretStorage.set(key, value); + onDidChangeSecrets.fire({ key }); + + return Promise.resolve(); + }); + when(secrets.delete(anything())).thenCall((key: string) => { + secretStorage.delete(key); + + return Promise.resolve(); + }); + }); + + teardown(() => { + sinon.restore(); + }); + + suite('generic getSecret', () => { + test('returns value when stored', async () => { + secretStorage.set('customKey', 'custom-value'); + + const value = await getSecret('customKey'); + + expect(value).to.equal('custom-value'); + }); + + test('returns undefined when not set', async () => { + const value = await getSecret('customKey'); + + expect(value).to.be.undefined; + }); + + test('returns undefined when value is empty string', async () => { + secretStorage.set('customKey', ''); + + const value = await getSecret('customKey'); + + expect(value).to.be.undefined; + }); + }); + + suite('generic setSecret', () => { + test('stores value in secrets', async () => { + await setSecret('customKey', 'custom-value'); + + expect(secretStorage.get('customKey')).to.equal('custom-value'); + }); + }); + + suite('generic clearSecret', () => { + test('deletes value from secrets', async () => { + secretStorage.set('customKey', 'custom-value'); + + await clearSecret('customKey'); + + expect(secretStorage.has('customKey')).to.be.false; + }); + }); + + suite('generic promptForSecret', () => { + test('stores and returns value when user enters input', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('user-input')); + + const value = await promptForSecret('customKey', { + prompt: 'Enter value', + placeHolder: 'placeholder', + password: false + }); + + expect(value).to.equal('user-input'); + expect(secretStorage.get('customKey')).to.equal('user-input'); + }); + + test('returns undefined when user cancels', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const value = await promptForSecret('customKey', { prompt: 'Enter value' }); + + expect(value).to.be.undefined; + }); + }); + + suite('generic getOrPromptSecret', () => { + test('returns value when present in store', async () => { + secretStorage.set('customKey', 'stored-value'); + + const value = await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); + + expect(value).to.equal('stored-value'); + }); + + test('throws when value missing and user cancels prompt', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + try { + await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); + expect.fail('Should have thrown'); + } catch (e) { + expect((e as Error).message).to.equal('Value is required'); + } + }); + }); + + suite('getOpenAiApiKey', () => { + test('returns key when stored', async () => { + secretStorage.set('openAiApiKey', 'test-key'); + + const key = await getOpenAiApiKey(); + + expect(key).to.equal('test-key'); + }); + + test('returns undefined when not set', async () => { + const key = await getOpenAiApiKey(); + + expect(key).to.be.undefined; + }); + + test('returns undefined when key is empty string', async () => { + secretStorage.set('openAiApiKey', ''); + + const key = await getOpenAiApiKey(); + + expect(key).to.be.undefined; + }); + }); + + suite('setOpenAiApiKey', () => { + test('stores key in secrets', async () => { + await setOpenAiApiKey('my-api-key'); + + expect(secretStorage.get('openAiApiKey')).to.equal('my-api-key'); + }); + }); + + suite('clearOpenAiApiKey', () => { + test('deletes key from secrets', async () => { + secretStorage.set('openAiApiKey', 'test-key'); + + await clearOpenAiApiKey(); + + expect(secretStorage.has('openAiApiKey')).to.be.false; + }); + }); + + suite('promptForOpenAiApiKey', () => { + test('stores and returns key when user enters value', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('sk-abc123')); + + const key = await promptForOpenAiApiKey(); + + expect(key).to.equal('sk-abc123'); + expect(secretStorage.get('openAiApiKey')).to.equal('sk-abc123'); + }); + + test('returns undefined when user cancels', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const key = await promptForOpenAiApiKey(); + + expect(key).to.be.undefined; + }); + + test('returns undefined when user enters empty string', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(' ')); + + const key = await promptForOpenAiApiKey(); + + expect(key).to.be.undefined; + }); + }); + + suite('getOrPromptOpenAiApiKey', () => { + test('returns key when present in store', async () => { + secretStorage.set('openAiApiKey', 'stored-key'); + + const key = await getOrPromptOpenAiApiKey(); + + expect(key).to.equal('stored-key'); + }); + + test('prompts and returns key when missing', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('prompted-key')); + + const key = await getOrPromptOpenAiApiKey(); + + expect(key).to.equal('prompted-key'); + }); + + test('throws when key missing and user cancels prompt', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + try { + await getOrPromptOpenAiApiKey(); + expect.fail('Should have thrown'); + } catch (e) { + expect((e as Error).message).to.include('OpenAI API key is not set'); + } + }); + }); +}); diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 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..e00762286e 100644 --- a/src/notebooks/deepnote/deepnoteTestHelpers.ts +++ b/src/notebooks/deepnote/deepnoteTestHelpers.ts @@ -47,6 +47,11 @@ export interface CreateMockNotebookOptions { notebookType?: string; uri?: Uri; metadata?: Record; + /** + * Backing cells. Pass the same array you mutate in the test — `cellAt`/`getCells`/`cellCount` + * read through to it, so edits applied during the test are visible to the code under test. + */ + cells?: NotebookCell[]; } /** @@ -56,13 +61,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 +156,8 @@ export function createMockCell(options?: CreateMockCellOptions): NotebookCell { positionAt: () => ({}) as unknown, validateRange: () => ({}) as unknown, validatePosition: () => ({}) as unknown, - getWordRangeAtPosition: () => undefined + getWordRangeAtPosition: () => undefined, + encoding: 'utf-8' } as unknown as TextDocument; return { diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts new file mode 100644 index 0000000000..19ca8b76fc --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts @@ -0,0 +1,126 @@ +import { + Disposable, + NotebookCell, + NotebookDocument, + OverviewRulerLane, + Range, + TextEditor, + TextEditorDecorationType, + ThemeColor, + window, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { isEphemeralCell } from './dataConversionUtils'; +import { IExtensionSyncActivationService } from '../../platform/activation/types'; + +const NOTEBOOK_CELL_SCHEME = 'vscode-notebook-cell'; + +/** + * Applies visual decorations (left border, background tint, reduced opacity) to + * code cell editors that belong to ephemeral blocks (`is_ephemeral: true`). + * + * The left border is rendered via a `before` pseudo-element on each line, + * which avoids overlapping or shifting the code text. + * + * Markup cells are handled separately by the markdown-it renderer plugin in + * `src/renderers/client/markdown.ts`. + */ +@injectable() +export class EphemeralCellDecorationProvider implements IExtensionSyncActivationService { + private readonly disposables: Disposable[] = []; + + private ephemeralDecorationType!: TextEditorDecorationType; + + public activate(): void { + this.ephemeralDecorationType = window.createTextEditorDecorationType({ + opacity: '0.8', + isWholeLine: true, + overviewRulerColor: new ThemeColor('charts.yellow'), + overviewRulerLane: OverviewRulerLane.Left, + before: { + contentText: '\u200B', + width: '3px', + backgroundColor: new ThemeColor('charts.yellow'), + margin: '0 8px 0 0' + } + }); + + this.disposables.push(this.ephemeralDecorationType); + + this.disposables.push( + window.onDidChangeVisibleTextEditors(() => { + this.updateDecorations(); + }) + ); + + this.disposables.push( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook.notebookType === 'deepnote') { + this.updateDecorations(); + } + }) + ); + + this.updateDecorations(); + } + + public dispose(): void { + for (const disposable of this.disposables) { + disposable.dispose(); + } + } + + private findCellForEditor(editor: TextEditor): NotebookCell | undefined { + const uri = editor.document.uri; + if (uri.scheme !== NOTEBOOK_CELL_SCHEME) { + return undefined; + } + + for (const notebook of workspace.notebookDocuments) { + if (notebook.notebookType !== 'deepnote') { + continue; + } + + const cell = this.findMatchingCell(notebook, editor); + if (cell) { + return cell; + } + } + + return undefined; + } + + private findMatchingCell(notebook: NotebookDocument, editor: TextEditor): NotebookCell | undefined { + for (const cell of notebook.getCells()) { + if (cell.document.uri.toString() === editor.document.uri.toString()) { + return cell; + } + } + + return undefined; + } + + private updateDecorations(): void { + for (const editor of window.visibleTextEditors) { + if (editor.document.uri.scheme !== NOTEBOOK_CELL_SCHEME) { + continue; + } + + const cell = this.findCellForEditor(editor); + if (!cell || !isEphemeralCell(cell)) { + editor.setDecorations(this.ephemeralDecorationType, []); + continue; + } + + const lineRanges: Range[] = []; + for (let i = 0; i < editor.document.lineCount; i++) { + const line = editor.document.lineAt(i); + lineRanges.push(line.range); + } + + editor.setDecorations(this.ephemeralDecorationType, lineRanges); + } + } +} diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts new file mode 100644 index 0000000000..60c0a67fba --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts @@ -0,0 +1,80 @@ +import { + CancellationToken, + Disposable, + EventEmitter, + NotebookCell, + NotebookCellStatusBarItem, + NotebookCellStatusBarItemProvider, + l10n, + notebooks, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { isEphemeralCell } from './dataConversionUtils'; +import { IExtensionSyncActivationService } from '../../platform/activation/types'; + +const EPHEMERAL_INDICATOR_PRIORITY = 1000; + +/** + * Provides a status bar indicator for ephemeral cells — blocks that were + * auto-generated by an agent and marked with `is_ephemeral: true` in metadata. + */ +@injectable() +export class EphemeralCellStatusBarProvider + implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService +{ + private readonly disposables: Disposable[] = []; + private readonly _onDidChangeCellStatusBarItems = new EventEmitter(); + + public readonly onDidChangeCellStatusBarItems = this._onDidChangeCellStatusBarItems.event; + + public activate(): void { + this.disposables.push(notebooks.registerNotebookCellStatusBarItemProvider('deepnote', this)); + + this.disposables.push( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook.notebookType === 'deepnote') { + this._onDidChangeCellStatusBarItems.fire(); + } + }) + ); + + this.disposables.push(this._onDidChangeCellStatusBarItems); + } + + public dispose(): void { + this.disposables.forEach((d) => d.dispose()); + } + + public provideCellStatusBarItems( + cell: NotebookCell, + token: CancellationToken + ): NotebookCellStatusBarItem | undefined { + if (token.isCancellationRequested) { + return undefined; + } + + if (!isEphemeralCell(cell)) { + return undefined; + } + + const agentSourceBlockId = cell.metadata?.agent_source_block_id as string | undefined; + + return this.createEphemeralIndicatorItem(agentSourceBlockId); + } + + private createEphemeralIndicatorItem(agentSourceBlockId?: string): NotebookCellStatusBarItem { + const tooltipLines = [l10n.t('Auto-generated ephemeral block')]; + if (agentSourceBlockId) { + tooltipLines.push(l10n.t('Source agent block: {0}', agentSourceBlockId)); + } + + return { + text: `$(sparkle) ${l10n.t('Ephemeral')}`, + alignment: 1, + priority: EPHEMERAL_INDICATOR_PRIORITY, + tooltip: tooltipLines.join('\n') + }; + } +} diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts new file mode 100644 index 0000000000..27e53dcdf2 --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts @@ -0,0 +1,169 @@ +import { expect } from 'chai'; +import { CancellationToken } from 'vscode'; + +import { EphemeralCellStatusBarProvider } from './ephemeralCellStatusBarProvider'; +import { createMockCell } from './deepnoteTestHelpers'; + +suite('EphemeralCellStatusBarProvider', () => { + let provider: EphemeralCellStatusBarProvider; + let mockToken: CancellationToken; + + setup(() => { + mockToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + provider = new EphemeralCellStatusBarProvider(); + }); + + teardown(() => { + provider.dispose(); + }); + + suite('Ephemeral Cell Detection', () => { + test('Should return a status bar item for ephemeral cell', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + + test('Should return undefined when is_ephemeral is false', () => { + const cell = createMockCell({ metadata: { is_ephemeral: false } }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined when is_ephemeral is not set', () => { + const cell = createMockCell({ metadata: {} }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined for cell without metadata', () => { + const cell = createMockCell({ metadata: undefined }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined when is_ephemeral is a non-boolean truthy value', () => { + const cell = createMockCell({ metadata: { is_ephemeral: 'true' } }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined when cancellation is requested', () => { + const cancelledToken: CancellationToken = { + isCancellationRequested: true, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, cancelledToken); + + expect(item).to.be.undefined; + }); + }); + + suite('Status Bar Item Properties', () => { + test('Should display sparkle icon with Ephemeral label', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.text).to.include('$(sparkle)'); + expect(item.text).to.include('Ephemeral'); + }); + + test('Should have left alignment', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.alignment).to.equal(1); + }); + + test('Should have priority 1000 to appear before all other items', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.priority).to.equal(1000); + }); + + test('Should not have a command', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.command).to.be.undefined; + }); + }); + + suite('Tooltip', () => { + test('Should include auto-generated description in tooltip', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.tooltip).to.include('Auto-generated ephemeral block'); + }); + + test('Should include agent source block ID in tooltip when present', () => { + const cell = createMockCell({ + metadata: { + is_ephemeral: true, + agent_source_block_id: 'a0000000000000000000000000000004' + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.tooltip).to.include('a0000000000000000000000000000004'); + expect(item.tooltip).to.include('Source agent block'); + }); + + test('Should not include source block line in tooltip when agent_source_block_id is absent', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.tooltip).to.not.include('Source agent block'); + }); + }); + + suite('Coexistence with other cell types', () => { + test('Should return item for ephemeral agent cell', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + is_ephemeral: true, + agent_source_block_id: 'source-id' + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + + test('Should return item for ephemeral code cell', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'code' }, + is_ephemeral: true + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + + test('Should return item for ephemeral markdown cell', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'markdown' }, + is_ephemeral: true + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + }); +}); diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index eb87ce58a3..ceadef4d78 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -94,7 +94,10 @@ import { DeepnoteExtensionSidecarWriter } from '../kernels/deepnote/environments import { DeepnoteNotebookEnvironmentMapper } from '../kernels/deepnote/environments/deepnoteNotebookEnvironmentMapper.node'; import { DeepnoteNotebookCommandListener } from './deepnote/deepnoteNotebookCommandListener'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; +import { AgentCellStatusBarProvider } from './deepnote/agentCellStatusBarProvider'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; +import { EphemeralCellDecorationProvider } from './deepnote/ephemeralCellDecorationProvider'; +import { EphemeralCellStatusBarProvider } from './deepnote/ephemeralCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIntegrationStartupCodeProvider'; import { DeepnoteCellCopyHandler } from './deepnote/deepnoteCellCopyHandler'; @@ -261,6 +264,18 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellDecorationProvider + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNewCellLanguageService diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index 28937d6b60..1ca01e8512 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/src/notebooks/serviceRegistry.web.ts @@ -50,7 +50,10 @@ import { IIntegrationWebviewProvider } from './deepnote/integrations/types'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; +import { AgentCellStatusBarProvider } from './deepnote/agentCellStatusBarProvider'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; +import { EphemeralCellDecorationProvider } from './deepnote/ephemeralCellDecorationProvider'; +import { EphemeralCellStatusBarProvider } from './deepnote/ephemeralCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler'; @@ -127,6 +130,18 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellDecorationProvider + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNewCellLanguageService diff --git a/src/renderers/client/markdown.ts b/src/renderers/client/markdown.ts index b5399de2df..f18a4392df 100644 --- a/src/renderers/client/markdown.ts +++ b/src/renderers/client/markdown.ts @@ -1,3 +1,31 @@ +import type { ActivationFunction } from 'vscode-notebook-renderer'; + +// markdown-it ships no type declarations and is only a transitive dependency, so describe the +// small surface this renderer touches rather than depending on its internals wholesale. +interface MarkdownItToken { + content: string; +} + +interface MarkdownItRuleState { + Token: new (type: string, tag: string, nesting: number) => MarkdownItToken; + env?: { + outputItem?: { + metadata?: Record; + }; + }; + tokens: MarkdownItToken[]; +} + +interface MarkdownIt { + core: { + ruler: { + push(name: string, rule: (state: MarkdownItRuleState) => void): void; + }; + }; +} + +type ExtendMarkdownIt = (callback: (md: MarkdownIt) => void) => void; + const styleContent = ` .alert { width: auto; @@ -31,13 +59,61 @@ 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'); + // RendererApi exposes extension hooks through an index signature, so extendMarkdownIt arrives + // as unknown and has to be narrowed before it can be called. + const extendMarkdownIt = markdownRenderer?.extendMarkdownIt as ExtendMarkdownIt | undefined; + + if (typeof extendMarkdownIt === 'function') { + extendMarkdownIt((md) => { + addEphemeralCellWrapper(md); + }); + } + + 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..9470a8a8ce 100644 --- a/src/test/mocks/deepnoteRuntimeCore.ts +++ b/src/test/mocks/deepnoteRuntimeCore.ts @@ -1,9 +1,11 @@ -import type { ServerInfo, ServerOptions } from '@deepnote/runtime-core'; +import type { AgentBlockContext, ServerInfo, ServerOptions } from '@deepnote/runtime-core'; +import type { AgentBlock } from '@deepnote/blocks'; import type { ChildProcess } from 'child_process'; /** * Mock of @deepnote/runtime-core for unit tests: the real startServer/stopServer spawn and - * kill Python processes, so this records calls and returns fake server info instead. + * kill Python processes, and the real executeAgentBlock calls the OpenAI API, so this records + * calls and returns fake results instead. * * build/mocha-esm-loader.js resolves the '@deepnote/runtime-core' specifier to this module, * so code under test and tests importing the __ helpers below share one module instance. @@ -12,6 +14,8 @@ import type { ChildProcess } from 'child_process'; type RuntimeCore = typeof import('@deepnote/runtime-core'); +const executeAgentBlockCalls: { block: AgentBlock; context: AgentBlockContext }[] = []; +const serializeNotebookContextFromBlocksCalls: { blockCount: number; notebookName: string }[] = []; const startServerCalls: ServerOptions[] = []; const stopServerCalls: ServerInfo[] = []; let nextServerId = 0; @@ -27,6 +31,21 @@ function makeFakeProcess(id: number): ChildProcess { } as unknown as ChildProcess; } +export const executeAgentBlock: RuntimeCore['executeAgentBlock'] = async (block, context) => { + executeAgentBlockCalls.push({ block, context }); + + return { finalOutput: '' }; +}; + +export const serializeNotebookContextFromBlocks: RuntimeCore['serializeNotebookContextFromBlocks'] = ({ + blocks, + notebookName +}) => { + serializeNotebookContextFromBlocksCalls.push({ blockCount: blocks.length, notebookName }); + + return `notebook:${notebookName} blocks:${blocks.length}`; +}; + export const startServer: RuntimeCore['startServer'] = async (options) => { startServerCalls.push(options); @@ -49,6 +68,14 @@ export const stopServer: RuntimeCore['stopServer'] = async (info) => { }; // Test-only helpers (prefixed with __ to signal they are not part of the real API). +export function __getExecuteAgentBlockCalls(): { block: AgentBlock; context: AgentBlockContext }[] { + return executeAgentBlockCalls; +} + +export function __getSerializeNotebookContextFromBlocksCalls(): { blockCount: number; notebookName: string }[] { + return serializeNotebookContextFromBlocksCalls; +} + export function __getStartServerCalls(): ServerOptions[] { return startServerCalls; } @@ -62,6 +89,8 @@ export function __setStartServerImpl(impl: RuntimeCore['startServer'] | null): v } export function __resetRuntimeCoreMock(): void { + executeAgentBlockCalls.length = 0; + serializeNotebookContextFromBlocksCalls.length = 0; startServerCalls.length = 0; stopServerCalls.length = 0; nextServerId = 0;