fix: Ensure that Code blocks are Code blocks

src/notebooks/deepnote/deepnoteNewCellLanguageService.ts

@@ -0,0 +1,52 @@
+import { inject, injectable } from 'inversify';
+import { languages, NotebookCellKind, NotebookDocumentChangeEvent, workspace } from 'vscode';
+
+import { DEEPNOTE_NOTEBOOK_TYPE } from '../../kernels/deepnote/types';
+import { IExtensionSyncActivationService } from '../../platform/activation/types';
+import { PYTHON_LANGUAGE } from '../../platform/common/constants';
+import { IDisposableRegistry } from '../../platform/common/types';
+import { noop } from '../../platform/common/utils/misc';
+
+/**
+ * Ensures newly added code cells in Deepnote notebooks default to Python language.
+ * VS Code copies the language from adjacent cells when inserting, which causes
+ * new cells after SQL blocks to be SQL. This service corrects that by resetting
+ * unintentional language inheritance to Python.
+ */
+@injectable()
+export class DeepnoteNewCellLanguageService implements IExtensionSyncActivationService {
+    constructor(@inject(IDisposableRegistry) private readonly disposables: IDisposableRegistry) {}
+
+    public activate(): void {
+        this.disposables.push(workspace.onDidChangeNotebookDocument(this.onDidChangeNotebookDocument, this));
+    }
+
+    private async onDidChangeNotebookDocument(e: NotebookDocumentChangeEvent): Promise<void> {
+        if (e.notebook.notebookType !== DEEPNOTE_NOTEBOOK_TYPE) {
+            return;
+        }
+
+        for (const change of e.contentChanges) {
+            for (const cell of change.addedCells) {
+                if (cell.kind !== NotebookCellKind.Code) {
+                    continue;
+                }
+
+                if (cell.document.getText().trim().length > 0) {
+                    continue;
+                }
+
+                const pocketType = cell.metadata?.__deepnotePocket?.type;
+
+                if (pocketType) {
+                    continue;
+                }
+
+                // TODO: This will have to be revisited if we add support for other languages in Deepnote.
+                if (cell.document.languageId !== PYTHON_LANGUAGE) {
+                    languages.setTextDocumentLanguage(cell.document, PYTHON_LANGUAGE).then(noop, noop);
+                }
+            }
+        }
+    }
+}

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

