Updates to cell status displayed in VS Code Notebooks

src/client/datascience/interactive-ipynb/nativeEditorStorage.ts

@@ -331,7 +331,8 @@ export class NativeEditorNotebookModel implements INotebookModel {
     private updateCellExecutionCount(cellId: string, executionCount?: number) {
         const index = this.cells.findIndex((v) => v.id === cellId);
         if (index >= 0) {
-            this._state.cells[index].data.execution_count = (executionCount || 0) > 0 ? executionCount : null;
+            this._state.cells[index].data.execution_count =
+                typeof executionCount === 'number' && executionCount > 0 ? executionCount : null;
             return true;
         }
         return false;

src/client/datascience/notebook/executionHelpers.ts

@@ -6,6 +6,7 @@
 import type { nbformat } from '@jupyterlab/coreutils';
 import type { KernelMessage } from '@jupyterlab/services';
 import { NotebookCell, NotebookCellRunState, NotebookDocument } from 'vscode';
+import { IBaseCellVSCodeMetadata } from '../../../../types/@jupyterlab_coreutils_nbformat';
 import { createErrorOutput } from '../../../datascience-ui/common/cellFactory';
 import {
     INotebookModelCellExecutionCountChange,
@@ -138,3 +139,34 @@ export function updateCellOutput(notebookCellModel: ICell, outputs: nbformat.IOu
     };
     model.update(updateCell);
 }
+
+export function updateCellMetadata(
+    notebookCellModel: ICell,
+    metadata: Partial<IBaseCellVSCodeMetadata>,
+    model: INotebookModel
+) {
+    const originalVscodeMetadata: IBaseCellVSCodeMetadata = notebookCellModel.data.metadata.vscode || {};
+    // Update our model with the new metadata stored in jupyter.
+    const newCell: ICell = {
+        ...notebookCellModel,
+        data: {
+            ...notebookCellModel.data,
+            metadata: {
+                ...notebookCellModel.data.metadata,
+                vscode: {
+                    ...originalVscodeMetadata,
+                    ...metadata
+                }
+            }
+        }
+    };
+    const updateCell: INotebookModelModifyChange = {
+        kind: 'modify',
+        newCells: [newCell],
+        oldCells: [notebookCellModel],
+        newDirty: true,
+        oldDirty: model.isDirty === true,
+        source: 'user'
+    };
+    model.update(updateCell);
+}

src/client/datascience/notebook/executionService.ts

@@ -155,9 +155,6 @@ export class NotebookExecutionService implements INotebookExecutionService {
             const observable = nb.executeObservable(cell.source, document.fileName, 0, cell.uri.fsPath, false);
             subscription = observable?.subscribe(
                 (cells) => {
-                    if (wrappedToken.isCancellationRequested) {
-                        return;
-                    }
                     const rawCellOutput = cells
                         .filter((item) => item.id === cell.uri.fsPath)
                         .flatMap((item) => (item.data.outputs as unknown) as nbformat.IOutput[])
@@ -183,11 +180,15 @@ export class NotebookExecutionService implements INotebookExecutionService {
                     this.errorHandler.handleError((error as unknown) as Error).ignoreErrors();
                 },
                 () => {
+                    cell.metadata.lastRunDuration = stopWatch.elapsedTime;
                     cell.metadata.runState = wrappedToken.isCancellationRequested
                         ? vscodeNotebookEnums.NotebookCellRunState.Idle
                         : vscodeNotebookEnums.NotebookCellRunState.Success;
-                    cell.metadata.lastRunDuration = stopWatch.elapsedTime;
                     cell.metadata.statusMessage = '';
+                    // If there are any errors in the cell, then change status to error.
+                    if (cell.outputs.some((output) => output.outputKind === vscodeNotebookEnums.CellOutputKind.Error)) {
+                        cell.metadata.runState = vscodeNotebookEnums.NotebookCellRunState.Error;
+                    }
                     deferred.resolve();
                 }
             );

src/client/datascience/notebook/helpers.ts

@@ -37,7 +37,7 @@ export function notebookModelToVSCNotebookData(model: INotebookModel): NotebookD
 
     return {
         cells,
-        languages: [PYTHON_LANGUAGE, MARKDOWN_LANGUAGE],
+        languages: [PYTHON_LANGUAGE],
         metadata: {
             cellEditable: true,
             cellRunnable: true,

src/test/datascience/notebook/contentProvider.unit.test.ts

@@ -61,7 +61,7 @@ suite('Data Science - NativeNotebook ContentProvider', () => {
         const notebook = await contentProvider.openNotebook(fileUri);
 
         assert.isOk(notebook);
-        assert.deepEqual(notebook.languages, [PYTHON_LANGUAGE, MARKDOWN_LANGUAGE]);
+        assert.deepEqual(notebook.languages, [PYTHON_LANGUAGE]);
         assert.deepEqual(notebook.cells, [
             {
                 cellKind: (vscodeNotebookEnums as any).CellKind.Code,

src/test/datascience/notebook/executionWithJupyter.ds.test.ts → src/test/datascience/notebook/executionService.ds.test.ts

@@ -4,16 +4,18 @@
 'use strict';
 
 // tslint:disable:no-require-imports no-var-requires
-import { assert } from 'chai';
+import { assert, expect } from 'chai';
 import * as sinon from 'sinon';
 import { commands } from 'vscode';
+import { CellErrorOutput } from '../../../../typings/vscode-proposed';
 import { IVSCodeNotebook } from '../../../client/common/application/types';
 import { IDisposable } from '../../../client/common/types';
 import { INotebookEditorProvider } from '../../../client/datascience/types';
 import { IExtensionTestApi, waitForCondition } from '../../common';
 import { initialize } from '../../initialize';
 import {
     assertHasExecutionCompletedSuccessfully,
+    assertHasExecutionCompletedWithErrors,
     assertHasTextOutputInICell,
     assertHasTextOutputInVSCode,
     canRunTests,
@@ -23,6 +25,8 @@ import {
     startJupyter,
     swallowSavingOfNotebooks
 } from './helper';
+// tslint:disable-next-line: no-var-requires no-require-imports
+const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed');
 
 // tslint:disable: no-any no-invalid-this
 suite('DataScience - VSCode Notebook - (Execution)', function () {
@@ -113,4 +117,48 @@ suite('DataScience - VSCode Notebook - (Execution)', function () {
         assert.ok(vscCell[0].metadata.executionOrder, 'Execution count should be > 0');
         assert.equal(vscCell[1].metadata.executionOrder! - 1, vscCell[0].metadata.executionOrder!);
     });
+    test('Verify metadata for successfully executed cell', async () => {
+        await insertPythonCellAndWait('print("Foo Bar")', 0);
+        const cell = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
+
+        await commands.executeCommand('notebook.execute');
+
+        // Wait till execution count changes and status is success.
+        await waitForCondition(
+            async () => assertHasExecutionCompletedSuccessfully(cell),
+            15_000,
+            'Cell did not get executed'
+        );
+
+        expect(cell.metadata.executionOrder).to.be.greaterThan(0, 'Execution count should be > 0');
+        expect(cell.metadata.runStartTime).to.be.greaterThan(0, 'Start time should be > 0');
+        expect(cell.metadata.lastRunDuration).to.be.greaterThan(0, 'Duration should be > 0');
+        assert.equal(cell.metadata.runState, vscodeNotebookEnums.NotebookCellRunState.Success, 'Incorrect State');
+        assert.equal(cell.metadata.statusMessage, '', 'Incorrect Status message');
+    });
+    test('Verify output & metadata for executed cell with errors', async () => {
+        await insertPythonCellAndWait('print(abcd)', 0);
+        const cell = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
+
+        await commands.executeCommand('notebook.execute');
+
+        // Wait till execution count changes and status is error.
+        await waitForCondition(
+            async () => assertHasExecutionCompletedWithErrors(cell),
+            15_000,
+            'Cell did not get executed'
+        );
+
+        assert.lengthOf(cell.outputs, 1, 'Incorrect output');
+        const errorOutput = cell.outputs[0] as CellErrorOutput;
+        assert.equal(errorOutput.outputKind, vscodeNotebookEnums.CellOutputKind.Error, 'Incorrect output');
+        assert.isNotEmpty(errorOutput.ename, 'Incorrect ename');
+        assert.isNotEmpty(errorOutput.evalue, 'Incorrect evalue');
+        assert.isNotEmpty(errorOutput.traceback, 'Incorrect traceback');
+        expect(cell.metadata.executionOrder).to.be.greaterThan(0, 'Execution count should be > 0');
+        expect(cell.metadata.runStartTime).to.be.greaterThan(0, 'Start time should be > 0');
+        expect(cell.metadata.lastRunDuration).to.be.greaterThan(0, 'Duration should be > 0');
+        assert.equal(cell.metadata.runState, vscodeNotebookEnums.NotebookCellRunState.Error, 'Incorrect State');
+        assert.equal(cell.metadata.statusMessage, '', 'Incorrect Status message');
+    });
 });

src/test/datascience/notebook/helper.ts

@@ -11,6 +11,7 @@ import * as sinon from 'sinon';
 import * as tmp from 'tmp';
 import { commands } from 'vscode';
 import { NotebookCell } from '../../../../types/vscode-proposed';
+import { CellDisplayOutput } from '../../../../typings/vscode-proposed';
 import { IApplicationEnvironment, IVSCodeNotebook } from '../../../client/common/application/types';
 import { MARKDOWN_LANGUAGE, PYTHON_LANGUAGE } from '../../../client/common/constants';
 import { IDisposable } from '../../../client/common/types';
@@ -211,6 +212,12 @@ export function assertHasExecutionCompletedSuccessfully(cell: NotebookCell) {
         cell.metadata.runState === vscodeNotebookEnums.NotebookCellRunState.Success
     );
 }
+export function assertHasExecutionCompletedWithErrors(cell: NotebookCell) {
+    return (
+        (cell.metadata.executionOrder ?? 0) > 0 &&
+        cell.metadata.runState === vscodeNotebookEnums.NotebookCellRunState.Error
+    );
+}
 export function hasOutputInVSCode(cell: NotebookCell) {
     assert.ok(cell.outputs.length, 'No output');
 }
@@ -220,8 +227,8 @@ export function hasOutputInICell(cell: ICell) {
 export function assertHasTextOutputInVSCode(cell: NotebookCell, text: string, index: number, isExactMatch = true) {
     const cellOutputs = cell.outputs;
     assert.ok(cellOutputs, 'No output');
-    assert.equal(cellOutputs[index].outputKind, vscodeNotebookEnums.CellOutputKind.Text, 'Incorrect output kind');
-    const outputText = (cellOutputs[index] as any).text.trim();
+    assert.equal(cellOutputs[index].outputKind, vscodeNotebookEnums.CellOutputKind.Rich, 'Incorrect output kind');
+    const outputText = (cellOutputs[index] as CellDisplayOutput).data['text/plain'].trim();
     if (isExactMatch) {
         assert.equal(outputText, text, 'Incorrect output');
     } else {
@@ -241,3 +248,7 @@ export function assertVSCCellIsIdle(cell: NotebookCell) {
     assert.equal(cell.metadata.runState, vscodeNotebookEnums.NotebookCellRunState.Idle);
     return true;
 }
+export function assertVSCCellHasErrors(cell: NotebookCell) {
+    assert.equal(cell.metadata.runState, vscodeNotebookEnums.NotebookCellRunState.Error);
+    return true;
+}

src/test/datascience/notebook/helpers.unit.test.ts

@@ -46,7 +46,7 @@ suite('Data Science - NativeNotebook helpers', () => {
         const notebook = notebookModelToVSCNotebookData((model as unknown) as INotebookModel);
 
         assert.isOk(notebook);
-        assert.deepEqual(notebook.languages, [PYTHON_LANGUAGE, MARKDOWN_LANGUAGE]);
+        assert.deepEqual(notebook.languages, [PYTHON_LANGUAGE]);
         assert.deepEqual(notebook.cells, [
             {
                 cellKind: vscodeNotebookEnums.CellKind.Code,

src/test/datascience/notebook/interrupRestart.ds.test.ts

@@ -6,6 +6,7 @@
 import { assert } from 'chai';
 import * as sinon from 'sinon';
 import { CancellationTokenSource, commands, NotebookEditor as VSCNotebookEditor } from 'vscode';
+import { CellErrorOutput, CellStreamOutput } from '../../../../typings/vscode-proposed';
 import { IVSCodeNotebook } from '../../../client/common/application/types';
 import { IDisposable } from '../../../client/common/types';
 import { createDeferredFromPromise, sleep } from '../../../client/common/utils/async';
@@ -14,6 +15,7 @@ import { INotebookEditorProvider } from '../../../client/datascience/types';
 import { IExtensionTestApi, waitForCondition } from '../../common';
 import { initialize } from '../../initialize';
 import {
+    assertVSCCellHasErrors,
     assertVSCCellIsIdle,
     assertVSCCellIsRunning,
     canRunTests,
@@ -24,6 +26,8 @@ import {
     startJupyter,
     swallowSavingOfNotebooks
 } from './helper';
+// tslint:disable-next-line: no-var-requires no-require-imports
+const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed');
 
 // tslint:disable: no-any no-invalid-this
 /*
@@ -47,6 +51,7 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors', functio
         if (!(await canRunTests())) {
             return this.skip();
         }
+        await closeNotebooksAndCleanUpAfterTests();
         await startJupyter();
         vscodeNotebook = api.serviceContainer.get<IVSCodeNotebook>(IVSCodeNotebook);
         editorProvider = api.serviceContainer.get<INotebookEditorProvider>(INotebookEditorProvider);
@@ -85,7 +90,7 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors', functio
         cancellation.cancel();
 
         await waitForCondition(async () => deferred.completed, 5_000, 'Execution not cancelled');
-        assertVSCCellIsIdle(cell);
+        assertVSCCellHasErrors(cell);
     });
     test('Cancelling using VSC Command for cell (slow)', async function () {
         // Fails due to VSC bugs.
@@ -105,7 +110,7 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors', functio
         // Interrupt the kernel.
         await commands.executeCommand('notebook.cell.cancelExecution', cell);
 
-        await waitForCondition(async () => assertVSCCellIsIdle(cell), 1_000, 'Execution not cancelled');
+        await waitForCondition(async () => assertVSCCellHasErrors(cell), 1_000, 'Execution not cancelled');
     });
     test('Cancelling using VSC Command in toolbar (slow)', async () => {
         await insertPythonCellAndWait('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
@@ -123,7 +128,7 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors', functio
         // Interrupt the kernel.
         await commands.executeCommand('notebook.cancelExecution');
 
-        await waitForCondition(async () => assertVSCCellIsIdle(cell), 1_000, 'Execution not cancelled');
+        await waitForCondition(async () => assertVSCCellHasErrors(cell), 1_000, 'Execution not cancelled');
     });
     test('When running entire notebook, clicking VSCode Stop button should trigger a single interrupt, not one per cell', async () => {
         await insertPythonCellAndWait('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
@@ -151,11 +156,20 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors', functio
 
         // Wait for ?s, and verify cells are not running.
         await waitForCondition(
-            async () => assertVSCCellIsIdle(cell1) && assertVSCCellIsIdle(cell2) && assertVSCCellIsIdle(cell3),
+            // First cell will have interrupt error and others will have just stopped even before they start.
+            async () => assertVSCCellHasErrors(cell1) && assertVSCCellIsIdle(cell2) && assertVSCCellIsIdle(cell3),
             15_000,
             'Cells are still running'
         );
 
+        assert.lengthOf(cell1.outputs, 2, 'First cell should have two outputs (stream and error)');
+        assert.lengthOf(cell2.outputs, 0, 'Should not have any output');
+        assert.lengthOf(cell3.outputs, 0, 'Should not have any output');
+
+        const cell1Output1 = cell1.outputs[0] as CellStreamOutput;
+        const cell1Output2 = cell1.outputs[1] as CellErrorOutput;
+        assert.equal(cell1Output1.outputKind, vscodeNotebookEnums.CellOutputKind.Rich, 'Should be text output');
+        assert.equal(cell1Output2.outputKind, vscodeNotebookEnums.CellOutputKind.Error, 'Should be error output');
         assert.equal(interrupt.callCount, 1, 'Interrupt should have been invoked only once');
     });
     test('Restarting kernel will cancel cell execution (slow)', async () => {
@@ -174,7 +188,7 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors', functio
         // Restart the kernel.
         const restartPromise = commands.executeCommand('python.datascience.notebookeditor.restartkernel');
 
-        await waitForCondition(async () => assertVSCCellIsIdle(cell), 1_000, 'Execution not cancelled');
+        await waitForCondition(async () => assertVSCCellIsIdle(cell), 15_000, 'Execution not cancelled');
 
         // Wait before we execute cells again.
         await restartPromise;
@@ -187,6 +201,6 @@ suite('DataScience - VSCode Notebook - Restart/Interrupt/Cancel/Errors', functio
         await commands.executeCommand('notebook.cancelExecution');
 
         // Wait for ?s, and verify cells are not running.
-        await waitForCondition(async () => assertVSCCellIsIdle(cell), 15_000, 'Cell is still running');
+        await waitForCondition(async () => assertVSCCellHasErrors(cell), 15_000, 'Cell is still running');
     });
 });