feat: Add a command for adding chart blocks

package.json

@@ -129,6 +129,12 @@
                 "category": "Deepnote",
                 "icon": "$(graph)"
             },
+            {
+                "command": "deepnote.addChartBlock",
+                "title": "%deepnote.commands.addChartBlock.title%",
+                "category": "Deepnote",
+                "icon": "$(graph-line)"
+            },
             {
                 "command": "deepnote.addInputTextBlock",
                 "title": "%deepnote.commands.addInputTextBlock.title%",
@@ -850,6 +856,51 @@
                     "group": "navigation@1",
                     "when": "notebookType == 'deepnote'"
                 },
+                {
+                    "command": "deepnote.addChartBlock",
+                    "group": "navigation@2",
+                    "when": "notebookType == 'deepnote'"
+                },
+                {
+                    "command": "deepnote.addBigNumberChartBlock",
+                    "group": "navigation@3",
+                    "when": "notebookType == 'deepnote'"
+                },
+                {
+                    "command": "deepnote.addInputTextBlock",
+                    "group": "navigation@4",
+                    "when": "notebookType == 'deepnote'"
+                },
+                {
+                    "command": "deepnote.addInputTextareaBlock",
+                    "group": "navigation@5",
+                    "when": "notebookType == 'deepnote'"
+                },
+                {
+                    "command": "deepnote.addInputSelectBlock",
+                    "group": "navigation@6",
+                    "when": "notebookType == 'deepnote'"
+                },
+                {
+                    "command": "deepnote.addInputSliderBlock",
+                    "group": "navigation@7",
+                    "when": "notebookType == 'deepnote'"
+                },
+                {
+                    "command": "deepnote.addInputCheckboxBlock",
+                    "group": "navigation@8",
+                    "when": "notebookType == 'deepnote'"
+                },
+                {
+                    "command": "deepnote.addInputDateBlock",
+                    "group": "navigation@9",
+                    "when": "notebookType == 'deepnote'"
+                },
+                {
+                    "command": "deepnote.addInputDateRangeBlock",
+                    "group": "navigation@10",
+                    "when": "notebookType == 'deepnote'"
+                },
                 {
                     "command": "jupyter.restartkernel",
                     "group": "navigation/execute@5",

package.nls.json

@@ -255,7 +255,8 @@
     "deepnote.commands.importNotebook.title": "Import Notebook",
     "deepnote.commands.importJupyterNotebook.title": "Import Jupyter Notebook",
     "deepnote.commands.addSqlBlock.title": "Add SQL Block",
-    "deepnote.commands.addBigNumberChartBlock.title": "Add Big Number Chart Block",
+    "deepnote.commands.addBigNumberChartBlock.title": "Add Big Number Block",
+    "deepnote.commands.addChartBlock.title": "Add Chart Block",
     "deepnote.commands.addInputTextBlock.title": "Add Input Text Block",
     "deepnote.commands.addInputTextareaBlock.title": "Add Input Textarea Block",
     "deepnote.commands.addInputSelectBlock.title": "Add Input Select Block",

src/commands.ts

@@ -188,6 +188,7 @@ export interface ICommandNameArgumentTypeMapping {
     [DSCommands.ContinueEditSessionInCodespace]: [];
     [DSCommands.AddSqlBlock]: [];
     [DSCommands.AddBigNumberChartBlock]: [];
+    [DSCommands.AddChartBlock]: [];
     [DSCommands.AddInputTextBlock]: [];
     [DSCommands.AddInputTextareaBlock]: [];
     [DSCommands.AddInputSelectBlock]: [];

src/notebooks/deepnote/deepnoteNotebookCommandListener.ts

@@ -17,6 +17,7 @@ import { IExtensionSyncActivationService } from '../../platform/activation/types
 import { IDisposableRegistry } from '../../platform/common/types';
 import { Commands } from '../../platform/common/constants';
 import { chainWithPendingUpdates } from '../../kernels/execution/notebookUpdater';
+import { WrappedError } from '../../platform/errors/types';
 import {
     DeepnoteBigNumberMetadataSchema,
     DeepnoteTextInputMetadataSchema,
@@ -76,7 +77,10 @@ export function getInputBlockMetadata(blockType: InputBlockType, variableName: s
 
 export function safeParseDeepnoteVariableNameFromContentJson(content: string): string | undefined {
     try {
-        const variableNameResult = z.string().safeParse(JSON.parse(content)['deepnote_variable_name']);
+        const parsed = JSON.parse(content);
+        // Chart blocks use 'variable' key, other blocks use 'deepnote_variable_name'
+        const variableName = parsed['variable'] ?? parsed['deepnote_variable_name'];
+        const variableNameResult = z.string().safeParse(variableName);
         return variableNameResult.success ? variableNameResult.data : undefined;
     } catch (error) {
         logger.error('Error parsing deepnote variable name from content JSON', error);
@@ -148,6 +152,7 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation
         this.disposableRegistry.push(
             commands.registerCommand(Commands.AddBigNumberChartBlock, () => this.addBigNumberChartBlock())
         );
+        this.disposableRegistry.push(commands.registerCommand(Commands.AddChartBlock, () => this.addChartBlock()));
         this.disposableRegistry.push(
             commands.registerCommand(Commands.AddInputTextBlock, () => this.addInputBlock('input-text'))
         );
@@ -262,6 +267,61 @@ export class DeepnoteNotebookCommandListener implements IExtensionSyncActivation
         await commands.executeCommand('notebook.cell.edit');
     }
 
+    public async addChartBlock(): Promise<void> {
+        const editor = window.activeNotebookEditor;
+
+        if (!editor) {
+            throw new WrappedError(l10n.t('No active notebook editor found'));
+        }
+
+        const document = editor.notebook;
+        const selection = editor.selection;
+        const insertIndex = selection ? selection.end : document.cellCount;
+
+        const defaultVisualizationSpec = {
+            mark: 'line',
+            $schema: 'https://vega.github.io/schema/vega-lite/v5.json',
+            data: { values: [] },
+            encoding: {
+                x: { field: 'x', type: 'quantitative' },
+                y: { field: 'y', type: 'quantitative' }
+            }
+        };
+
+        const cellContent = {
+            variable: 'df_1',
+            spec: defaultVisualizationSpec,
+            filters: []
+        };
+
+        const metadata = {
+            __deepnotePocket: {
+                type: 'visualization'
+            }
+        };
+
+        const result = await chainWithPendingUpdates(document, (edit) => {
+            const newCell = new NotebookCellData(NotebookCellKind.Code, JSON.stringify(cellContent, null, 2), 'json');
+
+            newCell.metadata = metadata;
+
+            const nbEdit = NotebookEdit.insertCells(insertIndex, [newCell]);
+
+            edit.set(document.uri, [nbEdit]);
+        });
+
+        if (result !== true) {
+            throw new WrappedError(l10n.t('Failed to insert chart block'));
+        }
+
+        const notebookRange = new NotebookRange(insertIndex, insertIndex + 1);
+
+        editor.revealRange(notebookRange, NotebookEditorRevealType.Default);
+        editor.selection = notebookRange;
+
+        await commands.executeCommand('notebook.cell.edit');
+    }
+
     public async addInputBlock(blockType: InputBlockType): Promise<void> {
         const editor = window.activeNotebookEditor;
         if (!editor) {

src/notebooks/deepnote/deepnoteNotebookCommandListener.unit.test.ts

@@ -21,6 +21,7 @@ import {
 import { IDisposable } from '../../platform/common/types';
 import * as notebookUpdater from '../../kernels/execution/notebookUpdater';
 import { createMockedNotebookDocument } from '../../test/datascience/editor-integration/helpers';
+import { WrappedError } from '../../platform/errors/types';
 import { DATAFRAME_SQL_INTEGRATION_ID } from '../../platform/notebooks/deepnote/integrationTypes';
 
 suite('DeepnoteNotebookCommandListener', () => {
@@ -976,5 +977,191 @@ suite('DeepnoteNotebookCommandListener', () => {
                 );
             });
         });
+
+        suite('addChartBlock', () => {
+            test('should add chart block at the end when no selection exists', async () => {
+                // Setup mocks
+                const { editor, document } = createMockEditor([], undefined);
+                const { chainStub, executeCommandStub, getCapturedNotebookEdits } =
+                    mockNotebookUpdateAndExecute(editor);
+
+                // Call the method
+                await commandListener.addChartBlock();
+
+                const capturedNotebookEdits = getCapturedNotebookEdits();
+
+                // Verify chainWithPendingUpdates was called
+                assert.isTrue(chainStub.calledOnce, 'chainWithPendingUpdates should be called once');
+                assert.equal(chainStub.firstCall.args[0], document, 'Should be called with correct document');
+
+                // Verify the edits were captured
+                assert.isNotNull(capturedNotebookEdits, 'Notebook edits should be captured');
+                assert.isDefined(capturedNotebookEdits, 'Notebook edits should be defined');
+
+                const editsArray = capturedNotebookEdits!;
+                assert.equal(editsArray.length, 1, 'Should have one notebook edit');
+
+                const notebookEdit = editsArray[0] as any;
+                assert.equal(notebookEdit.newCells.length, 1, 'Should insert one cell');
+
+                const newCell = notebookEdit.newCells[0];
+                assert.equal(newCell.kind, NotebookCellKind.Code, 'Should be a code cell');
+                assert.equal(newCell.languageId, 'json', 'Should have json language');
+
+                // Verify cell content is valid JSON with correct structure
+                const content = JSON.parse(newCell.value);
+                assert.equal(content.variable, 'df_1', 'Should have correct variable name');
+                assert.property(content, 'spec', 'Should have spec property');
+                assert.property(content, 'filters', 'Should have filters property');
+
+                // Verify the spec has the correct Vega-Lite structure
+                assert.equal(content.spec.mark, 'line', 'Should be a line chart');
+                assert.equal(
+                    content.spec.$schema,
+                    'https://vega.github.io/schema/vega-lite/v5.json',
+                    'Should have Vega-Lite schema'
+                );
+                assert.deepStrictEqual(content.spec.data, { values: [] }, 'Should have empty data array');
+                assert.property(content.spec, 'encoding', 'Should have encoding property');
+                assert.property(content.spec.encoding, 'x', 'Should have x encoding');
+                assert.property(content.spec.encoding, 'y', 'Should have y encoding');
+
+                // Verify metadata structure
+                assert.property(newCell.metadata, '__deepnotePocket', 'Should have __deepnotePocket metadata');
+                assert.equal(newCell.metadata.__deepnotePocket.type, 'visualization', 'Should have visualization type');
+
+                // Verify reveal and selection were set
+                assert.isTrue((editor.revealRange as sinon.SinonStub).calledOnce, 'Should reveal the new cell range');
+                const revealCall = (editor.revealRange as sinon.SinonStub).firstCall;
+                assert.equal(revealCall.args[0].start, 0, 'Should reveal correct range start');
+                assert.equal(revealCall.args[0].end, 1, 'Should reveal correct range end');
+                assert.equal(revealCall.args[1], 0, 'Should use NotebookEditorRevealType.Default (value 0)');
+
+                // Verify notebook.cell.edit command was executed
+                assert.isTrue(
+                    executeCommandStub.calledWith('notebook.cell.edit'),
+                    'Should execute notebook.cell.edit command'
+                );
+            });
+
+            test('should add chart block after selection when selection exists', async () => {
+                // Setup mocks
+                const existingCells = [createMockCell('{}'), createMockCell('{}')];
+                const selection = new NotebookRange(0, 1);
+                const { editor } = createMockEditor(existingCells, selection);
+                const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor);
+
+                // Call the method
+                await commandListener.addChartBlock();
+
+                const capturedNotebookEdits = getCapturedNotebookEdits();
+
+                // Verify chainWithPendingUpdates was called
+                assert.isTrue(chainStub.calledOnce, 'chainWithPendingUpdates should be called once');
+
+                // Verify a cell was inserted
+                assert.isNotNull(capturedNotebookEdits, 'Notebook edits should be captured');
+                const notebookEdit = capturedNotebookEdits![0] as any;
+                assert.equal(notebookEdit.newCells.length, 1, 'Should insert one cell');
+                assert.equal(notebookEdit.newCells[0].languageId, 'json', 'Should be JSON cell');
+            });
+
+            test('should use hardcoded variable name df_1', async () => {
+                // Setup mocks with existing df variables
+                const existingCells = [
+                    createMockCell('{ "deepnote_variable_name": "df_1" }'),
+                    createMockCell('{ "variable": "df_2" }')
+                ];
+                const { editor } = createMockEditor(existingCells, undefined);
+                const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor);
+
+                // Call the method
+                await commandListener.addChartBlock();
+
+                const capturedNotebookEdits = getCapturedNotebookEdits();
+                const notebookEdit = capturedNotebookEdits![0] as any;
+                const newCell = notebookEdit.newCells[0];
+
+                // Verify variable name is always df_1
+                const content = JSON.parse(newCell.value);
+                assert.equal(content.variable, 'df_1', 'Should always use df_1');
+            });
+
+            test('should always use df_1 regardless of existing variables', async () => {
+                // Setup mocks with various existing variables
+                const existingCells = [
+                    createMockCell('{ "deepnote_variable_name": "input_10" }'),
+                    createMockCell('{ "deepnote_variable_name": "df_5" }'),
+                    createMockCell('{ "variable": "df_2" }')
+                ];
+                const { editor } = createMockEditor(existingCells, undefined);
+                const { getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor);
+
+                // Call the method
+                await commandListener.addChartBlock();
+
+                const capturedNotebookEdits = getCapturedNotebookEdits();
+                const notebookEdit = capturedNotebookEdits![0] as any;
+                const newCell = notebookEdit.newCells[0];
+
+                // Verify variable name is always df_1
+                const content = JSON.parse(newCell.value);
+                assert.equal(content.variable, 'df_1', 'Should always use df_1');
+            });
+
+            test('should insert at correct position in the middle of notebook', async () => {
+                // Setup mocks
+                const existingCells = [createMockCell('{}'), createMockCell('{}'), createMockCell('{}')];
+                const selection = new NotebookRange(1, 2);
+                const { editor } = createMockEditor(existingCells, selection);
+                const { chainStub, getCapturedNotebookEdits } = mockNotebookUpdateAndExecute(editor);
+
+                // Call the method
+                await commandListener.addChartBlock();
+
+                const capturedNotebookEdits = getCapturedNotebookEdits();
+
+                // Verify chainWithPendingUpdates was called
+                assert.isTrue(chainStub.calledOnce, 'chainWithPendingUpdates should be called once');
+
+                // Verify a cell was inserted
+                assert.isNotNull(capturedNotebookEdits, 'Notebook edits should be captured');
+                const notebookEdit = capturedNotebookEdits![0] as any;
+                assert.equal(notebookEdit.newCells.length, 1, 'Should insert one cell');
+                assert.equal(notebookEdit.newCells[0].languageId, 'json', 'Should be JSON cell');
+            });
+
+            test('should throw error when no active editor exists', async () => {
+                // Setup: no active editor
+                Object.defineProperty(window, 'activeNotebookEditor', {
+                    value: undefined,
+                    configurable: true,
+                    writable: true
+                });
+
+                // Call the method and expect rejection
+                await assert.isRejected(
+                    commandListener.addChartBlock(),
+                    WrappedError,
+                    'No active notebook editor found'
+                );
+            });
+
+            test('should throw error when chainWithPendingUpdates fails', async () => {
+                // Setup mocks
+                const { editor } = createMockEditor([], undefined);
+                Object.defineProperty(window, 'activeNotebookEditor', {
+                    value: editor,
+                    configurable: true,
+                    writable: true
+                });
+
+                // Mock chainWithPendingUpdates to return false
+                sandbox.stub(notebookUpdater, 'chainWithPendingUpdates').resolves(false);
+
+                // Call the method and expect rejection
+                await assert.isRejected(commandListener.addChartBlock(), WrappedError, 'Failed to insert chart block');
+            });
+        });
     });
 });

src/platform/common/constants.ts

@@ -227,6 +227,7 @@ export namespace Commands {
     export const ManageIntegrations = 'deepnote.manageIntegrations';
     export const AddSqlBlock = 'deepnote.addSqlBlock';
     export const AddBigNumberChartBlock = 'deepnote.addBigNumberChartBlock';
+    export const AddChartBlock = 'deepnote.addChartBlock';
     export const AddInputTextBlock = 'deepnote.addInputTextBlock';
     export const AddInputTextareaBlock = 'deepnote.addInputTextareaBlock';
     export const AddInputSelectBlock = 'deepnote.addInputSelectBlock';