From ff6d44f4ee6715cb67cf37350d233374563469e9 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 11 Mar 2026 15:28:12 +0000 Subject: [PATCH 1/8] feat: Add Agent block visualization support --- .../deepnote/agentCellStatusBarProvider.ts | 249 +++++++++++++++++ .../agentCellStatusBarProvider.unit.test.ts | 251 ++++++++++++++++++ .../converters/agentBlockConverter.ts | 34 +++ .../agentBlockConverter.unit.test.ts | 198 ++++++++++++++ .../deepnote/deepnoteDataConverter.ts | 2 + .../ephemeralCellDecorationProvider.ts | 123 +++++++++ .../ephemeralCellStatusBarProvider.ts | 83 ++++++ ...phemeralCellStatusBarProvider.unit.test.ts | 169 ++++++++++++ src/notebooks/serviceRegistry.node.ts | 15 ++ src/notebooks/serviceRegistry.web.ts | 15 ++ src/renderers/client/markdown.ts | 48 +++- 11 files changed, 1186 insertions(+), 1 deletion(-) create mode 100644 src/notebooks/deepnote/agentCellStatusBarProvider.ts create mode 100644 src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts create mode 100644 src/notebooks/deepnote/converters/agentBlockConverter.ts create mode 100644 src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts create mode 100644 src/notebooks/deepnote/ephemeralCellDecorationProvider.ts create mode 100644 src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts create mode 100644 src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts new file mode 100644 index 0000000000..8d4ba8eba4 --- /dev/null +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -0,0 +1,249 @@ +import { + CancellationToken, + Disposable, + EventEmitter, + NotebookCell, + NotebookCellStatusBarItem, + NotebookCellStatusBarItemProvider, + NotebookEdit, + WorkspaceEdit, + commands, + l10n, + notebooks, + window, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; +import type { Pocket } from '../../platform/deepnote/pocket'; + +const DEFAULT_MAX_ITERATIONS = 20; +const MIN_ITERATIONS = 1; +const MAX_ITERATIONS = 100; +const AGENT_MODEL_OPTIONS = ['auto', 'gpt-4o', 'sonnet']; + +/** + * Provides status bar items for agent cells showing the block type indicator, + * AI model picker, and max iterations setting. + */ +@injectable() +export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService { + private readonly disposables: Disposable[] = []; + private readonly _onDidChangeCellStatusBarItems = new EventEmitter(); + + public readonly onDidChangeCellStatusBarItems = this._onDidChangeCellStatusBarItems.event; + + public activate(): void { + this.disposables.push(notebooks.registerNotebookCellStatusBarItemProvider('deepnote', this)); + + this.disposables.push( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook.notebookType === 'deepnote') { + this._onDidChangeCellStatusBarItems.fire(); + } + }) + ); + + this.disposables.push( + commands.registerCommand('deepnote.switchAgentModel', async (cell?: NotebookCell) => { + const activeCell = cell || this.getActiveCell(); + if (activeCell) { + await this.switchModel(activeCell); + } + }) + ); + + this.disposables.push( + commands.registerCommand('deepnote.setAgentMaxIterations', async (cell?: NotebookCell) => { + const activeCell = cell || this.getActiveCell(); + if (activeCell) { + await this.setMaxIterations(activeCell); + } + }) + ); + + this.disposables.push(this._onDidChangeCellStatusBarItems); + } + + public dispose(): void { + this.disposables.forEach((d) => d.dispose()); + } + + public provideCellStatusBarItems( + cell: NotebookCell, + token: CancellationToken + ): NotebookCellStatusBarItem[] | undefined { + if (token.isCancellationRequested) { + return undefined; + } + + if (!this.isAgentCell(cell)) { + return undefined; + } + + const metadata = cell.metadata as Record | undefined; + const model = this.getModel(metadata); + const maxIterations = this.getMaxIterations(metadata); + + return [ + this.createAgentIndicatorItem(), + this.createModelPickerItem(cell, model), + this.createMaxIterationsItem(cell, maxIterations) + ]; + } + + private createAgentIndicatorItem(): NotebookCellStatusBarItem { + return { + text: `$(hubot) ${l10n.t('Agent Block')}`, + alignment: 1, + priority: 100, + tooltip: l10n.t('Deepnote Agent Block\nAI-powered block that autonomously generates code and analysis') + }; + } + + private createMaxIterationsItem(cell: NotebookCell, maxIterations: number): NotebookCellStatusBarItem { + return { + text: l10n.t('$(iterations) Max iterations: {0}', maxIterations), + alignment: 1, + priority: 80, + tooltip: l10n.t('Maximum iterations for agent\nClick to change'), + command: { + title: l10n.t('Set Max Iterations'), + command: 'deepnote.setAgentMaxIterations', + arguments: [cell] + } + }; + } + + private createModelPickerItem(cell: NotebookCell, model: string): NotebookCellStatusBarItem { + return { + text: `$(symbol-enum) ${l10n.t('Model: {0}', model)}`, + alignment: 1, + priority: 90, + tooltip: l10n.t('AI Model: {0}\nClick to change', model), + command: { + title: l10n.t('Switch Model'), + command: 'deepnote.switchAgentModel', + arguments: [cell] + } + }; + } + + private getActiveCell(): NotebookCell | undefined { + const activeEditor = window.activeNotebookEditor; + if (activeEditor && activeEditor.selection) { + return activeEditor.notebook.cellAt(activeEditor.selection.start); + } + + return undefined; + } + + private getMaxIterations(metadata: Record | undefined): number { + const value = metadata?.deepnote_max_iterations; + if (typeof value === 'number' && Number.isInteger(value) && value >= MIN_ITERATIONS) { + return value; + } + + return DEFAULT_MAX_ITERATIONS; + } + + private getModel(metadata: Record | undefined): string { + const value = metadata?.deepnote_model; + if (typeof value === 'string' && value) { + return value; + } + + return 'auto'; + } + + private isAgentCell(cell: NotebookCell): boolean { + const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; + + return pocket?.type === 'agent'; + } + + private async setMaxIterations(cell: NotebookCell): Promise { + if (!this.isAgentCell(cell)) { + return; + } + + const metadata = cell.metadata as Record | undefined; + const currentValue = this.getMaxIterations(metadata); + + const input = await window.showInputBox({ + prompt: l10n.t('Enter maximum number of iterations ({0}-{1})', MIN_ITERATIONS, MAX_ITERATIONS), + value: String(currentValue), + validateInput: (value) => { + const num = parseInt(value, 10); + if (isNaN(num) || !Number.isInteger(num)) { + return l10n.t('Please enter a whole number'); + } + if (num < MIN_ITERATIONS || num > MAX_ITERATIONS) { + return l10n.t('Value must be between {0} and {1}', MIN_ITERATIONS, MAX_ITERATIONS); + } + + return undefined; + } + }); + + if (input === undefined) { + return; + } + + const newValue = parseInt(input, 10); + if (newValue === currentValue) { + return; + } + + await this.updateCellMetadata(cell, { deepnote_max_iterations: newValue }); + } + + private async switchModel(cell: NotebookCell): Promise { + if (!this.isAgentCell(cell)) { + return; + } + + const metadata = cell.metadata as Record | undefined; + const currentModel = this.getModel(metadata); + + const items = AGENT_MODEL_OPTIONS.map((option) => ({ + label: option, + description: option === currentModel ? l10n.t('Currently selected') : undefined + })); + + const selected = await window.showQuickPick(items, { + placeHolder: l10n.t('Select AI model for agent') + }); + + if (!selected || selected.label === currentModel) { + return; + } + + const newModel = selected.label === 'auto' ? undefined : selected.label; + + await this.updateCellMetadata(cell, { deepnote_model: newModel }); + } + + private async updateCellMetadata(cell: NotebookCell, updates: Record): Promise { + const updatedMetadata = { ...cell.metadata, ...updates }; + + // Remove keys set to undefined so they don't persist + for (const [key, value] of Object.entries(updates)) { + if (value === undefined) { + delete updatedMetadata[key]; + } + } + + const edit = new WorkspaceEdit(); + edit.set(cell.notebook.uri, [NotebookEdit.updateCellMetadata(cell.index, updatedMetadata)]); + + const success = await workspace.applyEdit(edit); + if (!success) { + void window.showErrorMessage(l10n.t('Failed to update agent cell metadata')); + return; + } + + this._onDidChangeCellStatusBarItems.fire(); + } +} diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts new file mode 100644 index 0000000000..e5397a1948 --- /dev/null +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -0,0 +1,251 @@ +import { expect } from 'chai'; +import { CancellationToken } from 'vscode'; + +import { AgentCellStatusBarProvider } from './agentCellStatusBarProvider'; +import { createMockCell } from './deepnoteTestHelpers'; + +suite('AgentCellStatusBarProvider', () => { + let provider: AgentCellStatusBarProvider; + let mockToken: CancellationToken; + + setup(() => { + mockToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + provider = new AgentCellStatusBarProvider(); + }); + + teardown(() => { + provider.dispose(); + }); + + suite('Agent Cell Detection', () => { + test('Should return status bar items for agent cell', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.not.be.undefined; + expect(items).to.have.lengthOf(3); + }); + + test('Should return undefined for code cell', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined for sql cell', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'sql' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined for markdown cell', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'markdown' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined for cell without pocket', () => { + const cell = createMockCell({ metadata: {} }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined for cell without metadata', () => { + const cell = createMockCell({ metadata: undefined }); + const items = provider.provideCellStatusBarItems(cell, mockToken); + + expect(items).to.be.undefined; + }); + + test('Should return undefined when cancellation is requested', () => { + const cancelledToken: CancellationToken = { + isCancellationRequested: true, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, cancelledToken); + + expect(items).to.be.undefined; + }); + }); + + suite('Agent Block Indicator', () => { + test('Should display agent block label with icon', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[0].text).to.include('$(hubot)'); + expect(items[0].text).to.include('Agent Block'); + expect(items[0].alignment).to.equal(1); + expect(items[0].priority).to.equal(100); + }); + + test('Should not have a command on the indicator', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[0].command).to.be.undefined; + }); + }); + + suite('Model Picker', () => { + test('Should display "auto" when no model is set', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: auto'); + expect(items[1].text).to.include('$(symbol-enum)'); + }); + + test('Should display configured model from metadata', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_model: 'gpt-4o' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: gpt-4o'); + }); + + test('Should display sonnet model', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_model: 'sonnet' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: sonnet'); + }); + + test('Should display "auto" when model is empty string', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_model: '' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].text).to.include('Model: auto'); + }); + + test('Should have switch model command', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].command).to.not.be.undefined; + const cmd = items[1].command as any; + expect(cmd.command).to.equal('deepnote.switchAgentModel'); + }); + + test('Should have priority 90', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[1].priority).to.equal(90); + }); + }); + + suite('Max Iterations', () => { + test('Should display default max iterations (20) when not set', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + expect(items[2].text).to.include('$(iterations)'); + }); + + test('Should display configured max iterations from metadata', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: 10 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 10'); + }); + + test('Should display default when max iterations is not a number', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: 'invalid' + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + }); + + test('Should display default when max iterations is zero', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: 0 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + }); + + test('Should display default when max iterations is a float', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: 5.5 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + }); + + test('Should have set max iterations command', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].command).to.not.be.undefined; + const cmd = items[2].command as any; + expect(cmd.command).to.equal('deepnote.setAgentMaxIterations'); + }); + + test('Should have priority 80', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].priority).to.equal(80); + }); + }); + + suite('Combined metadata', () => { + test('Should display both model and max iterations from metadata', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_model: 'gpt-4o', + deepnote_max_iterations: 50 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items).to.have.lengthOf(3); + expect(items[0].text).to.include('Agent Block'); + expect(items[1].text).to.include('Model: gpt-4o'); + expect(items[2].text).to.include('Max iterations: 50'); + }); + }); +}); diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.ts b/src/notebooks/deepnote/converters/agentBlockConverter.ts new file mode 100644 index 0000000000..357ab5d3c9 --- /dev/null +++ b/src/notebooks/deepnote/converters/agentBlockConverter.ts @@ -0,0 +1,34 @@ +import type { DeepnoteBlock } from '@deepnote/blocks'; +import { NotebookCellData, NotebookCellKind } from 'vscode'; + +import type { BlockConverter } from './blockConverter'; + +/** + * Converter for agent blocks. + * + * Agent blocks are rendered as code cells with markdown language so the natural-language + * prompt gets reasonable syntax highlighting while remaining visually distinct from + * Python code blocks. The prompt text is stored in `block.content`. + * + * Agent-specific metadata (model, MCP servers, max iterations, etc.) is preserved + * through the generic metadata pass-through in DeepnoteDataConverter. + */ +export class AgentBlockConverter implements BlockConverter { + applyChangesToBlock(block: DeepnoteBlock, cell: NotebookCellData): void { + block.content = cell.value || ''; + } + + canConvert(blockType: string): boolean { + return blockType.toLowerCase() === 'agent'; + } + + convertToCell(block: DeepnoteBlock): NotebookCellData { + const cell = new NotebookCellData(NotebookCellKind.Code, block.content || '', 'markdown'); + + return cell; + } + + getSupportedTypes(): string[] { + return ['agent']; + } +} diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts new file mode 100644 index 0000000000..43fac425db --- /dev/null +++ b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts @@ -0,0 +1,198 @@ +import type { DeepnoteBlock } from '@deepnote/blocks'; +import { assert } from 'chai'; +import { NotebookCellData, NotebookCellKind } from 'vscode'; +import { AgentBlockConverter } from './agentBlockConverter'; +import dedent from 'dedent'; + +suite('AgentBlockConverter', () => { + let converter: AgentBlockConverter; + + setup(() => { + converter = new AgentBlockConverter(); + }); + + suite('canConvert', () => { + test('returns true for "agent" type', () => { + assert.strictEqual(converter.canConvert('agent'), true); + }); + + test('returns true for "Agent" type (case insensitive)', () => { + assert.strictEqual(converter.canConvert('Agent'), true); + }); + + test('returns false for other types', () => { + assert.strictEqual(converter.canConvert('code'), false); + assert.strictEqual(converter.canConvert('markdown'), false); + assert.strictEqual(converter.canConvert('sql'), false); + }); + }); + + suite('getSupportedTypes', () => { + test('returns array with "agent"', () => { + const types = converter.getSupportedTypes(); + + assert.deepStrictEqual(types, ['agent']); + }); + }); + + suite('convertToCell', () => { + test('converts agent block to code cell with markdown language', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Analyze the dataset and create a summary report', + id: 'agent-block-123', + sortingKey: 'a0', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, 'Analyze the dataset and create a summary report'); + assert.strictEqual(cell.languageId, 'markdown'); + }); + + test('handles empty content', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: '', + id: 'agent-block-456', + sortingKey: 'a1', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, ''); + assert.strictEqual(cell.languageId, 'markdown'); + }); + + test('handles undefined content', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + id: 'agent-block-789', + sortingKey: 'a2', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, ''); + assert.strictEqual(cell.languageId, 'markdown'); + }); + + test('preserves multiline prompt', () => { + const prompt = dedent` + You are a senior data analyst. + + Perform a thorough exploratory analysis: + 1. Create a grouped bar chart of revenue by quarter + 2. Create a line chart showing churn rate trends + 3. Compute a pivot table of average revenue + `; + + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: prompt, + id: 'agent-block-multiline', + sortingKey: 'a3', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, prompt); + assert.strictEqual(cell.languageId, 'markdown'); + }); + + test('preserves agent block with metadata', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Analyze the data', + id: 'agent-block-with-metadata', + metadata: { + deepnote_agent_model: 'gpt-4o' + }, + sortingKey: 'a4', + type: 'agent' + }; + + const cell = converter.convertToCell(block); + + assert.strictEqual(cell.kind, NotebookCellKind.Code); + assert.strictEqual(cell.value, 'Analyze the data'); + assert.strictEqual(cell.languageId, 'markdown'); + }); + }); + + suite('applyChangesToBlock', () => { + test('updates block content from cell value', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Old prompt', + id: 'agent-block-123', + sortingKey: 'a0', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + const cell = new NotebookCellData( + NotebookCellKind.Code, + 'New prompt with updated instructions', + 'markdown' + ); + + converter.applyChangesToBlock(block, cell); + + assert.strictEqual(block.content, 'New prompt with updated instructions'); + }); + + test('handles empty cell value', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Some prompt', + id: 'agent-block-456', + sortingKey: 'a1', + metadata: { deepnote_agent_model: 'auto' }, + type: 'agent' + }; + const cell = new NotebookCellData(NotebookCellKind.Code, '', 'markdown'); + + converter.applyChangesToBlock(block, cell); + + assert.strictEqual(block.content, ''); + }); + + test('does not modify other block properties', () => { + const block: DeepnoteBlock = { + blockGroup: 'test-group', + content: 'Old prompt', + id: 'agent-block-789', + metadata: { + deepnote_agent_model: 'gpt-4o', + custom: 'value' + }, + sortingKey: 'a2', + type: 'agent' + }; + const cell = new NotebookCellData(NotebookCellKind.Code, 'New prompt', 'markdown'); + + converter.applyChangesToBlock(block, cell); + + assert.strictEqual(block.content, 'New prompt'); + assert.strictEqual(block.id, 'agent-block-789'); + assert.strictEqual(block.type, 'agent'); + assert.strictEqual(block.sortingKey, 'a2'); + assert.deepStrictEqual(block.metadata, { + deepnote_agent_model: 'gpt-4o', + custom: 'value' + }); + }); + }); +}); diff --git a/src/notebooks/deepnote/deepnoteDataConverter.ts b/src/notebooks/deepnote/deepnoteDataConverter.ts index 9f71700a71..51007cae9d 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.ts @@ -11,6 +11,7 @@ import { MarkdownBlockConverter } from './converters/markdownBlockConverter'; import { VisualizationBlockConverter } from './converters/visualizationBlockConverter'; import { compile as convertVegaLiteSpecToVega, ensureVegaLiteLoaded } from './vegaLiteWrapper'; import { produce } from 'immer'; +import { AgentBlockConverter } from './converters/agentBlockConverter'; import { SqlBlockConverter } from './converters/sqlBlockConverter'; import { TextBlockConverter } from './converters/textBlockConverter'; // @ts-ignore - types_unstable subpath requires moduleResolution: "node16" which mandates module: "node16" and .js extensions on all imports @@ -38,6 +39,7 @@ export class DeepnoteDataConverter { private readonly registry = new ConverterRegistry(); constructor() { + this.registry.register(new AgentBlockConverter()); this.registry.register(new CodeBlockConverter()); this.registry.register(new MarkdownBlockConverter()); this.registry.register(new ChartBigNumberBlockConverter()); diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts new file mode 100644 index 0000000000..5c913700ff --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts @@ -0,0 +1,123 @@ +import { + Disposable, + NotebookCell, + NotebookDocument, + OverviewRulerLane, + Range, + TextEditor, + TextEditorDecorationType, + ThemeColor, + window, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; + +const NOTEBOOK_CELL_SCHEME = 'vscode-notebook-cell'; + +/** + * Applies visual decorations (left border, background tint, reduced opacity) to + * code cell editors that belong to ephemeral blocks (`is_ephemeral: true`). + * + * The left border is rendered via a `before` pseudo-element on each line, + * which avoids overlapping or shifting the code text. + * + * Markup cells are handled separately by the markdown-it renderer plugin in + * `src/renderers/client/markdown.ts`. + */ +@injectable() +export class EphemeralCellDecorationProvider implements IExtensionSyncActivationService { + private readonly disposables: Disposable[] = []; + + private ephemeralDecorationType!: TextEditorDecorationType; + + public activate(): void { + this.ephemeralDecorationType = window.createTextEditorDecorationType({ + opacity: '0.8', + isWholeLine: true, + overviewRulerColor: new ThemeColor('charts.yellow'), + overviewRulerLane: OverviewRulerLane.Left, + before: { + contentText: '\u200B', + width: '3px', + backgroundColor: new ThemeColor('charts.yellow'), + margin: '0 8px 0 0' + } + }); + + this.disposables.push(this.ephemeralDecorationType); + + this.disposables.push( + window.onDidChangeVisibleTextEditors(() => { + this.updateDecorations(); + }) + ); + + this.disposables.push( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook.notebookType === 'deepnote') { + this.updateDecorations(); + } + }) + ); + + this.updateDecorations(); + } + + public dispose(): void { + this.disposables.forEach((d) => d.dispose()); + } + + private findCellForEditor(editor: TextEditor): NotebookCell | undefined { + const uri = editor.document.uri; + if (uri.scheme !== NOTEBOOK_CELL_SCHEME) { + return undefined; + } + + for (const notebook of workspace.notebookDocuments) { + if (notebook.notebookType !== 'deepnote') { + continue; + } + + const cell = this.findMatchingCell(notebook, editor); + if (cell) { + return cell; + } + } + + return undefined; + } + + private findMatchingCell(notebook: NotebookDocument, editor: TextEditor): NotebookCell | undefined { + for (const cell of notebook.getCells()) { + if (cell.document.uri.toString() === editor.document.uri.toString()) { + return cell; + } + } + + return undefined; + } + + private updateDecorations(): void { + for (const editor of window.visibleTextEditors) { + if (editor.document.uri.scheme !== NOTEBOOK_CELL_SCHEME) { + continue; + } + + const cell = this.findCellForEditor(editor); + if (!cell || cell.metadata?.is_ephemeral !== true) { + editor.setDecorations(this.ephemeralDecorationType, []); + continue; + } + + const lineRanges: Range[] = []; + for (let i = 0; i < editor.document.lineCount; i++) { + const line = editor.document.lineAt(i); + lineRanges.push(line.range); + } + + editor.setDecorations(this.ephemeralDecorationType, lineRanges); + } + } +} diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts new file mode 100644 index 0000000000..7089391e7c --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts @@ -0,0 +1,83 @@ +import { + CancellationToken, + Disposable, + EventEmitter, + NotebookCell, + NotebookCellStatusBarItem, + NotebookCellStatusBarItemProvider, + l10n, + notebooks, + workspace +} from 'vscode'; +import { injectable } from 'inversify'; + +import { IExtensionSyncActivationService } from '../../platform/activation/types'; + +const EPHEMERAL_INDICATOR_PRIORITY = 1000; + +/** + * Provides a status bar indicator for ephemeral cells — blocks that were + * auto-generated by an agent and marked with `is_ephemeral: true` in metadata. + */ +@injectable() +export class EphemeralCellStatusBarProvider + implements NotebookCellStatusBarItemProvider, IExtensionSyncActivationService +{ + private readonly disposables: Disposable[] = []; + private readonly _onDidChangeCellStatusBarItems = new EventEmitter(); + + public readonly onDidChangeCellStatusBarItems = this._onDidChangeCellStatusBarItems.event; + + public activate(): void { + this.disposables.push(notebooks.registerNotebookCellStatusBarItemProvider('deepnote', this)); + + this.disposables.push( + workspace.onDidChangeNotebookDocument((e) => { + if (e.notebook.notebookType === 'deepnote') { + this._onDidChangeCellStatusBarItems.fire(); + } + }) + ); + + this.disposables.push(this._onDidChangeCellStatusBarItems); + } + + public dispose(): void { + this.disposables.forEach((d) => d.dispose()); + } + + public provideCellStatusBarItems( + cell: NotebookCell, + token: CancellationToken + ): NotebookCellStatusBarItem | undefined { + if (token.isCancellationRequested) { + return undefined; + } + + if (!this.isEphemeralCell(cell)) { + return undefined; + } + + const agentSourceBlockId = cell.metadata?.agent_source_block_id as string | undefined; + + return this.createEphemeralIndicatorItem(agentSourceBlockId); + } + + private createEphemeralIndicatorItem(agentSourceBlockId?: string): NotebookCellStatusBarItem { + const tooltipLines = [l10n.t('Auto-generated ephemeral block')]; + if (agentSourceBlockId) { + tooltipLines.push(l10n.t('Source agent block: {0}', agentSourceBlockId)); + } + + return { + text: `$(sparkle) ${l10n.t('Ephemeral')}`, + alignment: 1, + priority: EPHEMERAL_INDICATOR_PRIORITY, + tooltip: tooltipLines.join('\n') + }; + } + + private isEphemeralCell(cell: NotebookCell): boolean { + return cell.metadata?.is_ephemeral === true; + } +} diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts new file mode 100644 index 0000000000..27e53dcdf2 --- /dev/null +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.unit.test.ts @@ -0,0 +1,169 @@ +import { expect } from 'chai'; +import { CancellationToken } from 'vscode'; + +import { EphemeralCellStatusBarProvider } from './ephemeralCellStatusBarProvider'; +import { createMockCell } from './deepnoteTestHelpers'; + +suite('EphemeralCellStatusBarProvider', () => { + let provider: EphemeralCellStatusBarProvider; + let mockToken: CancellationToken; + + setup(() => { + mockToken = { + isCancellationRequested: false, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + provider = new EphemeralCellStatusBarProvider(); + }); + + teardown(() => { + provider.dispose(); + }); + + suite('Ephemeral Cell Detection', () => { + test('Should return a status bar item for ephemeral cell', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + + test('Should return undefined when is_ephemeral is false', () => { + const cell = createMockCell({ metadata: { is_ephemeral: false } }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined when is_ephemeral is not set', () => { + const cell = createMockCell({ metadata: {} }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined for cell without metadata', () => { + const cell = createMockCell({ metadata: undefined }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined when is_ephemeral is a non-boolean truthy value', () => { + const cell = createMockCell({ metadata: { is_ephemeral: 'true' } }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.be.undefined; + }); + + test('Should return undefined when cancellation is requested', () => { + const cancelledToken: CancellationToken = { + isCancellationRequested: true, + onCancellationRequested: () => ({ dispose: () => undefined }) + } as any; + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, cancelledToken); + + expect(item).to.be.undefined; + }); + }); + + suite('Status Bar Item Properties', () => { + test('Should display sparkle icon with Ephemeral label', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.text).to.include('$(sparkle)'); + expect(item.text).to.include('Ephemeral'); + }); + + test('Should have left alignment', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.alignment).to.equal(1); + }); + + test('Should have priority 1000 to appear before all other items', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.priority).to.equal(1000); + }); + + test('Should not have a command', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.command).to.be.undefined; + }); + }); + + suite('Tooltip', () => { + test('Should include auto-generated description in tooltip', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.tooltip).to.include('Auto-generated ephemeral block'); + }); + + test('Should include agent source block ID in tooltip when present', () => { + const cell = createMockCell({ + metadata: { + is_ephemeral: true, + agent_source_block_id: 'a0000000000000000000000000000004' + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.tooltip).to.include('a0000000000000000000000000000004'); + expect(item.tooltip).to.include('Source agent block'); + }); + + test('Should not include source block line in tooltip when agent_source_block_id is absent', () => { + const cell = createMockCell({ metadata: { is_ephemeral: true } }); + const item = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(item.tooltip).to.not.include('Source agent block'); + }); + }); + + suite('Coexistence with other cell types', () => { + test('Should return item for ephemeral agent cell', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + is_ephemeral: true, + agent_source_block_id: 'source-id' + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + + test('Should return item for ephemeral code cell', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'code' }, + is_ephemeral: true + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + + test('Should return item for ephemeral markdown cell', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'markdown' }, + is_ephemeral: true + } + }); + const item = provider.provideCellStatusBarItems(cell, mockToken); + + expect(item).to.not.be.undefined; + }); + }); +}); diff --git a/src/notebooks/serviceRegistry.node.ts b/src/notebooks/serviceRegistry.node.ts index cbd8b860fe..5de4c2c4d9 100644 --- a/src/notebooks/serviceRegistry.node.ts +++ b/src/notebooks/serviceRegistry.node.ts @@ -85,7 +85,10 @@ import { DeepnoteExtensionSidecarWriter } from '../kernels/deepnote/environments import { DeepnoteNotebookEnvironmentMapper } from '../kernels/deepnote/environments/deepnoteNotebookEnvironmentMapper.node'; import { DeepnoteNotebookCommandListener } from './deepnote/deepnoteNotebookCommandListener'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; +import { AgentCellStatusBarProvider } from './deepnote/agentCellStatusBarProvider'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; +import { EphemeralCellDecorationProvider } from './deepnote/ephemeralCellDecorationProvider'; +import { EphemeralCellStatusBarProvider } from './deepnote/ephemeralCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIntegrationStartupCodeProvider'; import { DeepnoteCellCopyHandler } from './deepnote/deepnoteCellCopyHandler'; @@ -230,6 +233,18 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellDecorationProvider + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNewCellLanguageService diff --git a/src/notebooks/serviceRegistry.web.ts b/src/notebooks/serviceRegistry.web.ts index 2488ff73d7..4be669266c 100644 --- a/src/notebooks/serviceRegistry.web.ts +++ b/src/notebooks/serviceRegistry.web.ts @@ -50,7 +50,10 @@ import { IIntegrationWebviewProvider } from './deepnote/integrations/types'; import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider'; +import { AgentCellStatusBarProvider } from './deepnote/agentCellStatusBarProvider'; import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider'; +import { EphemeralCellDecorationProvider } from './deepnote/ephemeralCellDecorationProvider'; +import { EphemeralCellStatusBarProvider } from './deepnote/ephemeralCellStatusBarProvider'; import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService'; import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider'; import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler'; @@ -125,6 +128,18 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea IExtensionSyncActivationService, DeepnoteBigNumberCellStatusBarProvider ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + AgentCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellStatusBarProvider + ); + serviceManager.addSingleton( + IExtensionSyncActivationService, + EphemeralCellDecorationProvider + ); serviceManager.addSingleton( IExtensionSyncActivationService, DeepnoteNewCellLanguageService diff --git a/src/renderers/client/markdown.ts b/src/renderers/client/markdown.ts index b5399de2df..8c69bfd618 100644 --- a/src/renderers/client/markdown.ts +++ b/src/renderers/client/markdown.ts @@ -1,3 +1,5 @@ +import type { ActivationFunction } from 'vscode-notebook-renderer'; + const styleContent = ` .alert { width: auto; @@ -31,13 +33,57 @@ const styleContent = ` background-color: rgb(255,205,210); color: rgb(183,28,28); } + +.ephemeral-cell { + border-left: 3px solid var(--vscode-charts-yellow, #cca700); + padding-left: 8px; + opacity: 0.8; +} +.ephemeral-badge { + display: inline-block; + font-size: 0.75em; + padding: 1px 6px; + border-radius: 3px; + background: var(--vscode-charts-yellow, #cca700); + color: var(--vscode-editor-background, #1e1e1e); + margin-bottom: 4px; + font-weight: 600; + letter-spacing: 0.03em; +} `; -export async function activate() { +export const activate: ActivationFunction = async (ctx) => { const style = document.createElement('style'); style.textContent = styleContent; const template = document.createElement('template'); template.classList.add('markdown-style'); template.content.appendChild(style); document.head.appendChild(template); + + const markdownRenderer = await ctx.getRenderer('vscode.markdown-it-renderer'); + if (markdownRenderer) { + (markdownRenderer as any).extendMarkdownIt((md: any) => { + addEphemeralCellWrapper(md); + }); + } + + return undefined; +}; + +function addEphemeralCellWrapper(md: any): void { + md.core.ruler.push('ephemeral_wrapper', (state: any) => { + const metadata = state.env?.outputItem?.metadata; + if (!metadata || metadata.is_ephemeral !== true) { + return; + } + + const openToken = new state.Token('html_block', '', 0); + openToken.content = '
\u2728 Ephemeral\n'; + + const closeToken = new state.Token('html_block', '', 0); + closeToken.content = '
\n'; + + state.tokens.unshift(openToken); + state.tokens.push(closeToken); + }); } From 957fdcdc43b55f014fbcc4b21763b716457e57b2 Mon Sep 17 00:00:00 2001 From: tomas Date: Thu, 12 Mar 2026 11:32:19 +0000 Subject: [PATCH 2/8] Add a dummy agent block execution handler --- .../controllers/vscodeNotebookController.ts | 26 +- .../deepnote/agentCellExecutionHandler.ts | 51 ++++ .../agentCellExecutionHandler.unit.test.ts | 257 ++++++++++++++++++ .../converters/agentBlockConverter.ts | 8 +- .../agentBlockConverter.unit.test.ts | 18 +- .../deepnoteKernelAutoSelector.node.ts | 32 ++- 6 files changed, 368 insertions(+), 24 deletions(-) create mode 100644 src/notebooks/deepnote/agentCellExecutionHandler.ts create mode 100644 src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts diff --git a/src/notebooks/controllers/vscodeNotebookController.ts b/src/notebooks/controllers/vscodeNotebookController.ts index eeeb2614e8..8b4f86db78 100644 --- a/src/notebooks/controllers/vscodeNotebookController.ts +++ b/src/notebooks/controllers/vscodeNotebookController.ts @@ -90,6 +90,7 @@ import { RemoteKernelReconnectBusyIndicator } from './remoteKernelReconnectBusyI import { IConnectionDisplayData, IConnectionDisplayDataProvider, IVSCodeNotebookController } from './types'; import { notebookPathToDeepnoteProjectFilePath } from '../../platform/deepnote/deepnoteProjectUtils'; import { DEEPNOTE_NOTEBOOK_TYPE, IDeepnoteKernelAutoSelector } from '../../kernels/deepnote/types'; +import { executeAgentCell, isAgentCell } from '../deepnote/agentCellExecutionHandler'; /** * Our implementation of the VSCode Notebook Controller. Called by VS code to execute cells in a notebook. Also displayed @@ -626,16 +627,29 @@ export class VSCodeNotebookController implements Disposable, IVSCodeNotebookCont // Start execution now (from the user's point of view) // Creating these execution objects marks the cell as queued for execution (vscode will update cell UI). type CellExec = { cell: NotebookCell; exec: NotebookCellExecution }; - const cellExecs: CellExec[] = (this.cellQueue.get(doc) || []).map((cell) => { - const exec = this.createCellExecutionIfNecessary(cell, new KernelController(this.controller)); - return { cell, exec }; - }); + const allCells = this.cellQueue.get(doc) || []; this.cellQueue.delete(doc); - const firstCell = cellExecs.length ? cellExecs[0].cell : undefined; - if (!firstCell) { + + const agentCells = allCells.filter((cell) => isAgentCell(cell)); + const kernelCells = allCells.filter((cell) => !isAgentCell(cell)); + + // Execute agent cells directly without kernel involvement + if (agentCells.length > 0) { + logger.trace(`Executing ${agentCells.length} agent cell(s) for ${getDisplayPath(doc.uri)} without kernel`); + await Promise.all(agentCells.map((cell) => executeAgentCell(cell, this.controller))).catch(noop); + } + + if (kernelCells.length === 0) { return; } + const cellExecs: CellExec[] = kernelCells.map((cell) => { + const exec = this.createCellExecutionIfNecessary(cell, new KernelController(this.controller)); + return { cell, exec }; + }); + + const firstCell = cellExecs[0].cell; + logger.trace(`Execute Notebook ${getDisplayPath(doc.uri)}. Step 1`); // Connect to a matching kernel if possible (but user may pick a different one) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts new file mode 100644 index 0000000000..51ab4b0ce5 --- /dev/null +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -0,0 +1,51 @@ +import { NotebookCell, NotebookCellOutput, NotebookCellOutputItem, NotebookController } from 'vscode'; + +import type { Pocket } from '../../platform/deepnote/pocket'; +import { logger } from '../../platform/logging'; + +export function isAgentCell(cell: NotebookCell): boolean { + const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; + + return pocket?.type === 'agent'; +} + +export async function executeAgentCell(cell: NotebookCell, controller: NotebookController): Promise { + const execution = controller.createNotebookCellExecution(cell); + execution.start(Date.now()); + + try { + await execution.clearOutput(); + const prompt = cell.document.getText(); + + const output = new NotebookCellOutput([ + NotebookCellOutputItem.text(`[Agent] Received prompt (${prompt.length} chars)...\n`) + ]); + await execution.replaceOutput([output]); + + const chunks = [ + { delay: 500, text: '[Agent] Analyzing prompt...\n' }, + { delay: 1000, text: '[Agent] Generating plan...\n' }, + { delay: 2000, text: '[Agent] Executing steps...\n' }, + { delay: 3000, text: `[Agent] Done.\n\nPrompt: ${prompt}\n` } + ]; + + let accumulated = `[Agent] Received prompt (${prompt.length} chars)...\n`; + for (const chunk of chunks) { + await delay(chunk.delay); + accumulated += chunk.text; + await execution.replaceOutputItems(NotebookCellOutputItem.text(accumulated), output); + } + + execution.end(true, Date.now()); + } catch (error) { + logger.error('Agent cell execution failed', error); + const message = error instanceof Error ? error.message : String(error); + const stderrOutput = new NotebookCellOutput([NotebookCellOutputItem.stderr(message)]); + await execution.replaceOutput([stderrOutput]).then(undefined, () => undefined); + execution.end(false, Date.now()); + } +} + +function delay(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts new file mode 100644 index 0000000000..77544aeafd --- /dev/null +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -0,0 +1,257 @@ +import { expect } from 'chai'; +import * as sinon from 'sinon'; +import { NotebookCellOutput, NotebookCellOutputItem, NotebookController } from 'vscode'; + +import { executeAgentCell, isAgentCell } from './agentCellExecutionHandler'; +import { createMockCell } from './deepnoteTestHelpers'; + +suite('AgentCellExecutionHandler', () => { + suite('isAgentCell', () => { + test('returns true for cell with agent pocket type', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); + + expect(isAgentCell(cell)).to.be.true; + }); + + test('returns false for cell with code pocket type', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'code' } } }); + + expect(isAgentCell(cell)).to.be.false; + }); + + test('returns false for cell with markdown pocket type', () => { + const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'markdown' } } }); + + expect(isAgentCell(cell)).to.be.false; + }); + + test('returns false for cell without pocket', () => { + const cell = createMockCell({ metadata: {} }); + + expect(isAgentCell(cell)).to.be.false; + }); + + test('returns false for cell without metadata', () => { + const cell = createMockCell({ metadata: undefined }); + + expect(isAgentCell(cell)).to.be.false; + }); + }); + + suite('executeAgentCell', () => { + let clock: sinon.SinonFakeTimers; + let mockExecution: { + clearOutput: sinon.SinonStub; + end: sinon.SinonStub; + replaceOutput: sinon.SinonStub; + replaceOutputItems: sinon.SinonStub; + start: sinon.SinonStub; + }; + let mockController: NotebookController; + + setup(() => { + clock = sinon.useFakeTimers(); + + mockExecution = { + clearOutput: sinon.stub().resolves(), + end: sinon.stub(), + replaceOutput: sinon.stub().resolves(), + replaceOutputItems: sinon.stub().resolves(), + start: sinon.stub() + }; + + mockController = { + createNotebookCellExecution: sinon.stub().returns(mockExecution) + } as unknown as NotebookController; + }); + + teardown(() => { + clock.restore(); + }); + + async function runToCompletion(promise: Promise): Promise { + // Total delay across all chunks: 500 + 1000 + 2000 + 3000 = 6500ms + await clock.tickAsync(7000); + await promise; + } + + test('creates execution and starts it', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Analyze data' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect((mockController.createNotebookCellExecution as sinon.SinonStub).calledOnceWith(cell)).to.be.true; + expect(mockExecution.start.calledOnce).to.be.true; + }); + + test('clears output before streaming', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Analyze data' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.clearOutput.calledOnce).to.be.true; + expect(mockExecution.clearOutput.calledBefore(mockExecution.replaceOutput)).to.be.true; + }); + + test('sets initial output via replaceOutput', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Hello world' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.replaceOutput.calledOnce).to.be.true; + + const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; + expect(outputs).to.have.lengthOf(1); + expect(outputs[0].items).to.have.lengthOf(1); + + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('[Agent] Received prompt (11 chars)'); + }); + + test('streams 4 chunks via replaceOutputItems', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test prompt' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.replaceOutputItems.callCount).to.equal(4); + }); + + test('streaming chunks accumulate text progressively', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + const getChunkText = (callIndex: number): string => { + const item = mockExecution.replaceOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; + + return Buffer.from(item.data).toString('utf-8'); + }; + + const chunk1 = getChunkText(0); + const chunk2 = getChunkText(1); + const chunk3 = getChunkText(2); + const chunk4 = getChunkText(3); + + expect(chunk1).to.include('Analyzing prompt'); + expect(chunk2).to.include('Generating plan'); + expect(chunk2).to.include('Analyzing prompt'); + expect(chunk3).to.include('Executing steps'); + expect(chunk3).to.include('Generating plan'); + expect(chunk4).to.include('Done'); + expect(chunk4).to.include('Prompt: Test'); + }); + + test('streaming chunks fire at correct intervals', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test' + }); + + const promise = executeAgentCell(cell, mockController); + + expect(mockExecution.replaceOutputItems.callCount).to.equal(0); + + await clock.tickAsync(500); + expect(mockExecution.replaceOutputItems.callCount).to.equal(1); + + await clock.tickAsync(1000); + expect(mockExecution.replaceOutputItems.callCount).to.equal(2); + + await clock.tickAsync(2000); + expect(mockExecution.replaceOutputItems.callCount).to.equal(3); + + await clock.tickAsync(3000); + expect(mockExecution.replaceOutputItems.callCount).to.equal(4); + + await promise; + }); + + test('ends execution with success', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.true; + }); + + test('ends execution with failure when error occurs', async () => { + mockExecution.clearOutput.rejects(new Error('Test error')); + + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.false; + }); + + test('writes error message to stderr output on failure', async () => { + mockExecution.clearOutput.rejects(new Error('Something went wrong')); + + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: 'Test' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.replaceOutput.calledOnce).to.be.true; + + const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; + expect(outputs).to.have.lengthOf(1); + + const item = outputs[0].items[0]; + expect(item.mime).to.equal('application/vnd.code.notebook.stderr'); + + const text = Buffer.from(item.data).toString('utf-8'); + expect(text).to.equal('Something went wrong'); + }); + + test('handles empty prompt', async () => { + const cell = createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text: '' + }); + + const promise = executeAgentCell(cell, mockController); + await runToCompletion(promise); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.true; + + const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('(0 chars)'); + }); + }); +}); diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.ts b/src/notebooks/deepnote/converters/agentBlockConverter.ts index 357ab5d3c9..6f9ebbd31c 100644 --- a/src/notebooks/deepnote/converters/agentBlockConverter.ts +++ b/src/notebooks/deepnote/converters/agentBlockConverter.ts @@ -6,9 +6,9 @@ import type { BlockConverter } from './blockConverter'; /** * Converter for agent blocks. * - * Agent blocks are rendered as code cells with markdown language so the natural-language - * prompt gets reasonable syntax highlighting while remaining visually distinct from - * Python code blocks. The prompt text is stored in `block.content`. + * Agent blocks are rendered as code cells with plaintext language so the + * natural-language prompt appears without syntax highlighting while remaining + * executable. The prompt text is stored in `block.content`. * * Agent-specific metadata (model, MCP servers, max iterations, etc.) is preserved * through the generic metadata pass-through in DeepnoteDataConverter. @@ -23,7 +23,7 @@ export class AgentBlockConverter implements BlockConverter { } convertToCell(block: DeepnoteBlock): NotebookCellData { - const cell = new NotebookCellData(NotebookCellKind.Code, block.content || '', 'markdown'); + const cell = new NotebookCellData(NotebookCellKind.Code, block.content || '', 'plaintext'); return cell; } diff --git a/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts index 43fac425db..a3ce26acf8 100644 --- a/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts +++ b/src/notebooks/deepnote/converters/agentBlockConverter.unit.test.ts @@ -36,7 +36,7 @@ suite('AgentBlockConverter', () => { }); suite('convertToCell', () => { - test('converts agent block to code cell with markdown language', () => { + test('converts agent block to code cell with plaintext language', () => { const block: DeepnoteBlock = { blockGroup: 'test-group', content: 'Analyze the dataset and create a summary report', @@ -50,7 +50,7 @@ suite('AgentBlockConverter', () => { assert.strictEqual(cell.kind, NotebookCellKind.Code); assert.strictEqual(cell.value, 'Analyze the dataset and create a summary report'); - assert.strictEqual(cell.languageId, 'markdown'); + assert.strictEqual(cell.languageId, 'plaintext'); }); test('handles empty content', () => { @@ -67,7 +67,7 @@ suite('AgentBlockConverter', () => { assert.strictEqual(cell.kind, NotebookCellKind.Code); assert.strictEqual(cell.value, ''); - assert.strictEqual(cell.languageId, 'markdown'); + assert.strictEqual(cell.languageId, 'plaintext'); }); test('handles undefined content', () => { @@ -83,7 +83,7 @@ suite('AgentBlockConverter', () => { assert.strictEqual(cell.kind, NotebookCellKind.Code); assert.strictEqual(cell.value, ''); - assert.strictEqual(cell.languageId, 'markdown'); + assert.strictEqual(cell.languageId, 'plaintext'); }); test('preserves multiline prompt', () => { @@ -109,7 +109,7 @@ suite('AgentBlockConverter', () => { assert.strictEqual(cell.kind, NotebookCellKind.Code); assert.strictEqual(cell.value, prompt); - assert.strictEqual(cell.languageId, 'markdown'); + assert.strictEqual(cell.languageId, 'plaintext'); }); test('preserves agent block with metadata', () => { @@ -128,7 +128,7 @@ suite('AgentBlockConverter', () => { assert.strictEqual(cell.kind, NotebookCellKind.Code); assert.strictEqual(cell.value, 'Analyze the data'); - assert.strictEqual(cell.languageId, 'markdown'); + assert.strictEqual(cell.languageId, 'plaintext'); }); }); @@ -145,7 +145,7 @@ suite('AgentBlockConverter', () => { const cell = new NotebookCellData( NotebookCellKind.Code, 'New prompt with updated instructions', - 'markdown' + 'plaintext' ); converter.applyChangesToBlock(block, cell); @@ -162,7 +162,7 @@ suite('AgentBlockConverter', () => { metadata: { deepnote_agent_model: 'auto' }, type: 'agent' }; - const cell = new NotebookCellData(NotebookCellKind.Code, '', 'markdown'); + const cell = new NotebookCellData(NotebookCellKind.Code, '', 'plaintext'); converter.applyChangesToBlock(block, cell); @@ -181,7 +181,7 @@ suite('AgentBlockConverter', () => { sortingKey: 'a2', type: 'agent' }; - const cell = new NotebookCellData(NotebookCellKind.Code, 'New prompt', 'markdown'); + const cell = new NotebookCellData(NotebookCellKind.Code, 'New prompt', 'plaintext'); converter.applyChangesToBlock(block, cell); diff --git a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts index 63e7129200..2a5964b6c8 100644 --- a/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts +++ b/src/notebooks/deepnote/deepnoteKernelAutoSelector.node.ts @@ -56,6 +56,7 @@ import { logger } from '../../platform/logging'; import { PythonEnvironment } from '../../platform/pythonEnvironments/info'; import { IControllerRegistration, IVSCodeNotebookController } from '../controllers/types'; import { IDeepnoteNotebookManager } from '../types'; +import { executeAgentCell, isAgentCell } from './agentCellExecutionHandler'; import { IDeepnoteInitNotebookRunner } from './deepnoteInitNotebookRunner.node'; import { computeRequirementsHash } from './deepnoteProjectUtils'; import { IDeepnoteRequirementsHelper } from './deepnoteRequirementsHelper.node'; @@ -1204,7 +1205,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, ); controller.supportsExecutionOrder = true; - controller.supportedLanguages = ['python', 'sql', 'markdown']; + controller.supportedLanguages = ['python', 'sql', 'markdown', 'plaintext']; // Execution handler that shows environment picker when user tries to run without an environment controller.executeHandler = async (cells, doc) => { @@ -1214,6 +1215,28 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, } cells` ); + const agentCells = cells.filter((cell) => isAgentCell(cell)); + const kernelCells = cells.filter((cell) => !isAgentCell(cell)); + + // Execute agent cells directly without kernel involvement + if (agentCells.length > 0) { + logger.info( + `Executing ${agentCells.length} agent cell(s) for ${getDisplayPath(doc.uri)} without kernel` + ); + + for (const cell of agentCells) { + try { + await executeAgentCell(cell, controller); + } catch (cellError) { + logger.error(`Error executing agent cell ${cell.index}`, cellError); + } + } + } + + if (kernelCells.length === 0) { + return; + } + // Create a cancellation token that cancels when the notebook is closed const cts = new CancellationTokenSource(); const closeListener = workspace.onDidCloseNotebookDocument((closedDoc) => { @@ -1242,7 +1265,7 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, return; } - logger.info(`Executing ${cells.length} cells through kernel after environment configuration`); + logger.info(`Executing ${kernelCells.length} cells through kernel after environment configuration`); // Get or create a kernel for this notebook with the new connection const kernel = this.kernelProvider.getOrCreate(doc, { @@ -1254,16 +1277,15 @@ export class DeepnoteKernelAutoSelector implements IDeepnoteKernelAutoSelector, // Execute cells through the kernel const kernelExecution = this.kernelProvider.getKernelExecution(kernel); - for (const cell of cells) { + for (const cell of kernelCells) { try { await kernelExecution.executeCell(cell); } catch (cellError) { logger.error(`Error executing cell ${cell.index}`, cellError); - // Continue with remaining cells } } - logger.info(`Finished executing ${cells.length} cells`); + logger.info(`Finished executing ${kernelCells.length} cells`); } catch (error) { if (isCancellationError(error)) { logger.info(`Environment setup cancelled for ${getDisplayPath(doc.uri)}`); From 9fbe622953b843ad4c9b57c6663041188d1f5bf9 Mon Sep 17 00:00:00 2001 From: tomas Date: Fri, 13 Mar 2026 13:13:55 +0000 Subject: [PATCH 3/8] feat(agent-block): Integrate deepnote runtime-core to execute Agent blocks --- .../deepnote/agentCellExecutionHandler.ts | 296 +++++++++++++++++- .../deepnote/deepnoteDataConverter.ts | 2 +- 2 files changed, 281 insertions(+), 17 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index 51ab4b0ce5..6de5454b4a 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -1,7 +1,32 @@ -import { NotebookCell, NotebookCellOutput, NotebookCellOutputItem, NotebookController } from 'vscode'; +import { + NotebookCell, + NotebookCellOutput, + NotebookCellOutputItem, + NotebookController, + NotebookDocument, + NotebookEdit, + NotebookRange, + WorkspaceEdit, + commands, + workspace +} from 'vscode'; +import { AgentBlock, DeepnoteBlock } from '@deepnote/blocks'; +import { + AgentBlockContext, + AgentStreamEvent, + executeAgentBlock, + serializeNotebookContextFromBlocks +} from '@deepnote/runtime-core'; + +import { translateCellDisplayOutput } from '../../kernels/execution/helpers'; +import { createDeferred } from '../../platform/common/utils/async'; +import { uuidUtils } from '../../platform/common/uuid'; import type { Pocket } from '../../platform/deepnote/pocket'; import { logger } from '../../platform/logging'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; +import { generateBlockId, generateSortingKey } from './dataConversionUtils'; +import { DeepnoteDataConverter } from './deepnoteDataConverter'; export function isAgentCell(cell: NotebookCell): boolean { const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; @@ -9,43 +34,282 @@ export function isAgentCell(cell: NotebookCell): boolean { return pocket?.type === 'agent'; } +export function serializeNotebookContext({ cells }: { cells: NotebookCell[] }): string { + const converter = new DeepnoteDataConverter(); + + const blocks = cells.reduce((acc, cell) => { + try { + const block = converter.convertCellToBlock( + { + kind: cell.kind, + value: cell.document.getText(), + languageId: cell.document.languageId, + metadata: cell.metadata, + outputs: [...(cell.outputs || [])] + }, + cell.index + ); + acc.push(block); + } catch (error) { + logger.error(`Error converting cell to block: ${error}`); + } + return acc; + }, []); + + return serializeNotebookContextFromBlocks({ blocks, notebookName: null }); +} + export async function executeAgentCell(cell: NotebookCell, controller: NotebookController): Promise { const execution = controller.createNotebookCellExecution(cell); execution.start(Date.now()); try { await execution.clearOutput(); + const prompt = cell.document.getText(); - const output = new NotebookCellOutput([ - NotebookCellOutputItem.text(`[Agent] Received prompt (${prompt.length} chars)...\n`) - ]); + let accumulated = `[Agent] Planning next steps...`; + const output = new NotebookCellOutput([NotebookCellOutputItem.text(accumulated)]); await execution.replaceOutput([output]); - const chunks = [ - { delay: 500, text: '[Agent] Analyzing prompt...\n' }, - { delay: 1000, text: '[Agent] Generating plan...\n' }, - { delay: 2000, text: '[Agent] Executing steps...\n' }, - { delay: 3000, text: `[Agent] Done.\n\nPrompt: ${prompt}\n` } - ]; + await removeEphemeralCellsForAgent(cell); - let accumulated = `[Agent] Received prompt (${prompt.length} chars)...\n`; - for (const chunk of chunks) { - await delay(chunk.delay); - accumulated += chunk.text; - await execution.replaceOutputItems(NotebookCellOutputItem.text(accumulated), output); + const dataConverter = new DeepnoteDataConverter(); + const deepnoteBlock = dataConverter.convertCellToBlock( + { + kind: cell.kind, + value: cell.document.getText(), + languageId: cell.document.languageId, + metadata: cell.metadata, + outputs: [...(cell.outputs || [])] + }, + cell.index + ); + const agentBlock: AgentBlock | null = deepnoteBlock.type === 'agent' ? deepnoteBlock : null; + + if (agentBlock == null) { + // TODO: better DX error handling + throw new Error('Cell is not an agent cell'); } + let lastAgentEventType: AgentStreamEvent['type'] | undefined; + + const notebookContext = serializeNotebookContext({ + cells: cell.notebook.getCells().filter((c) => c.index !== cell.index) + }); + + const openAiToken = process.env.OPENAI_API_KEY; + if (openAiToken == null) { + throw new Error('OPENAI_API_KEY is not set'); + } + + const context: AgentBlockContext = { + openAiToken, + mcpServers: [], + notebookContext, + addMarkdownBlock: async ({ content }: { content: string }) => { + await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'markdown', content); + return { success: true }; + }, + addAndExecuteCodeBlock: async ({ code }: { code: string }) => { + const cellIndex = await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'code', code); + + const { success } = await executeEphemeralCell(cell.notebook, cellIndex); + return success ? { success } : { success: false, error: new Error('Ephemeral cell execution failed') }; + }, + onLog: (message: string) => { + logger.info('Agent log', message); + // accumulated += message; + // TODO: replaceOutputItems is Async function + // execution.replaceOutputItems(NotebookCellOutputItem.text(accumulated), output); + }, + onAgentEvent: async (event: AgentStreamEvent) => { + logger.info('Agent event', JSON.stringify(event)); + if (lastAgentEventType != null && lastAgentEventType !== event.type) { + accumulated += `\n\n`; + } + switch (event.type) { + case 'tool_called': + // Ignore calling tool_called events + // accumulated += `[Agent] Tool called: ${event.toolName}`; + break; + case 'tool_output': + accumulated += `[Agent] Tool output: ${event.toolName}`; + accumulated += `[Agent] Tool output length: ${event.output?.length}`; + break; + case 'text_delta': + if (lastAgentEventType !== 'text_delta') { + accumulated += `[Agent] Text:\n`; + } + accumulated += event.text; + break; + case 'reasoning_delta': + if (lastAgentEventType !== 'reasoning_delta') { + accumulated += `[Agent] Reasoning:\n`; + } + accumulated += event.text; + break; + default: + event satisfies never; + } + lastAgentEventType = event.type; + + await execution.replaceOutputItems(NotebookCellOutputItem.text(accumulated), output); + } + }; + + logger.info( + `Agent cell: starting executeAgentBlock, model=${agentBlock.metadata.deepnote_agent_model}, prompt length=${prompt.length}` + ); + const result = await executeAgentBlock(agentBlock, context); + logger.info(`Agent cell: executeAgentBlock completed, finalOutput length=${result.finalOutput.length}`); + execution.end(true, Date.now()); } catch (error) { logger.error('Agent cell execution failed', error); + if (error instanceof Error) { + logger.error(`Agent error name=${error.name}, message=${error.message}`); + if (error.cause) { + logger.error('Agent error cause:', error.cause); + } + if (error.stack) { + logger.error('Agent error stack:', error.stack); + } + } + const message = error instanceof Error ? error.message : String(error); const stderrOutput = new NotebookCellOutput([NotebookCellOutputItem.stderr(message)]); - await execution.replaceOutput([stderrOutput]).then(undefined, () => undefined); + await execution.appendOutput([stderrOutput]).then(undefined, () => undefined); execution.end(false, Date.now()); } } +function getInsertIndexAfterAgentCell( + notebook: NotebookDocument, + agentCellIndex: number, + agentBlockId: string +): number { + let index = agentCellIndex + 1; + + while (index < notebook.cellCount) { + const cell = notebook.cellAt(index); + if (cell.metadata?.is_ephemeral === true && cell.metadata?.agent_source_block_id === agentBlockId) { + index++; + } else { + break; + } + } + + return index; +} + +async function insertEphemeralCell( + notebook: NotebookDocument, + agentCellIndex: number, + agentBlockId: string, + blockType: 'code' | 'markdown', + content: string +): Promise { + const insertIndex = getInsertIndexAfterAgentCell(notebook, agentCellIndex, agentBlockId); + + const block: DeepnoteBlock = { + type: blockType, + id: generateBlockId(), + blockGroup: uuidUtils.generateUuid(), + sortingKey: generateSortingKey(insertIndex), + content, + metadata: { + is_ephemeral: true, + agent_source_block_id: agentBlockId + } + }; + + const converter = new DeepnoteDataConverter(); + const [cellData] = converter.convertBlocksToCells([block]); + + const edit = new WorkspaceEdit(); + edit.set(notebook.uri, [NotebookEdit.insertCells(insertIndex, [cellData])]); + await workspace.applyEdit(edit); + + return insertIndex; +} + +const EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS = 5 * 60 * 1000; + +async function executeEphemeralCell( + notebook: NotebookDocument, + cellIndex: number +): Promise<{ success: boolean; outputs: unknown[]; executionCount: number | null }> { + const cell = notebook.cellAt(cellIndex); + const completionDeferred = createDeferred(); + + const disposable = notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell === cell && e.state === NotebookCellExecutionState.Idle) { + completionDeferred.resolve(); + } + }); + + const timeout = setTimeout(() => { + completionDeferred.reject(new Error('Ephemeral cell execution timed out')); + }, EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS); + + try { + await commands.executeCommand('notebook.cell.execute', { + ranges: [{ start: cellIndex, end: cellIndex + 1 }], + document: notebook.uri + }); + + await completionDeferred.promise; + + return { + success: cell.executionSummary?.success !== false, + outputs: cell.outputs.map(translateCellDisplayOutput), + executionCount: cell.executionSummary?.executionOrder ?? null + }; + } catch (error) { + return { + success: false, + outputs: [], + executionCount: null + }; + } finally { + disposable.dispose(); + clearTimeout(timeout); + } +} + +async function removeEphemeralCellsForAgent(agentCell: NotebookCell): Promise { + const agentBlockId = (agentCell.metadata?.id ?? agentCell.metadata?.__deepnoteBlockId) as string | undefined; + if (!agentBlockId) { + return; + } + + const notebook = agentCell.notebook; + const deletions: NotebookEdit[] = []; + + for (let i = notebook.cellCount - 1; i >= 0; i--) { + const cell = notebook.cellAt(i); + + if (cell.metadata?.is_ephemeral === true && cell.metadata?.agent_source_block_id === agentBlockId) { + deletions.push(NotebookEdit.deleteCells(new NotebookRange(i, i + 1))); + } + } + + if (deletions.length === 0) { + return; + } + + const edit = new WorkspaceEdit(); + edit.set(notebook.uri, deletions); + + const success = await workspace.applyEdit(edit); + if (success) { + logger.info(`Removed ${deletions.length} ephemeral cell(s) for agent block ${agentBlockId}`); + } else { + logger.warn(`Failed to remove ephemeral cells for agent block ${agentBlockId}`); + } +} + function delay(ms: number): Promise { return new Promise((resolve) => setTimeout(resolve, ms)); } diff --git a/src/notebooks/deepnote/deepnoteDataConverter.ts b/src/notebooks/deepnote/deepnoteDataConverter.ts index 51007cae9d..d6ef93b52f 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.ts @@ -1,5 +1,5 @@ import { isExecutableBlock, type DeepnoteBlock } from '@deepnote/blocks'; -import { NotebookCellData, NotebookCellKind, NotebookCellOutput, NotebookCellOutputItem } from 'vscode'; +import { NotebookCell, NotebookCellData, NotebookCellKind, NotebookCellOutput, NotebookCellOutputItem } from 'vscode'; import { generateBlockId, generateSortingKey } from './dataConversionUtils'; import type { DeepnoteOutput } from '../../platform/deepnote/deepnoteTypes'; From 46f9a4c1184229c9d750d17373f91f3cda76c7c9 Mon Sep 17 00:00:00 2001 From: tomas Date: Sat, 14 Mar 2026 09:41:46 +0000 Subject: [PATCH 4/8] feat(agent-cell): Enhance executeAgentCell with options for custom execution functions and improve ephemeral cell handling --- .../deepnote/agentCellExecutionHandler.ts | 46 ++-- .../agentCellExecutionHandler.unit.test.ts | 208 ++++++++++-------- .../deepnote/deepnoteDataConverter.ts | 2 +- src/notebooks/deepnote/deepnoteTestHelpers.ts | 15 +- 4 files changed, 151 insertions(+), 120 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index 6de5454b4a..b84ee63255 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -59,7 +59,16 @@ export function serializeNotebookContext({ cells }: { cells: NotebookCell[] }): return serializeNotebookContextFromBlocks({ blocks, notebookName: null }); } -export async function executeAgentCell(cell: NotebookCell, controller: NotebookController): Promise { +export interface ExecuteAgentCellOptions { + executeAgentBlockFn?: typeof executeAgentBlock; +} + +export async function executeAgentCell( + cell: NotebookCell, + controller: NotebookController, + options?: ExecuteAgentCellOptions +): Promise { + const executeAgentBlockFn = options?.executeAgentBlockFn ?? executeAgentBlock; const execution = controller.createNotebookCellExecution(cell); execution.start(Date.now()); @@ -72,8 +81,6 @@ export async function executeAgentCell(cell: NotebookCell, controller: NotebookC const output = new NotebookCellOutput([NotebookCellOutputItem.text(accumulated)]); await execution.replaceOutput([output]); - await removeEphemeralCellsForAgent(cell); - const dataConverter = new DeepnoteDataConverter(); const deepnoteBlock = dataConverter.convertCellToBlock( { @@ -92,6 +99,8 @@ export async function executeAgentCell(cell: NotebookCell, controller: NotebookC throw new Error('Cell is not an agent cell'); } + await removeEphemeralCellsForAgent(cell.notebook, agentBlock.id); + let lastAgentEventType: AgentStreamEvent['type'] | undefined; const notebookContext = serializeNotebookContext({ @@ -113,8 +122,9 @@ export async function executeAgentCell(cell: NotebookCell, controller: NotebookC }, addAndExecuteCodeBlock: async ({ code }: { code: string }) => { const cellIndex = await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'code', code); + const insertedCell = cell.notebook.cellAt(cellIndex); - const { success } = await executeEphemeralCell(cell.notebook, cellIndex); + const { success } = await executeEphemeralCell(insertedCell); return success ? { success } : { success: false, error: new Error('Ephemeral cell execution failed') }; }, onLog: (message: string) => { @@ -131,10 +141,10 @@ export async function executeAgentCell(cell: NotebookCell, controller: NotebookC switch (event.type) { case 'tool_called': // Ignore calling tool_called events - // accumulated += `[Agent] Tool called: ${event.toolName}`; + accumulated += `[Agent] Tool called: ${event.toolName}`; break; case 'tool_output': - accumulated += `[Agent] Tool output: ${event.toolName}`; + accumulated += `[Agent] Tool output: ${event.toolName}\n`; accumulated += `[Agent] Tool output length: ${event.output?.length}`; break; case 'text_delta': @@ -161,7 +171,7 @@ export async function executeAgentCell(cell: NotebookCell, controller: NotebookC logger.info( `Agent cell: starting executeAgentBlock, model=${agentBlock.metadata.deepnote_agent_model}, prompt length=${prompt.length}` ); - const result = await executeAgentBlock(agentBlock, context); + const result = await executeAgentBlockFn(agentBlock, context); logger.info(`Agent cell: executeAgentBlock completed, finalOutput length=${result.finalOutput.length}`); execution.end(true, Date.now()); @@ -236,11 +246,9 @@ async function insertEphemeralCell( const EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS = 5 * 60 * 1000; -async function executeEphemeralCell( - notebook: NotebookDocument, - cellIndex: number +export async function executeEphemeralCell( + cell: NotebookCell ): Promise<{ success: boolean; outputs: unknown[]; executionCount: number | null }> { - const cell = notebook.cellAt(cellIndex); const completionDeferred = createDeferred(); const disposable = notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { @@ -254,9 +262,11 @@ async function executeEphemeralCell( }, EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS); try { + const cellIndex = cell.index; + await commands.executeCommand('notebook.cell.execute', { ranges: [{ start: cellIndex, end: cellIndex + 1 }], - document: notebook.uri + document: cell.notebook.uri }); await completionDeferred.promise; @@ -278,13 +288,7 @@ async function executeEphemeralCell( } } -async function removeEphemeralCellsForAgent(agentCell: NotebookCell): Promise { - const agentBlockId = (agentCell.metadata?.id ?? agentCell.metadata?.__deepnoteBlockId) as string | undefined; - if (!agentBlockId) { - return; - } - - const notebook = agentCell.notebook; +async function removeEphemeralCellsForAgent(notebook: NotebookDocument, agentBlockId: string): Promise { const deletions: NotebookEdit[] = []; for (let i = notebook.cellCount - 1; i >= 0; i--) { @@ -309,7 +313,3 @@ async function removeEphemeralCellsForAgent(agentCell: NotebookCell): Promise { - return new Promise((resolve) => setTimeout(resolve, ms)); -} diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index 77544aeafd..b84bb6286c 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -1,8 +1,17 @@ import { expect } from 'chai'; import * as sinon from 'sinon'; +import { anything, capture, reset, when } from 'ts-mockito'; import { NotebookCellOutput, NotebookCellOutputItem, NotebookController } from 'vscode'; -import { executeAgentCell, isAgentCell } from './agentCellExecutionHandler'; +import type { AgentBlock } from '@deepnote/blocks'; +import type { AgentBlockContext, AgentBlockResult } from '@deepnote/runtime-core'; + +import { + NotebookCellExecutionState, + notebookCellExecutions +} from '../../platform/notebooks/cellExecutionStateService'; +import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; +import { executeAgentCell, executeEphemeralCell, isAgentCell } from './agentCellExecutionHandler'; import { createMockCell } from './deepnoteTestHelpers'; suite('AgentCellExecutionHandler', () => { @@ -39,8 +48,8 @@ suite('AgentCellExecutionHandler', () => { }); suite('executeAgentCell', () => { - let clock: sinon.SinonFakeTimers; let mockExecution: { + appendOutput: sinon.SinonStub; clearOutput: sinon.SinonStub; end: sinon.SinonStub; replaceOutput: sinon.SinonStub; @@ -48,11 +57,15 @@ suite('AgentCellExecutionHandler', () => { start: sinon.SinonStub; }; let mockController: NotebookController; + let executeAgentBlockStub: sinon.SinonStub; + let savedOpenAiKey: string | undefined; setup(() => { - clock = sinon.useFakeTimers(); + savedOpenAiKey = process.env.OPENAI_API_KEY; + process.env.OPENAI_API_KEY = 'test-key'; mockExecution = { + appendOutput: sinon.stub().resolves(), clearOutput: sinon.stub().resolves(), end: sinon.stub(), replaceOutput: sinon.stub().resolves(), @@ -63,52 +76,47 @@ suite('AgentCellExecutionHandler', () => { mockController = { createNotebookCellExecution: sinon.stub().returns(mockExecution) } as unknown as NotebookController; + + executeAgentBlockStub = sinon.stub().resolves({ finalOutput: 'done' } as AgentBlockResult); }); teardown(() => { - clock.restore(); + if (savedOpenAiKey !== undefined) { + process.env.OPENAI_API_KEY = savedOpenAiKey; + } else { + delete process.env.OPENAI_API_KEY; + } }); - async function runToCompletion(promise: Promise): Promise { - // Total delay across all chunks: 500 + 1000 + 2000 + 3000 = 6500ms - await clock.tickAsync(7000); - await promise; + function createAgentCell(text: string = 'Test prompt') { + return createMockCell({ + metadata: { __deepnotePocket: { type: 'agent' } }, + text + }); } test('creates execution and starts it', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Analyze data' - }); + const cell = createAgentCell('Analyze data'); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect((mockController.createNotebookCellExecution as sinon.SinonStub).calledOnceWith(cell)).to.be.true; expect(mockExecution.start.calledOnce).to.be.true; }); test('clears output before streaming', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Analyze data' - }); + const cell = createAgentCell('Analyze data'); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect(mockExecution.clearOutput.calledOnce).to.be.true; expect(mockExecution.clearOutput.calledBefore(mockExecution.replaceOutput)).to.be.true; }); test('sets initial output via replaceOutput', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Hello world' - }); + const cell = createAgentCell('Hello world'); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect(mockExecution.replaceOutput.calledOnce).to.be.true; @@ -117,29 +125,35 @@ suite('AgentCellExecutionHandler', () => { expect(outputs[0].items).to.have.lengthOf(1); const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); - expect(text).to.include('[Agent] Received prompt (11 chars)'); + expect(text).to.include('[Agent] Planning next steps...'); }); - test('streams 4 chunks via replaceOutputItems', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Test prompt' + test('streams events via replaceOutputItems using onAgentEvent callback', async () => { + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.onAgentEvent?.({ type: 'text_delta', text: 'Hello ' }); + await context.onAgentEvent?.({ type: 'text_delta', text: 'world' }); + + return { finalOutput: 'Hello world' } as AgentBlockResult; }); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - expect(mockExecution.replaceOutputItems.callCount).to.equal(4); + expect(mockExecution.replaceOutputItems.callCount).to.equal(2); }); test('streaming chunks accumulate text progressively', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Test' + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.onAgentEvent?.({ type: 'text_delta', text: 'first' }); + await context.onAgentEvent?.({ type: 'text_delta', text: ' second' }); + + return { finalOutput: 'first second' } as AgentBlockResult; }); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); const getChunkText = (callIndex: number): string => { const item = mockExecution.replaceOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; @@ -149,51 +163,39 @@ suite('AgentCellExecutionHandler', () => { const chunk1 = getChunkText(0); const chunk2 = getChunkText(1); - const chunk3 = getChunkText(2); - const chunk4 = getChunkText(3); - - expect(chunk1).to.include('Analyzing prompt'); - expect(chunk2).to.include('Generating plan'); - expect(chunk2).to.include('Analyzing prompt'); - expect(chunk3).to.include('Executing steps'); - expect(chunk3).to.include('Generating plan'); - expect(chunk4).to.include('Done'); - expect(chunk4).to.include('Prompt: Test'); - }); - test('streaming chunks fire at correct intervals', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Test' - }); + expect(chunk1).to.include('[Agent] Text:'); + expect(chunk1).to.include('first'); + expect(chunk2).to.include('first second'); + }); - const promise = executeAgentCell(cell, mockController); + test('separates different event types with blank lines', async () => { + executeAgentBlockStub.callsFake(async (_block: AgentBlock, context: AgentBlockContext) => { + await context.onAgentEvent?.({ type: 'text_delta', text: 'thinking...' }); + await context.onAgentEvent?.({ type: 'tool_called', toolName: 'search' }); - expect(mockExecution.replaceOutputItems.callCount).to.equal(0); + return { finalOutput: '' } as AgentBlockResult; + }); - await clock.tickAsync(500); - expect(mockExecution.replaceOutputItems.callCount).to.equal(1); + const cell = createAgentCell(); - await clock.tickAsync(1000); - expect(mockExecution.replaceOutputItems.callCount).to.equal(2); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - await clock.tickAsync(2000); - expect(mockExecution.replaceOutputItems.callCount).to.equal(3); + const getChunkText = (callIndex: number): string => { + const item = mockExecution.replaceOutputItems.getCall(callIndex).args[0] as NotebookCellOutputItem; - await clock.tickAsync(3000); - expect(mockExecution.replaceOutputItems.callCount).to.equal(4); + return Buffer.from(item.data).toString('utf-8'); + }; - await promise; + const chunk2 = getChunkText(1); + expect(chunk2).to.include('\n\n'); + expect(chunk2).to.include('[Agent] Tool called: search'); }); test('ends execution with success', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Test' - }); + const cell = createAgentCell(); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect(mockExecution.end.calledOnce).to.be.true; expect(mockExecution.end.firstCall.args[0]).to.be.true; @@ -202,13 +204,9 @@ suite('AgentCellExecutionHandler', () => { test('ends execution with failure when error occurs', async () => { mockExecution.clearOutput.rejects(new Error('Test error')); - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Test' - }); + const cell = createAgentCell(); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect(mockExecution.end.calledOnce).to.be.true; expect(mockExecution.end.firstCall.args[0]).to.be.false; @@ -217,17 +215,13 @@ suite('AgentCellExecutionHandler', () => { test('writes error message to stderr output on failure', async () => { mockExecution.clearOutput.rejects(new Error('Something went wrong')); - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: 'Test' - }); + const cell = createAgentCell(); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); - expect(mockExecution.replaceOutput.calledOnce).to.be.true; + expect(mockExecution.appendOutput.calledOnce).to.be.true; - const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; + const outputs = mockExecution.appendOutput.firstCall.args[0] as NotebookCellOutput[]; expect(outputs).to.have.lengthOf(1); const item = outputs[0].items[0]; @@ -238,20 +232,48 @@ suite('AgentCellExecutionHandler', () => { }); test('handles empty prompt', async () => { - const cell = createMockCell({ - metadata: { __deepnotePocket: { type: 'agent' } }, - text: '' - }); + const cell = createAgentCell(''); - const promise = executeAgentCell(cell, mockController); - await runToCompletion(promise); + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); expect(mockExecution.end.calledOnce).to.be.true; expect(mockExecution.end.firstCall.args[0]).to.be.true; const outputs = mockExecution.replaceOutput.firstCall.args[0] as NotebookCellOutput[]; const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); - expect(text).to.include('(0 chars)'); + expect(text).to.include('[Agent] Planning next steps...'); + }); + }); + + suite('executeEphemeralCell', () => { + teardown(() => { + reset(mockedVSCodeNamespaces.commands); + }); + + test('uses current cell index, not stale index from insertion time', async () => { + const staleIndex = 5; + const currentIndex = 6; + + const cell = createMockCell({ index: staleIndex }); + + // Simulate a concurrent insertion shifting the cell's index + (cell as { index: number }).index = currentIndex; + + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenCall(async () => { + notebookCellExecutions.changeCellState(cell, NotebookCellExecutionState.Idle); + }); + + await executeEphemeralCell(cell); + + const [commandName, commandArg] = capture( + mockedVSCodeNamespaces.commands.executeCommand as (cmd: string, arg: unknown) => Thenable + ).last(); + + expect(commandName).to.equal('notebook.cell.execute'); + expect(commandArg).to.deep.equal({ + ranges: [{ start: currentIndex, end: currentIndex + 1 }], + document: cell.notebook.uri + }); }); }); }); diff --git a/src/notebooks/deepnote/deepnoteDataConverter.ts b/src/notebooks/deepnote/deepnoteDataConverter.ts index d6ef93b52f..51007cae9d 100644 --- a/src/notebooks/deepnote/deepnoteDataConverter.ts +++ b/src/notebooks/deepnote/deepnoteDataConverter.ts @@ -1,5 +1,5 @@ import { isExecutableBlock, type DeepnoteBlock } from '@deepnote/blocks'; -import { NotebookCell, NotebookCellData, NotebookCellKind, NotebookCellOutput, NotebookCellOutputItem } from 'vscode'; +import { NotebookCellData, NotebookCellKind, NotebookCellOutput, NotebookCellOutputItem } from 'vscode'; import { generateBlockId, generateSortingKey } from './dataConversionUtils'; import type { DeepnoteOutput } from '../../platform/deepnote/deepnoteTypes'; diff --git a/src/notebooks/deepnote/deepnoteTestHelpers.ts b/src/notebooks/deepnote/deepnoteTestHelpers.ts index 18d03e558d..eb0dc26464 100644 --- a/src/notebooks/deepnote/deepnoteTestHelpers.ts +++ b/src/notebooks/deepnote/deepnoteTestHelpers.ts @@ -47,8 +47,16 @@ export function createMockNotebook(options?: CreateMockNotebookOptions): Noteboo return { uri, notebookType, - metadata - } as NotebookDocument; + metadata, + cellCount: 0, + cellAt: () => ({}) as NotebookCell, + getCells: () => [], + version: 1, + isDirty: false, + isUntitled: false, + isClosed: false, + save: async () => true + } satisfies NotebookDocument; } /** @@ -121,7 +129,8 @@ export function createMockCell(options?: CreateMockCellOptions): NotebookCell { positionAt: () => ({}) as unknown, validateRange: () => ({}) as unknown, validatePosition: () => ({}) as unknown, - getWordRangeAtPosition: () => undefined + getWordRangeAtPosition: () => undefined, + encoding: 'utf-8' } as unknown as TextDocument; return { From 58e653d2d5efc12e750c324e5b509f71e1cf6dd0 Mon Sep 17 00:00:00 2001 From: tomas Date: Mon, 16 Mar 2026 17:10:52 +0000 Subject: [PATCH 5/8] feat(ephemeral-cells): Introduce isEphemeralCell utility and enhance handling of ephemeral cells in serialization and decoration --- build/esbuild/build.ts | 3 +- .../deepnote/agentCellExecutionHandler.ts | 7 +- src/notebooks/deepnote/dataConversionUtils.ts | 9 +++ src/notebooks/deepnote/deepnoteSerializer.ts | 17 +++-- .../deepnote/deepnoteSerializer.unit.test.ts | 65 +++++++++++++++++++ .../ephemeralCellDecorationProvider.ts | 3 +- .../ephemeralCellStatusBarProvider.ts | 7 +- 7 files changed, 96 insertions(+), 15 deletions(-) diff --git a/build/esbuild/build.ts b/build/esbuild/build.ts index c313ce8cf8..40fdd60cc3 100644 --- a/build/esbuild/build.ts +++ b/build/esbuild/build.ts @@ -72,7 +72,8 @@ const commonExternals = [ const webExternals = [ ...commonExternals, 'canvas', // Native module used by vega for server-side rendering, not needed in browser - 'mathjax-electron' // Uses Node.js path module, MathJax rendering handled differently in browser + 'mathjax-electron', // Uses Node.js path module, MathJax rendering handled differently in browser + '@deepnote/runtime-core' // Uses tcp-port-used → net, only needed in desktop for agent block execution ]; const desktopExternals = [...commonExternals, ...deskTopNodeModulesToExternalize]; const bundleConfig = getBundleConfiguration(); diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index b84ee63255..ea8485e8a4 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -25,7 +25,7 @@ import { uuidUtils } from '../../platform/common/uuid'; import type { Pocket } from '../../platform/deepnote/pocket'; import { logger } from '../../platform/logging'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; -import { generateBlockId, generateSortingKey } from './dataConversionUtils'; +import { generateBlockId, generateSortingKey, isEphemeralCell } from './dataConversionUtils'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; export function isAgentCell(cell: NotebookCell): boolean { @@ -107,6 +107,7 @@ export async function executeAgentCell( cells: cell.notebook.getCells().filter((c) => c.index !== cell.index) }); + // eslint-disable-next-line local-rules/dont-use-process const openAiToken = process.env.OPENAI_API_KEY; if (openAiToken == null) { throw new Error('OPENAI_API_KEY is not set'); @@ -203,7 +204,7 @@ function getInsertIndexAfterAgentCell( while (index < notebook.cellCount) { const cell = notebook.cellAt(index); - if (cell.metadata?.is_ephemeral === true && cell.metadata?.agent_source_block_id === agentBlockId) { + if (isEphemeralCell(cell) && cell.metadata?.agent_source_block_id === agentBlockId) { index++; } else { break; @@ -294,7 +295,7 @@ async function removeEphemeralCellsForAgent(notebook: NotebookDocument, agentBlo for (let i = notebook.cellCount - 1; i >= 0; i--) { const cell = notebook.cellAt(i); - if (cell.metadata?.is_ephemeral === true && cell.metadata?.agent_source_block_id === agentBlockId) { + if (isEphemeralCell(cell) && cell.metadata?.agent_source_block_id === agentBlockId) { deletions.push(NotebookEdit.deleteCells(new NotebookRange(i, i + 1))); } } diff --git a/src/notebooks/deepnote/dataConversionUtils.ts b/src/notebooks/deepnote/dataConversionUtils.ts index 1b30484770..8b01da1256 100644 --- a/src/notebooks/deepnote/dataConversionUtils.ts +++ b/src/notebooks/deepnote/dataConversionUtils.ts @@ -2,6 +2,8 @@ * Utility functions for Deepnote block ID and sorting key generation */ +import { NotebookCell, NotebookCellData } from 'vscode'; + export function parseJsonWithFallback(value: string, fallback?: unknown): unknown | null { try { return JSON.parse(value); @@ -22,6 +24,13 @@ export function generateBlockId(): string { return id; } +/** + * Returns true if the cell metadata indicates an ephemeral cell (auto-generated by agent). + */ +export function isEphemeralCell(cell: NotebookCell | NotebookCellData): boolean { + return cell.metadata?.is_ephemeral === true; +} + /** * Generate sorting key based on index (format: a0, a1, ..., a99, b0, b1, ...) */ diff --git a/src/notebooks/deepnote/deepnoteSerializer.ts b/src/notebooks/deepnote/deepnoteSerializer.ts index 17443e93b9..e9372f94fc 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.ts @@ -6,6 +6,7 @@ import { l10n, window, workspace, type CancellationToken, type NotebookData, typ import { logger } from '../../platform/logging'; import { IDeepnoteNotebookManager } from '../types'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; +import { isEphemeralCell } from './dataConversionUtils'; import type { DeepnoteNotebook } from '../../platform/deepnote/deepnoteTypes'; import { SnapshotService } from './snapshots/snapshotService'; import { computeHash } from '../../platform/common/crypto'; @@ -273,11 +274,17 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { throw new Error(`Notebook with ID ${notebookId} not found in project`); } - logger.debug(`SerializeNotebook: Found notebook, converting ${data.cells.length} cells to blocks`); + // Exclude ephemeral cells (agent-generated) from persistence + const nonEphemeralCells = data.cells.filter((cell) => !isEphemeralCell(cell)); + + logger.debug( + `SerializeNotebook: Found notebook, converting ${nonEphemeralCells.length} cells to blocks ` + + `(${data.cells.length - nonEphemeralCells.length} ephemeral excluded)` + ); // Log cell metadata IDs before conversion - for (let i = 0; i < data.cells.length; i++) { - const cell = data.cells[i]; + for (let i = 0; i < nonEphemeralCells.length; i++) { + const cell = nonEphemeralCells[i]; logger.trace( `SerializeNotebook: cell[${i}] metadata.id=${cell.metadata?.id}, metadata keys=${ cell.metadata ? Object.keys(cell.metadata).join(',') : 'none' @@ -287,7 +294,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { // Clone blocks while removing circular references that may have been // introduced by VS Code's notebook cell/output handling - const blocks = this.converter.convertCellsToBlocks(data.cells); + const blocks = this.converter.convertCellsToBlocks(nonEphemeralCells); logger.debug(`SerializeNotebook: Converted to ${blocks.length} blocks`); @@ -301,7 +308,7 @@ export class DeepnoteNotebookSerializer implements NotebookSerializer { } // Add snapshot metadata to blocks (contentHash and execution timing) - await this.addSnapshotMetadataToBlocks(blocks, data); + await this.addSnapshotMetadataToBlocks(blocks, { ...data, cells: nonEphemeralCells }); // Handle snapshot mode: strip outputs and execution metadata from main file if (this.snapshotService?.isSnapshotsEnabled()) { diff --git a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts index c3332f974d..2813f4acd0 100644 --- a/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts +++ b/src/notebooks/deepnote/deepnoteSerializer.unit.test.ts @@ -205,6 +205,71 @@ project: assert.include(yamlString, 'project-123'); assert.include(yamlString, 'notebook-1'); }); + + test('should exclude ephemeral cells from serialized output', async () => { + const projectData: DeepnoteFile = { + version: '1.0.0', + metadata: { + createdAt: '2023-01-01T00:00:00Z', + modifiedAt: '2023-01-02T00:00:00Z' + }, + project: { + id: 'project-ephemeral-exclude', + name: 'Ephemeral Exclude Test', + notebooks: [ + { + id: 'notebook-1', + name: 'Test Notebook', + blocks: [ + { + id: 'block-1', + content: 'print("persisted")', + blockGroup: 'group-1', + metadata: {}, + sortingKey: 'a0', + type: 'code' + } + ], + executionMode: 'block', + isModule: false + } + ], + settings: {} + } + }; + + manager.storeOriginalProject('project-ephemeral-exclude', projectData, 'notebook-1'); + + const mockNotebookData = { + cells: [ + { + kind: 2, + value: 'print("persisted")', + languageId: 'python', + metadata: { id: 'block-1' } + }, + { + kind: 2, + value: 'print("ephemeral - should not persist")', + languageId: 'python', + metadata: { id: 'ephemeral-block', is_ephemeral: true } + } + ], + metadata: { + deepnoteProjectId: 'project-ephemeral-exclude', + deepnoteNotebookId: 'notebook-1' + } + }; + + const result = await serializer.serializeNotebook(mockNotebookData as any, {} as any); + const yamlString = new TextDecoder().decode(result); + const parsedResult = deserializeDeepnoteFile(yamlString); + + const notebook = parsedResult.project.notebooks.find((nb) => nb.id === 'notebook-1'); + assert.isDefined(notebook); + assert.strictEqual(notebook!.blocks.length, 1, 'Ephemeral cell should be excluded'); + assert.strictEqual(notebook!.blocks[0].content, 'print("persisted")'); + }); }); suite('findCurrentNotebookId', () => { diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts index 5c913700ff..36b5d24053 100644 --- a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts @@ -12,6 +12,7 @@ import { } from 'vscode'; import { injectable } from 'inversify'; +import { isEphemeralCell } from './dataConversionUtils'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; const NOTEBOOK_CELL_SCHEME = 'vscode-notebook-cell'; @@ -106,7 +107,7 @@ export class EphemeralCellDecorationProvider implements IExtensionSyncActivation } const cell = this.findCellForEditor(editor); - if (!cell || cell.metadata?.is_ephemeral !== true) { + if (!cell || !isEphemeralCell(cell)) { editor.setDecorations(this.ephemeralDecorationType, []); continue; } diff --git a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts index 7089391e7c..60c0a67fba 100644 --- a/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellStatusBarProvider.ts @@ -11,6 +11,7 @@ import { } from 'vscode'; import { injectable } from 'inversify'; +import { isEphemeralCell } from './dataConversionUtils'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; const EPHEMERAL_INDICATOR_PRIORITY = 1000; @@ -54,7 +55,7 @@ export class EphemeralCellStatusBarProvider return undefined; } - if (!this.isEphemeralCell(cell)) { + if (!isEphemeralCell(cell)) { return undefined; } @@ -76,8 +77,4 @@ export class EphemeralCellStatusBarProvider tooltip: tooltipLines.join('\n') }; } - - private isEphemeralCell(cell: NotebookCell): boolean { - return cell.metadata?.is_ephemeral === true; - } } From 75d02207abc9a1c2dfabe43798918f414bce9ef1 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 18 Mar 2026 10:16:58 +0000 Subject: [PATCH 6/8] Enhance agent cell execution handling and status bar provider - Introduced `getOpenAiApiKey` function to retrieve the OpenAI API key from configuration, improving error handling when the key is not set. - Updated `executeAgentCell` and `executeEphemeralCell` functions to utilize the new API key retrieval method and handle cancellation tokens. - Enhanced `AgentCellStatusBarProvider` to validate max iterations using Zod schema, ensuring robust input handling and defaulting to safe values. - Added unit tests for new functionality and edge cases in both execution handling and status bar provider. --- .../deepnote/agentCellExecutionHandler.ts | 49 ++++++++--- .../agentCellExecutionHandler.unit.test.ts | 83 +++++++++++++++---- .../deepnote/agentCellStatusBarProvider.ts | 19 ++++- .../agentCellStatusBarProvider.unit.test.ts | 60 ++++++++++++++ 4 files changed, 177 insertions(+), 34 deletions(-) diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index ea8485e8a4..e8a1268194 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -1,4 +1,6 @@ import { + CancellationError, + CancellationToken, NotebookCell, NotebookCellOutput, NotebookCellOutputItem, @@ -20,7 +22,9 @@ import { } from '@deepnote/runtime-core'; import { translateCellDisplayOutput } from '../../kernels/execution/helpers'; +import type { IDisposable } from '../../platform/common/types'; import { createDeferred } from '../../platform/common/utils/async'; +import { dispose } from '../../platform/common/utils/lifecycle'; import { uuidUtils } from '../../platform/common/uuid'; import type { Pocket } from '../../platform/deepnote/pocket'; import { logger } from '../../platform/logging'; @@ -59,6 +63,17 @@ export function serializeNotebookContext({ cells }: { cells: NotebookCell[] }): return serializeNotebookContextFromBlocks({ blocks, notebookName: null }); } +export function getOpenAiApiKey(): string { + const config = workspace.getConfiguration('deepnote'); + const key = config.get('agent.openAiApiKey', ''); + + if (!key) { + throw new Error('deepnote.agent.openAiApiKey is not set. Configure it in VS Code settings.'); + } + + return key; +} + export interface ExecuteAgentCellOptions { executeAgentBlockFn?: typeof executeAgentBlock; } @@ -107,11 +122,7 @@ export async function executeAgentCell( cells: cell.notebook.getCells().filter((c) => c.index !== cell.index) }); - // eslint-disable-next-line local-rules/dont-use-process - const openAiToken = process.env.OPENAI_API_KEY; - if (openAiToken == null) { - throw new Error('OPENAI_API_KEY is not set'); - } + const openAiToken = getOpenAiApiKey(); const context: AgentBlockContext = { openAiToken, @@ -125,7 +136,7 @@ export async function executeAgentCell( const cellIndex = await insertEphemeralCell(cell.notebook, cell.index, agentBlock.id, 'code', code); const insertedCell = cell.notebook.cellAt(cellIndex); - const { success } = await executeEphemeralCell(insertedCell); + const { success } = await executeEphemeralCell(insertedCell, execution.token); return success ? { success } : { success: false, error: new Error('Ephemeral cell execution failed') }; }, onLog: (message: string) => { @@ -248,15 +259,27 @@ async function insertEphemeralCell( const EPHEMERAL_CELL_EXECUTION_TIMEOUT_MS = 5 * 60 * 1000; export async function executeEphemeralCell( - cell: NotebookCell + cell: NotebookCell, + token?: CancellationToken ): Promise<{ success: boolean; outputs: unknown[]; executionCount: number | null }> { const completionDeferred = createDeferred(); + const disposables: IDisposable[] = []; + + disposables.push( + notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { + if (e.cell === cell && e.state === NotebookCellExecutionState.Idle) { + completionDeferred.resolve(); + } + }) + ); - const disposable = notebookCellExecutions.onDidChangeNotebookCellExecutionState((e) => { - if (e.cell === cell && e.state === NotebookCellExecutionState.Idle) { - completionDeferred.resolve(); + if (token) { + if (token.isCancellationRequested) { + completionDeferred.reject(new CancellationError()); + } else { + disposables.push(token.onCancellationRequested(() => completionDeferred.reject(new CancellationError()))); } - }); + } const timeout = setTimeout(() => { completionDeferred.reject(new Error('Ephemeral cell execution timed out')); @@ -273,7 +296,7 @@ export async function executeEphemeralCell( await completionDeferred.promise; return { - success: cell.executionSummary?.success !== false, + success: cell.executionSummary?.success === true, outputs: cell.outputs.map(translateCellDisplayOutput), executionCount: cell.executionSummary?.executionOrder ?? null }; @@ -284,7 +307,7 @@ export async function executeEphemeralCell( executionCount: null }; } finally { - disposable.dispose(); + dispose(disposables); clearTimeout(timeout); } } diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index b84bb6286c..51b9d089b2 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -1,17 +1,20 @@ import { expect } from 'chai'; import * as sinon from 'sinon'; -import { anything, capture, reset, when } from 'ts-mockito'; -import { NotebookCellOutput, NotebookCellOutputItem, NotebookController } from 'vscode'; +import { anything, capture, instance, mock, reset, when } from 'ts-mockito'; +import { + CancellationTokenSource, + NotebookCellOutput, + NotebookCellOutputItem, + NotebookController, + WorkspaceConfiguration +} from 'vscode'; import type { AgentBlock } from '@deepnote/blocks'; import type { AgentBlockContext, AgentBlockResult } from '@deepnote/runtime-core'; -import { - NotebookCellExecutionState, - notebookCellExecutions -} from '../../platform/notebooks/cellExecutionStateService'; +import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; -import { executeAgentCell, executeEphemeralCell, isAgentCell } from './agentCellExecutionHandler'; +import { executeAgentCell, executeEphemeralCell, getOpenAiApiKey, isAgentCell } from './agentCellExecutionHandler'; import { createMockCell } from './deepnoteTestHelpers'; suite('AgentCellExecutionHandler', () => { @@ -47,6 +50,24 @@ suite('AgentCellExecutionHandler', () => { }); }); + suite('getOpenAiApiKey', () => { + test('returns key when configured', () => { + const mockConfig = mock(); + when(mockConfig.get('agent.openAiApiKey', '')).thenReturn('test-key'); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + + expect(getOpenAiApiKey()).to.equal('test-key'); + }); + + test('throws when key is not set', () => { + const mockConfig = mock(); + when(mockConfig.get('agent.openAiApiKey', '')).thenReturn(''); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + + expect(() => getOpenAiApiKey()).to.throw('deepnote.agent.openAiApiKey is not set'); + }); + }); + suite('executeAgentCell', () => { let mockExecution: { appendOutput: sinon.SinonStub; @@ -58,11 +79,11 @@ suite('AgentCellExecutionHandler', () => { }; let mockController: NotebookController; let executeAgentBlockStub: sinon.SinonStub; - let savedOpenAiKey: string | undefined; setup(() => { - savedOpenAiKey = process.env.OPENAI_API_KEY; - process.env.OPENAI_API_KEY = 'test-key'; + const mockConfig = mock(); + when(mockConfig.get('agent.openAiApiKey', '')).thenReturn('test-key'); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); mockExecution = { appendOutput: sinon.stub().resolves(), @@ -80,14 +101,6 @@ suite('AgentCellExecutionHandler', () => { executeAgentBlockStub = sinon.stub().resolves({ finalOutput: 'done' } as AgentBlockResult); }); - teardown(() => { - if (savedOpenAiKey !== undefined) { - process.env.OPENAI_API_KEY = savedOpenAiKey; - } else { - delete process.env.OPENAI_API_KEY; - } - }); - function createAgentCell(text: string = 'Test prompt') { return createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } }, @@ -243,6 +256,24 @@ suite('AgentCellExecutionHandler', () => { const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); expect(text).to.include('[Agent] Planning next steps...'); }); + + test('ends with failure and writes error when API key is not set', async () => { + const mockConfig = mock(); + when(mockConfig.get('agent.openAiApiKey', '')).thenReturn(''); + when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + + const cell = createAgentCell(); + + await executeAgentCell(cell, mockController, { executeAgentBlockFn: executeAgentBlockStub }); + + expect(mockExecution.end.calledOnce).to.be.true; + expect(mockExecution.end.firstCall.args[0]).to.be.false; + expect(mockExecution.appendOutput.calledOnce).to.be.true; + + const outputs = mockExecution.appendOutput.firstCall.args[0] as NotebookCellOutput[]; + const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); + expect(text).to.include('deepnote.agent.openAiApiKey is not set'); + }); }); suite('executeEphemeralCell', () => { @@ -275,5 +306,21 @@ suite('AgentCellExecutionHandler', () => { document: cell.notebook.uri }); }); + + test('returns success false immediately when token is pre-cancelled', async () => { + const cell = createMockCell({ index: 0 }); + const tokenSource = new CancellationTokenSource(); + tokenSource.cancel(); + + when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); + + const result = await executeEphemeralCell(cell, tokenSource.token); + + expect(result).to.deep.equal({ + success: false, + outputs: [], + executionCount: null + }); + }); }); }); diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index 8d4ba8eba4..75b3cc1e4e 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -14,14 +14,17 @@ import { workspace } from 'vscode'; import { injectable } from 'inversify'; +import { z } from 'zod'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import type { Pocket } from '../../platform/deepnote/pocket'; +import { logger } from '../../platform/logging'; const DEFAULT_MAX_ITERATIONS = 20; const MIN_ITERATIONS = 1; const MAX_ITERATIONS = 100; const AGENT_MODEL_OPTIONS = ['auto', 'gpt-4o', 'sonnet']; +const MaxIterationsSchema = z.coerce.number().int().min(MIN_ITERATIONS); /** * Provides status bar items for agent cells showing the block type indicator, @@ -67,7 +70,9 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv } public dispose(): void { - this.disposables.forEach((d) => d.dispose()); + for (const disposable of this.disposables) { + disposable.dispose(); + } } public provideCellStatusBarItems( @@ -141,8 +146,16 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv private getMaxIterations(metadata: Record | undefined): number { const value = metadata?.deepnote_max_iterations; - if (typeof value === 'number' && Number.isInteger(value) && value >= MIN_ITERATIONS) { - return value; + const result = MaxIterationsSchema.safeParse(value); + + if (result.success) { + return result.data; + } + + if (value !== undefined) { + logger.debug( + `getMaxIterations: invalid value ${JSON.stringify(value)}, using default ${DEFAULT_MAX_ITERATIONS}` + ); } return DEFAULT_MAX_ITERATIONS; diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts index e5397a1948..62adc562cc 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.unit.test.ts @@ -214,6 +214,66 @@ suite('AgentCellStatusBarProvider', () => { expect(items[2].text).to.include('Max iterations: 20'); }); + test('Should display default when max iterations is negative', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: -5 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + }); + + test('Should display 1 when max iterations is MIN_ITERATIONS boundary', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: 1 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 1'); + }); + + test('Should display 100 when max iterations is at upper bound', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: 100 + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 100'); + }); + + test('Should display default when max iterations is null', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: null + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + }); + + test('Should display default when max iterations is boolean', () => { + const cell = createMockCell({ + metadata: { + __deepnotePocket: { type: 'agent' }, + deepnote_max_iterations: true + } + }); + const items = provider.provideCellStatusBarItems(cell, mockToken)!; + + expect(items[2].text).to.include('Max iterations: 20'); + }); + test('Should have set max iterations command', () => { const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); const items = provider.provideCellStatusBarItems(cell, mockToken)!; From f7bec65bf12b7a41abcdafaa9f3eb55d5d9f7437 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 18 Mar 2026 10:27:20 +0000 Subject: [PATCH 7/8] Add Agent OpenAI API key extension configuration --- package.json | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/package.json b/package.json index 1b31582e7b..dab239bb98 100644 --- a/package.json +++ b/package.json @@ -1638,6 +1638,12 @@ "type": "object", "title": "Deepnote", "properties": { + "deepnote.agent.openAiApiKey": { + "type": "string", + "default": "", + "description": "OpenAI API key for agent cell execution", + "scope": "application" + }, "deepnote.domain": { "type": "string", "default": "deepnote.com", From ea715e798b02090841f6f0a64b7ab23b43e52215 Mon Sep 17 00:00:00 2001 From: tomas Date: Wed, 18 Mar 2026 17:32:49 +0000 Subject: [PATCH 8/8] Implement OpenAI API key management in Deepnote - Added commands to set and clear the OpenAI API key, enhancing user interaction. - Introduced a new `deepnoteSecretStore` module for managing secrets, including functions to get, set, and clear the OpenAI API key. - Updated `agentCellExecutionHandler` to utilize the new secret management functions, improving error handling when the API key is not set. - Enhanced unit tests to cover the new secret management functionality and ensure robust error handling. --- package.json | 16 +- package.nls.json | 2 + .../deepnote/agentCellExecutionHandler.ts | 24 +- .../agentCellExecutionHandler.unit.test.ts | 110 ++++++-- .../deepnote/agentCellStatusBarProvider.ts | 19 +- src/notebooks/deepnote/deepnoteSecretStore.ts | 122 +++++++++ .../deepnote/deepnoteSecretStore.unit.test.ts | 241 ++++++++++++++++++ .../ephemeralCellDecorationProvider.ts | 4 +- 8 files changed, 487 insertions(+), 51 deletions(-) create mode 100644 src/notebooks/deepnote/deepnoteSecretStore.ts create mode 100644 src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts diff --git a/package.json b/package.json index dab239bb98..d800875f53 100644 --- a/package.json +++ b/package.json @@ -341,6 +341,16 @@ "title": "%deepnote.command.manageAccessToKernels%", "category": "Jupyter" }, + { + "command": "deepnote.setOpenAiApiKey", + "title": "%deepnote.command.setOpenAiApiKey%", + "category": "Deepnote" + }, + { + "command": "deepnote.clearOpenAiApiKey", + "title": "%deepnote.command.clearOpenAiApiKey%", + "category": "Deepnote" + }, { "command": "dataScience.ClearUserProviderJupyterServerCache", "title": "%deepnote.command.dataScience.clearUserProviderJupyterServerCache.title%", @@ -1638,12 +1648,6 @@ "type": "object", "title": "Deepnote", "properties": { - "deepnote.agent.openAiApiKey": { - "type": "string", - "default": "", - "description": "OpenAI API key for agent cell execution", - "scope": "application" - }, "deepnote.domain": { "type": "string", "default": "deepnote.com", diff --git a/package.nls.json b/package.nls.json index 35ee95ae66..07f3b2ba4a 100644 --- a/package.nls.json +++ b/package.nls.json @@ -116,6 +116,8 @@ "deepnote.command.deepnote.openOutlineView.title": "Show Table Of Contents (Outline View)", "deepnote.command.deepnote.openOutlineView.shorttitle": "Outline", "deepnote.command.manageAccessToKernels": "Manage Access To Jupyter Kernels", + "deepnote.command.setOpenAiApiKey": "Set OpenAI API Key", + "deepnote.command.clearOpenAiApiKey": "Clear OpenAI API Key", "deepnote.commandPalette.deepnote.replayPylanceLog.title": "Replay Pylance Log", "deepnote.notebookRenderer.IPyWidget.displayName": "Jupyter IPyWidget Renderer", "deepnote.notebookRenderer.Error.displayName": "Jupyter Error Renderer", diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.ts b/src/notebooks/deepnote/agentCellExecutionHandler.ts index e8a1268194..3ed36cf0e2 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.ts @@ -31,6 +31,11 @@ import { logger } from '../../platform/logging'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; import { generateBlockId, generateSortingKey, isEphemeralCell } from './dataConversionUtils'; import { DeepnoteDataConverter } from './deepnoteDataConverter'; +import { getOrPromptOpenAiApiKey } from './deepnoteSecretStore'; + +export async function getOpenAiApiKey(): Promise { + return getOrPromptOpenAiApiKey(); +} export function isAgentCell(cell: NotebookCell): boolean { const pocket = cell.metadata?.__deepnotePocket as Pocket | undefined; @@ -63,17 +68,6 @@ export function serializeNotebookContext({ cells }: { cells: NotebookCell[] }): return serializeNotebookContextFromBlocks({ blocks, notebookName: null }); } -export function getOpenAiApiKey(): string { - const config = workspace.getConfiguration('deepnote'); - const key = config.get('agent.openAiApiKey', ''); - - if (!key) { - throw new Error('deepnote.agent.openAiApiKey is not set. Configure it in VS Code settings.'); - } - - return key; -} - export interface ExecuteAgentCellOptions { executeAgentBlockFn?: typeof executeAgentBlock; } @@ -122,7 +116,7 @@ export async function executeAgentCell( cells: cell.notebook.getCells().filter((c) => c.index !== cell.index) }); - const openAiToken = getOpenAiApiKey(); + const openAiToken = await getOpenAiApiKey(); const context: AgentBlockContext = { openAiToken, @@ -139,12 +133,6 @@ export async function executeAgentCell( const { success } = await executeEphemeralCell(insertedCell, execution.token); return success ? { success } : { success: false, error: new Error('Ephemeral cell execution failed') }; }, - onLog: (message: string) => { - logger.info('Agent log', message); - // accumulated += message; - // TODO: replaceOutputItems is Async function - // execution.replaceOutputItems(NotebookCellOutputItem.text(accumulated), output); - }, onAgentEvent: async (event: AgentStreamEvent) => { logger.info('Agent event', JSON.stringify(event)); if (lastAgentEventType != null && lastAgentEventType !== event.type) { diff --git a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts index 51b9d089b2..b124742244 100644 --- a/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts +++ b/src/notebooks/deepnote/agentCellExecutionHandler.unit.test.ts @@ -3,21 +3,32 @@ import * as sinon from 'sinon'; import { anything, capture, instance, mock, reset, when } from 'ts-mockito'; import { CancellationTokenSource, + Disposable, + EventEmitter, + ExtensionMode, NotebookCellOutput, NotebookCellOutputItem, NotebookController, - WorkspaceConfiguration + SecretStorage, + SecretStorageChangeEvent } from 'vscode'; import type { AgentBlock } from '@deepnote/blocks'; import type { AgentBlockContext, AgentBlockResult } from '@deepnote/runtime-core'; +import type { IDisposable } from '../../platform/common/types'; +import { IExtensionContext } from '../../platform/common/types'; import { NotebookCellExecutionState, notebookCellExecutions } from '../../platform/notebooks/cellExecutionStateService'; +import { dispose } from '../../platform/common/utils/lifecycle'; import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; +import { ServiceContainer } from '../../platform/ioc/container'; import { executeAgentCell, executeEphemeralCell, getOpenAiApiKey, isAgentCell } from './agentCellExecutionHandler'; import { createMockCell } from './deepnoteTestHelpers'; suite('AgentCellExecutionHandler', () => { + const secretStorage = new Map(); + let disposables: IDisposable[] = []; + suite('isAgentCell', () => { test('returns true for cell with agent pocket type', () => { const cell = createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } } }); @@ -51,20 +62,47 @@ suite('AgentCellExecutionHandler', () => { }); suite('getOpenAiApiKey', () => { - test('returns key when configured', () => { - const mockConfig = mock(); - when(mockConfig.get('agent.openAiApiKey', '')).thenReturn('test-key'); - when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + setup(() => { + secretStorage.clear(); + const context = mock(); + const secrets = mock(); + const onDidChangeSecrets = new EventEmitter(); + const serviceContainer = mock(); + sinon.stub(ServiceContainer, 'instance').get(() => instance(serviceContainer)); + when(serviceContainer.get(IExtensionContext)).thenReturn(instance(context)); + when(context.extensionMode).thenReturn(ExtensionMode.Production); + when(context.secrets).thenReturn(instance(secrets)); + when(secrets.onDidChange).thenReturn(onDidChangeSecrets.event); + when(secrets.get(anything())).thenCall((key: string) => Promise.resolve(secretStorage.get(key))); + when(secrets.store(anything(), anything())).thenCall((key: string, value: string) => { + secretStorage.set(key, value); + + return Promise.resolve(); + }); + disposables.push(new Disposable(() => sinon.restore())); + }); - expect(getOpenAiApiKey()).to.equal('test-key'); + teardown(() => { + disposables = dispose(disposables); }); - test('throws when key is not set', () => { - const mockConfig = mock(); - when(mockConfig.get('agent.openAiApiKey', '')).thenReturn(''); - when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + test('returns key when configured', async () => { + secretStorage.set('openAiApiKey', 'test-key'); + + const key = await getOpenAiApiKey(); - expect(() => getOpenAiApiKey()).to.throw('deepnote.agent.openAiApiKey is not set'); + expect(key).to.equal('test-key'); + }); + + test('throws when key is not set', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + try { + await getOpenAiApiKey(); + expect.fail('Should have thrown'); + } catch (e) { + expect((e as Error).message).to.include('OpenAI API key is not set'); + } }); }); @@ -81,9 +119,24 @@ suite('AgentCellExecutionHandler', () => { let executeAgentBlockStub: sinon.SinonStub; setup(() => { - const mockConfig = mock(); - when(mockConfig.get('agent.openAiApiKey', '')).thenReturn('test-key'); - when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + secretStorage.clear(); + secretStorage.set('openAiApiKey', 'test-key'); + const context = mock(); + const secrets = mock(); + const onDidChangeSecrets = new EventEmitter(); + const serviceContainer = mock(); + sinon.stub(ServiceContainer, 'instance').get(() => instance(serviceContainer)); + when(serviceContainer.get(IExtensionContext)).thenReturn(instance(context)); + when(context.extensionMode).thenReturn(ExtensionMode.Production); + when(context.secrets).thenReturn(instance(secrets)); + when(secrets.onDidChange).thenReturn(onDidChangeSecrets.event); + when(secrets.get(anything())).thenCall((key: string) => Promise.resolve(secretStorage.get(key))); + when(secrets.store(anything(), anything())).thenCall((key: string, value: string) => { + secretStorage.set(key, value); + + return Promise.resolve(); + }); + disposables.push(new Disposable(() => sinon.restore())); mockExecution = { appendOutput: sinon.stub().resolves(), @@ -101,6 +154,10 @@ suite('AgentCellExecutionHandler', () => { executeAgentBlockStub = sinon.stub().resolves({ finalOutput: 'done' } as AgentBlockResult); }); + teardown(() => { + disposables = dispose(disposables); + }); + function createAgentCell(text: string = 'Test prompt') { return createMockCell({ metadata: { __deepnotePocket: { type: 'agent' } }, @@ -258,9 +315,8 @@ suite('AgentCellExecutionHandler', () => { }); test('ends with failure and writes error when API key is not set', async () => { - const mockConfig = mock(); - when(mockConfig.get('agent.openAiApiKey', '')).thenReturn(''); - when(mockedVSCodeNamespaces.workspace.getConfiguration('deepnote')).thenReturn(instance(mockConfig)); + secretStorage.clear(); + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); const cell = createAgentCell(); @@ -272,7 +328,7 @@ suite('AgentCellExecutionHandler', () => { const outputs = mockExecution.appendOutput.firstCall.args[0] as NotebookCellOutput[]; const text = Buffer.from(outputs[0].items[0].data).toString('utf-8'); - expect(text).to.include('deepnote.agent.openAiApiKey is not set'); + expect(text).to.include('OpenAI API key is not set'); }); }); @@ -314,13 +370,17 @@ suite('AgentCellExecutionHandler', () => { when(mockedVSCodeNamespaces.commands.executeCommand(anything(), anything())).thenResolve(); - const result = await executeEphemeralCell(cell, tokenSource.token); - - expect(result).to.deep.equal({ - success: false, - outputs: [], - executionCount: null - }); + try { + const result = await executeEphemeralCell(cell, tokenSource.token); + + expect(result).to.deep.equal({ + success: false, + outputs: [], + executionCount: null + }); + } finally { + tokenSource.dispose(); + } }); }); }); diff --git a/src/notebooks/deepnote/agentCellStatusBarProvider.ts b/src/notebooks/deepnote/agentCellStatusBarProvider.ts index 75b3cc1e4e..0bbc73f3b7 100644 --- a/src/notebooks/deepnote/agentCellStatusBarProvider.ts +++ b/src/notebooks/deepnote/agentCellStatusBarProvider.ts @@ -19,12 +19,13 @@ import { z } from 'zod'; import { IExtensionSyncActivationService } from '../../platform/activation/types'; import type { Pocket } from '../../platform/deepnote/pocket'; import { logger } from '../../platform/logging'; +import { clearOpenAiApiKey, promptForOpenAiApiKey } from './deepnoteSecretStore'; const DEFAULT_MAX_ITERATIONS = 20; const MIN_ITERATIONS = 1; const MAX_ITERATIONS = 100; const AGENT_MODEL_OPTIONS = ['auto', 'gpt-4o', 'sonnet']; -const MaxIterationsSchema = z.coerce.number().int().min(MIN_ITERATIONS); +const MaxIterationsSchema = z.coerce.number().int().min(MIN_ITERATIONS).max(MAX_ITERATIONS); /** * Provides status bar items for agent cells showing the block type indicator, @@ -66,6 +67,22 @@ export class AgentCellStatusBarProvider implements NotebookCellStatusBarItemProv }) ); + this.disposables.push( + commands.registerCommand('deepnote.setOpenAiApiKey', async () => { + const key = await promptForOpenAiApiKey(); + if (key) { + void window.showInformationMessage(l10n.t('OpenAI API key has been saved.')); + } + }) + ); + + this.disposables.push( + commands.registerCommand('deepnote.clearOpenAiApiKey', async () => { + await clearOpenAiApiKey(); + void window.showInformationMessage(l10n.t('OpenAI API key has been cleared.')); + }) + ); + this.disposables.push(this._onDidChangeCellStatusBarItems); } diff --git a/src/notebooks/deepnote/deepnoteSecretStore.ts b/src/notebooks/deepnote/deepnoteSecretStore.ts new file mode 100644 index 0000000000..fadf11bd42 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteSecretStore.ts @@ -0,0 +1,122 @@ +import { ExtensionMode, l10n, window } from 'vscode'; + +import { ServiceContainer } from '../../platform/ioc/container'; +import { IExtensionContext } from '../../platform/common/types'; + +export interface SecretPromptOptions { + prompt: string; + placeHolder?: string; + password?: boolean; +} + +function getContext(): IExtensionContext | null { + const context = ServiceContainer.instance.get(IExtensionContext); + + if (context.extensionMode === ExtensionMode.Test) { + return null; + } + + return context; +} + +export async function getSecret(key: string): Promise { + const context = getContext(); + + if (!context) { + return undefined; + } + + const value = await context.secrets.get(key); + + return value && value.length > 0 ? value : undefined; +} + +export async function setSecret(key: string, value: string): Promise { + const context = getContext(); + + if (!context) { + return; + } + + await context.secrets.store(key, value); +} + +export async function clearSecret(key: string): Promise { + const context = getContext(); + + if (!context) { + return; + } + + await context.secrets.delete(key); +} + +export async function promptForSecret(key: string, options: SecretPromptOptions): Promise { + const input = await window.showInputBox({ + prompt: options.prompt, + placeHolder: options.placeHolder, + password: options.password ?? true, + ignoreFocusOut: true + }); + + if (!input || input.trim().length === 0) { + return undefined; + } + + const trimmed = input.trim(); + await setSecret(key, trimmed); + + return trimmed; +} + +export async function getOrPromptSecret( + key: string, + options: SecretPromptOptions, + errorMessage: string +): Promise { + let value = await getSecret(key); + + if (!value) { + value = await promptForSecret(key, options); + } + + if (!value) { + throw new Error(errorMessage); + } + + return value; +} + +// OpenAI API key - specific wrappers + +const OPENAI_API_KEY = 'openAiApiKey'; + +const OPENAI_PROMPT_OPTIONS: SecretPromptOptions = { + prompt: l10n.t('Enter your OpenAI API key'), + placeHolder: l10n.t('sk-...'), + password: true +}; + +export async function getOpenAiApiKey(): Promise { + return getSecret(OPENAI_API_KEY); +} + +export async function setOpenAiApiKey(key: string): Promise { + return setSecret(OPENAI_API_KEY, key); +} + +export async function clearOpenAiApiKey(): Promise { + return clearSecret(OPENAI_API_KEY); +} + +export async function promptForOpenAiApiKey(): Promise { + return promptForSecret(OPENAI_API_KEY, OPENAI_PROMPT_OPTIONS); +} + +export async function getOrPromptOpenAiApiKey(): Promise { + return getOrPromptSecret( + OPENAI_API_KEY, + OPENAI_PROMPT_OPTIONS, + l10n.t('OpenAI API key is not set. Use the command "Deepnote: Set OpenAI API Key" to configure it.') + ); +} diff --git a/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts new file mode 100644 index 0000000000..e99bafc638 --- /dev/null +++ b/src/notebooks/deepnote/deepnoteSecretStore.unit.test.ts @@ -0,0 +1,241 @@ +import { expect } from 'chai'; +import * as sinon from 'sinon'; +import { anything, instance, mock, when } from 'ts-mockito'; +import { EventEmitter, ExtensionMode, SecretStorage, SecretStorageChangeEvent } from 'vscode'; + +import { IExtensionContext } from '../../platform/common/types'; +import { ServiceContainer } from '../../platform/ioc/container'; +import { + clearOpenAiApiKey, + clearSecret, + getOpenAiApiKey, + getOrPromptOpenAiApiKey, + getOrPromptSecret, + getSecret, + promptForOpenAiApiKey, + promptForSecret, + setOpenAiApiKey, + setSecret +} from './deepnoteSecretStore'; +import { mockedVSCodeNamespaces } from '../../test/vscode-mock'; + +suite('deepnoteSecretStore', () => { + const secretStorage = new Map(); + let context: IExtensionContext; + let secrets: SecretStorage; + let onDidChangeSecrets: EventEmitter; + + setup(() => { + secretStorage.clear(); + context = mock(); + secrets = mock(); + onDidChangeSecrets = new EventEmitter(); + + const serviceContainer = mock(); + sinon.stub(ServiceContainer, 'instance').get(() => instance(serviceContainer)); + when(serviceContainer.get(IExtensionContext)).thenReturn(instance(context)); + when(context.extensionMode).thenReturn(ExtensionMode.Production); + when(context.secrets).thenReturn(instance(secrets)); + when(secrets.onDidChange).thenReturn(onDidChangeSecrets.event); + when(secrets.get(anything())).thenCall((key: string) => Promise.resolve(secretStorage.get(key))); + when(secrets.store(anything(), anything())).thenCall((key: string, value: string) => { + secretStorage.set(key, value); + onDidChangeSecrets.fire({ key }); + + return Promise.resolve(); + }); + when(secrets.delete(anything())).thenCall((key: string) => { + secretStorage.delete(key); + + return Promise.resolve(); + }); + }); + + teardown(() => { + sinon.restore(); + }); + + suite('generic getSecret', () => { + test('returns value when stored', async () => { + secretStorage.set('customKey', 'custom-value'); + + const value = await getSecret('customKey'); + + expect(value).to.equal('custom-value'); + }); + + test('returns undefined when not set', async () => { + const value = await getSecret('customKey'); + + expect(value).to.be.undefined; + }); + + test('returns undefined when value is empty string', async () => { + secretStorage.set('customKey', ''); + + const value = await getSecret('customKey'); + + expect(value).to.be.undefined; + }); + }); + + suite('generic setSecret', () => { + test('stores value in secrets', async () => { + await setSecret('customKey', 'custom-value'); + + expect(secretStorage.get('customKey')).to.equal('custom-value'); + }); + }); + + suite('generic clearSecret', () => { + test('deletes value from secrets', async () => { + secretStorage.set('customKey', 'custom-value'); + + await clearSecret('customKey'); + + expect(secretStorage.has('customKey')).to.be.false; + }); + }); + + suite('generic promptForSecret', () => { + test('stores and returns value when user enters input', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('user-input')); + + const value = await promptForSecret('customKey', { + prompt: 'Enter value', + placeHolder: 'placeholder', + password: false + }); + + expect(value).to.equal('user-input'); + expect(secretStorage.get('customKey')).to.equal('user-input'); + }); + + test('returns undefined when user cancels', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const value = await promptForSecret('customKey', { prompt: 'Enter value' }); + + expect(value).to.be.undefined; + }); + }); + + suite('generic getOrPromptSecret', () => { + test('returns value when present in store', async () => { + secretStorage.set('customKey', 'stored-value'); + + const value = await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); + + expect(value).to.equal('stored-value'); + }); + + test('throws when value missing and user cancels prompt', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + try { + await getOrPromptSecret('customKey', { prompt: 'Enter value' }, 'Value is required'); + expect.fail('Should have thrown'); + } catch (e) { + expect((e as Error).message).to.equal('Value is required'); + } + }); + }); + + suite('getOpenAiApiKey', () => { + test('returns key when stored', async () => { + secretStorage.set('openAiApiKey', 'test-key'); + + const key = await getOpenAiApiKey(); + + expect(key).to.equal('test-key'); + }); + + test('returns undefined when not set', async () => { + const key = await getOpenAiApiKey(); + + expect(key).to.be.undefined; + }); + + test('returns undefined when key is empty string', async () => { + secretStorage.set('openAiApiKey', ''); + + const key = await getOpenAiApiKey(); + + expect(key).to.be.undefined; + }); + }); + + suite('setOpenAiApiKey', () => { + test('stores key in secrets', async () => { + await setOpenAiApiKey('my-api-key'); + + expect(secretStorage.get('openAiApiKey')).to.equal('my-api-key'); + }); + }); + + suite('clearOpenAiApiKey', () => { + test('deletes key from secrets', async () => { + secretStorage.set('openAiApiKey', 'test-key'); + + await clearOpenAiApiKey(); + + expect(secretStorage.has('openAiApiKey')).to.be.false; + }); + }); + + suite('promptForOpenAiApiKey', () => { + test('stores and returns key when user enters value', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('sk-abc123')); + + const key = await promptForOpenAiApiKey(); + + expect(key).to.equal('sk-abc123'); + expect(secretStorage.get('openAiApiKey')).to.equal('sk-abc123'); + }); + + test('returns undefined when user cancels', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + const key = await promptForOpenAiApiKey(); + + expect(key).to.be.undefined; + }); + + test('returns undefined when user enters empty string', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(' ')); + + const key = await promptForOpenAiApiKey(); + + expect(key).to.be.undefined; + }); + }); + + suite('getOrPromptOpenAiApiKey', () => { + test('returns key when present in store', async () => { + secretStorage.set('openAiApiKey', 'stored-key'); + + const key = await getOrPromptOpenAiApiKey(); + + expect(key).to.equal('stored-key'); + }); + + test('prompts and returns key when missing', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve('prompted-key')); + + const key = await getOrPromptOpenAiApiKey(); + + expect(key).to.equal('prompted-key'); + }); + + test('throws when key missing and user cancels prompt', async () => { + when(mockedVSCodeNamespaces.window.showInputBox(anything())).thenReturn(Promise.resolve(undefined)); + + try { + await getOrPromptOpenAiApiKey(); + expect.fail('Should have thrown'); + } catch (e) { + expect((e as Error).message).to.include('OpenAI API key is not set'); + } + }); + }); +}); diff --git a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts index 36b5d24053..19ca8b76fc 100644 --- a/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts +++ b/src/notebooks/deepnote/ephemeralCellDecorationProvider.ts @@ -67,7 +67,9 @@ export class EphemeralCellDecorationProvider implements IExtensionSyncActivation } public dispose(): void { - this.disposables.forEach((d) => d.dispose()); + for (const disposable of this.disposables) { + disposable.dispose(); + } } private findCellForEditor(editor: TextEditor): NotebookCell | undefined {