@@ -0,0 +1,275 @@
+import { expect } from 'chai';
+import { anything, verify, when } from 'ts-mockito';
+import { Disposable, NotebookCell, NotebookCellKind, NotebookDocument, TextDocument, Uri } from 'vscode';
+
+import { IDisposableRegistry } from '../../platform/common/types';
+import { mockedVSCodeNamespaces, resetVSCodeMocks } from '../../test/vscode-mock';
+import { DeepnoteNewCellLanguageService } from './deepnoteNewCellLanguageService';
+
+suite('DeepnoteNewCellLanguageService', () => {
+    let service: DeepnoteNewCellLanguageService;
+    let disposables: Disposable[];
+    let notebookChangeHandler: ((e: any) => void) | undefined;
+
+    function createMockNotebook(notebookType: string): NotebookDocument {
+        return {
+            uri: Uri.file('/test/notebook.deepnote'),
+            notebookType
+        } as NotebookDocument;
+    }
+
+    function createMockCell(options: {
+        kind?: NotebookCellKind;
+        languageId?: string;
+        content?: string;
+        metadata?: Record<string, unknown>;
+    }): NotebookCell {
+        const { kind = NotebookCellKind.Code, languageId = 'python', content = '', metadata = {} } = options;
+
+        return {
+            index: 0,
+            notebook: createMockNotebook('deepnote'),
+            kind,
+            document: {
+                uri: Uri.file('/test/notebook.deepnote#cell0'),
+                languageId,
+                getText: () => content
+            } as TextDocument,
+            metadata,
+            outputs: [],
+            executionSummary: undefined
+        } as unknown as NotebookCell;
+    }
+
+    setup(() => {
+        resetVSCodeMocks();
+        disposables = [];
+
+        when(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything(), anything())).thenCall(
+            (handler, thisArg) => {
+                notebookChangeHandler = (e: any) => handler.call(thisArg, e);
+
+                return { dispose: () => undefined };
+            }
+        );
+
+        when(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).thenReturn(
+            Promise.resolve({} as TextDocument)
+        );
+
+        service = new DeepnoteNewCellLanguageService(disposables as unknown as IDisposableRegistry);
+    });
+
+    teardown(() => {
+        notebookChangeHandler = undefined;
+        disposables.forEach((d) => d.dispose());
+    });
+
+    test('activate registers onDidChangeNotebookDocument listener', () => {
+        service.activate();
+
+        verify(mockedVSCodeNamespaces.workspace.onDidChangeNotebookDocument(anything(), anything())).once();
+    });
+
+    test('activate adds disposable to registry', () => {
+        service.activate();
+
+        expect(disposables.length).to.be.greaterThan(0);
+    });
+
+    test('ignores non-deepnote notebooks', async () => {
+        service.activate();
+        const jupyterNotebook = createMockNotebook('jupyter-notebook');
+        const cell = createMockCell({ languageId: 'sql' });
+
+        notebookChangeHandler!({
+            notebook: jupyterNotebook,
+            contentChanges: [{ addedCells: [cell] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
+    });
+
+    test('ignores markdown cells', async () => {
+        service.activate();
+        const notebook = createMockNotebook('deepnote');
+        const cell = createMockCell({ kind: NotebookCellKind.Markup, languageId: 'markdown' });
+
+        notebookChangeHandler!({
+            notebook,
+            contentChanges: [{ addedCells: [cell] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
+    });
+
+    test('ignores cells with content', async () => {
+        service.activate();
+        const notebook = createMockNotebook('deepnote');
+        const cell = createMockCell({ languageId: 'sql', content: 'SELECT * FROM table' });
+
+        notebookChangeHandler!({
+            notebook,
+            contentChanges: [{ addedCells: [cell] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
+    });
+
+    test('ignores cells that already have Python language', async () => {
+        service.activate();
+        const notebook = createMockNotebook('deepnote');
+        const cell = createMockCell({ languageId: 'python' });
+
+        notebookChangeHandler!({
+            notebook,
+            contentChanges: [{ addedCells: [cell] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
+    });
+
+    test('ignores intentional SQL blocks (with __deepnotePocket.type)', async () => {
+        service.activate();
+        const notebook = createMockNotebook('deepnote');
+        const cell = createMockCell({
+            languageId: 'sql',
+            metadata: { __deepnotePocket: { type: 'sql' } }
+        });
+
+        notebookChangeHandler!({
+            notebook,
+            contentChanges: [{ addedCells: [cell] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
+    });
+
+    test('ignores intentional chart blocks (with __deepnotePocket.type)', async () => {
+        service.activate();
+        const notebook = createMockNotebook('deepnote');
+        const cell = createMockCell({
+            languageId: 'json',
+            metadata: { __deepnotePocket: { type: 'chart-vega' } }
+        });
+
+        notebookChangeHandler!({
+            notebook,
+            contentChanges: [{ addedCells: [cell] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
+    });
+
+    test('ignores intentional input blocks (with __deepnotePocket.type)', async () => {
+        service.activate();
+        const notebook = createMockNotebook('deepnote');
+        const cell = createMockCell({
+            languageId: 'plaintext',
+            metadata: { __deepnotePocket: { type: 'input-text' } }
+        });
+
+        notebookChangeHandler!({
+            notebook,
+            contentChanges: [{ addedCells: [cell] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
+    });
+
+    test('changes SQL cell to Python when no __deepnotePocket metadata', async () => {
+        service.activate();
+        const notebook = createMockNotebook('deepnote');
+        const cell = createMockCell({ languageId: 'sql' });
+
+        notebookChangeHandler!({
+            notebook,
+            contentChanges: [{ addedCells: [cell] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell.document, 'python')).once();
+    });
+
+    test('changes JSON cell to Python when no __deepnotePocket metadata', async () => {
+        service.activate();
+        const notebook = createMockNotebook('deepnote');
+        const cell = createMockCell({ languageId: 'json' });
+
+        notebookChangeHandler!({
+            notebook,
+            contentChanges: [{ addedCells: [cell] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell.document, 'python')).once();
+    });
+
+    test('handles multiple added cells', async () => {
+        service.activate();
+        const notebook = createMockNotebook('deepnote');
+        const sqlCell = createMockCell({ languageId: 'sql' });
+        const pythonCell = createMockCell({ languageId: 'python' });
+        const jsonCell = createMockCell({ languageId: 'json' });
+
+        notebookChangeHandler!({
+            notebook,
+            contentChanges: [{ addedCells: [sqlCell, pythonCell, jsonCell] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(sqlCell.document, 'python')).once();
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(pythonCell.document, anything())).never();
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(jsonCell.document, 'python')).once();
+    });
+
+    test('handles multiple content changes', async () => {
+        service.activate();
+        const notebook = createMockNotebook('deepnote');
+        const cell1 = createMockCell({ languageId: 'sql' });
+        const cell2 = createMockCell({ languageId: 'javascript' });
+
+        notebookChangeHandler!({
+            notebook,
+            contentChanges: [{ addedCells: [cell1] }, { addedCells: [cell2] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell1.document, 'python')).once();
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell2.document, 'python')).once();
+    });
+
+    test('ignores content changes with no added cells', async () => {
+        service.activate();
+        const notebook = createMockNotebook('deepnote');
+
+        notebookChangeHandler!({
+            notebook,
+            contentChanges: [{ addedCells: [] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(anything(), anything())).never();
+    });
+
+    test('changes cells with whitespace-only content to Python', async () => {
+        service.activate();
+        const notebook = createMockNotebook('deepnote');
+        const cell = createMockCell({ languageId: 'sql', content: '   \n\t  ' });
+
+        notebookChangeHandler!({
+            notebook,
+            contentChanges: [{ addedCells: [cell] }]
+        });
+        await new Promise((resolve) => setTimeout(resolve, 10));
+
+        verify(mockedVSCodeNamespaces.languages.setTextDocumentLanguage(cell.document, 'python')).once();
+    });
+});

src/notebooks/serviceRegistry.node.ts

@@ -82,6 +82,7 @@ import { DeepnoteNotebookEnvironmentMapper } from '../kernels/deepnote/environme
 import { DeepnoteNotebookCommandListener } from './deepnote/deepnoteNotebookCommandListener';
 import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider';
 import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider';
+import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService';
 import { SqlIntegrationStartupCodeProvider } from './deepnote/integrations/sqlIntegrationStartupCodeProvider';
 import { DeepnoteCellCopyHandler } from './deepnote/deepnoteCellCopyHandler';
 import { DeepnoteEnvironmentTreeDataProvider } from '../kernels/deepnote/environments/deepnoteEnvironmentTreeDataProvider.node';
@@ -213,6 +214,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea
         IExtensionSyncActivationService,
         DeepnoteBigNumberCellStatusBarProvider
     );
+    serviceManager.addSingleton<IExtensionSyncActivationService>(
+        IExtensionSyncActivationService,
+        DeepnoteNewCellLanguageService
+    );
 
     // Deepnote configuration services
     serviceManager.addSingleton<DeepnoteEnvironmentStorage>(DeepnoteEnvironmentStorage, DeepnoteEnvironmentStorage);

src/notebooks/serviceRegistry.web.ts

@@ -51,6 +51,7 @@ import {
 } from './deepnote/integrations/types';
 import { DeepnoteInputBlockCellStatusBarItemProvider } from './deepnote/deepnoteInputBlockCellStatusBarProvider';
 import { DeepnoteBigNumberCellStatusBarProvider } from './deepnote/deepnoteBigNumberCellStatusBarProvider';
+import { DeepnoteNewCellLanguageService } from './deepnote/deepnoteNewCellLanguageService';
 import { SqlCellStatusBarProvider } from './deepnote/sqlCellStatusBarProvider';
 import { IntegrationKernelRestartHandler } from './deepnote/integrations/integrationKernelRestartHandler';
 
@@ -123,6 +124,10 @@ export function registerTypes(serviceManager: IServiceManager, isDevMode: boolea
         IExtensionSyncActivationService,
         DeepnoteBigNumberCellStatusBarProvider
     );
+    serviceManager.addSingleton<IExtensionSyncActivationService>(
+        IExtensionSyncActivationService,
+        DeepnoteNewCellLanguageService
+    );
     serviceManager.addSingleton<IExtensionSyncActivationService>(
         IExtensionSyncActivationService,
         SqlCellStatusBarProvider

src/test/mocks/vsc/extHostedTypes.ts

@@ -103,6 +103,7 @@ export namespace vscMockExtHostedTypes {
     }
     export enum NotebookCellKind {
         Markdown = 1,
+        Markup = 1, // VS Code uses 'Markup' but some older code uses 'Markdown' - both have value 1
         Code = 2
     }
     export enum NotebookCellExecutionState {