Include error in status and persist exec times

src/client/datascience/notebook/executionHelpers.ts

@@ -140,6 +140,27 @@ export function updateCellOutput(notebookCellModel: ICell, outputs: nbformat.IOu
     model.update(updateCell);
 }
 
+/**
+ * Store execution start and end times in ISO format for portability.
+ */
+export function updateCellExecutionTimes(
+    notebookCellModel: ICell,
+    model: INotebookModel,
+    startTime?: number,
+    duration?: number
+) {
+    const startTimeISO = startTime ? new Date(startTime).toISOString() : undefined;
+    const endTimeISO = duration && startTime ? new Date(startTime + duration).toISOString() : undefined;
+    updateCellMetadata(
+        notebookCellModel,
+        {
+            end_execution_time: endTimeISO,
+            start_execution_time: startTimeISO
+        },
+        model
+    );
+}
+
 export function updateCellMetadata(
     notebookCellModel: ICell,
     metadata: Partial<IBaseCellVSCodeMetadata>,

src/client/datascience/notebook/executionService.ts

@@ -22,10 +22,12 @@ import { findMappedNotebookCellModel } from './cellUpdateHelpers';
 import {
     handleUpdateDisplayDataMessage,
     hasTransientOutputForAnotherCell,
+    updateCellExecutionTimes,
     updateCellExecutionCount,
     updateCellOutput,
     updateCellWithErrorStatus
 } from './executionHelpers';
+import { getCellStatusMessageBasedOnFirstErrorOutput } from './helpers';
 import { INotebookExecutionService } from './types';
 // tslint:disable-next-line: no-var-requires no-require-imports
 const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed');
@@ -185,9 +187,23 @@ export class NotebookExecutionService implements INotebookExecutionService {
                         ? vscodeNotebookEnums.NotebookCellRunState.Idle
                         : vscodeNotebookEnums.NotebookCellRunState.Success;
                     cell.metadata.statusMessage = '';
+
+                    // Update metadata in our model.
+                    const notebookCellModel = findMappedNotebookCellModel(cell, model.cells);
+                    updateCellExecutionTimes(
+                        notebookCellModel,
+                        model,
+                        cell.metadata.runStartTime,
+                        cell.metadata.lastRunDuration
+                    );
+
                     // 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;
+                        cell.metadata.statusMessage = getCellStatusMessageBasedOnFirstErrorOutput(
+                            // tslint:disable-next-line: no-any
+                            notebookCellModel.data.outputs as any
+                        );
                     }
                     deferred.resolve();
                 }

src/client/datascience/notebook/helpers.ts

@@ -16,6 +16,7 @@ import type {
 // tslint:disable-next-line: no-var-requires no-require-imports
 const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed');
 import * as uuid from 'uuid/v4';
+import { NotebookCellRunState } from '../../../../typings/vscode-proposed';
 import {
     concatMultilineStringInput,
     concatMultilineStringOutput,
@@ -90,24 +91,62 @@ export function cellToVSCNotebookCellData(cell: ICell): NotebookCellData | undef
         return;
     }
 
-    return {
+    // tslint:disable-next-line: no-any
+    const outputs = cellOutputsToVSCCellOutputs(cell.data.outputs as any);
+
+    // If we have an execution count & no errors, then success state.
+    // If we have an execution count &  errors, then error state.
+    // Else idle state.
+    const hasErrors = outputs.some((output) => output.outputKind === vscodeNotebookEnums.CellOutputKind.Error);
+    const hasExecutionCount = typeof cell.data.execution_count === 'number' && cell.data.execution_count > 0;
+    let runState: NotebookCellRunState;
+    let statusMessage: string | undefined;
+    if (!hasExecutionCount) {
+        runState = vscodeNotebookEnums.NotebookCellRunState.Idle;
+    } else if (hasErrors) {
+        runState = vscodeNotebookEnums.NotebookCellRunState.Error;
+        // Error details are stripped from the output, get raw output.
+        // tslint:disable-next-line: no-any
+        statusMessage = getCellStatusMessageBasedOnFirstErrorOutput(cell.data.outputs as any);
+    } else {
+        runState = vscodeNotebookEnums.NotebookCellRunState.Success;
+    }
+
+    const notebookCellData: NotebookCellData = {
         cellKind:
             cell.data.cell_type === 'code' ? vscodeNotebookEnums.CellKind.Code : vscodeNotebookEnums.CellKind.Markdown,
         language: cell.data.cell_type === 'code' ? PYTHON_LANGUAGE : MARKDOWN_LANGUAGE,
         metadata: {
             editable: true,
             executionOrder: typeof cell.data.execution_count === 'number' ? cell.data.execution_count : undefined,
             hasExecutionOrder: cell.data.cell_type === 'code',
-            runState: vscodeNotebookEnums.NotebookCellRunState.Idle,
+            runState,
             runnable: cell.data.cell_type === 'code',
             custom: {
                 cellId: cell.id
             }
         },
         source: concatMultilineStringInput(cell.data.source),
-        // tslint:disable-next-line: no-any
-        outputs: cellOutputsToVSCCellOutputs(cell.data.outputs as any)
+        outputs
     };
+
+    if (statusMessage) {
+        notebookCellData.metadata.statusMessage = statusMessage;
+    }
+
+    const startExecutionTime = cell.data.metadata.vscode?.start_execution_time
+        ? new Date(Date.parse(cell.data.metadata.vscode.start_execution_time)).getTime()
+        : undefined;
+    const endExecutionTime = cell.data.metadata.vscode?.end_execution_time
+        ? new Date(Date.parse(cell.data.metadata.vscode.end_execution_time)).getTime()
+        : undefined;
+
+    if (startExecutionTime && typeof endExecutionTime === 'number') {
+        notebookCellData.metadata.runStartTime = startExecutionTime;
+        notebookCellData.metadata.lastRunDuration = endExecutionTime - startExecutionTime;
+    }
+
+    return notebookCellData;
 }
 
 export function cellOutputsToVSCCellOutputs(outputs?: nbformat.IOutput[]): CellOutput[] {
@@ -205,11 +244,32 @@ function translateStreamOutput(output: nbformat.IStream): CellStreamOutput | Cel
         }
     };
 }
+
+/**
+ * We will display the error message in the status of the cell.
+ * We will display the error message in the status of the cell.
+ * The `ename` & `evalue` is displayed at the top of the output by VS Code.
+ * As we're displaying the error in the statusbar, we don't want this dup error in output.
+ * Hence remove this.
+ */
 export function translateErrorOutput(output: nbformat.IError): CellErrorOutput {
     return {
-        ename: output.ename,
-        evalue: output.evalue,
+        ename: '',
+        evalue: '',
         outputKind: vscodeNotebookEnums.CellOutputKind.Error,
         traceback: output.traceback
     };
 }
+
+export function getCellStatusMessageBasedOnFirstErrorOutput(outputs?: nbformat.IOutput[]): string {
+    if (!Array.isArray(outputs)) {
+        return '';
+    }
+    const errorOutput = (outputs.find((output) => output.output_type === 'error') as unknown) as
+        | nbformat.IError
+        | undefined;
+    if (!errorOutput) {
+        return '';
+    }
+    return `${errorOutput.ename}${errorOutput.evalue ? ': ' : ''}${errorOutput.evalue}`;
+}

src/test/common.ts

@@ -652,6 +652,9 @@ export class TestEventHandler<T extends void | any = any> implements IDisposable
     public get count(): number {
         return this.handledEvents.length;
     }
+    public get all(): T[] {
+        return this.handledEvents;
+    }
     private readonly handler: IDisposable;
     // tslint:disable-next-line: no-any
     private readonly handledEvents: any[] = [];

src/test/datascience/notebook/helper.ts

@@ -18,7 +18,7 @@ import { IDisposable } from '../../../client/common/types';
 import { noop, swallowExceptions } from '../../../client/common/utils/misc';
 import { NotebookContentProvider } from '../../../client/datascience/notebook/contentProvider';
 import { ICell, INotebookEditorProvider, INotebookProvider } from '../../../client/datascience/types';
-import { waitForCondition } from '../../common';
+import { createEventHandler, waitForCondition } from '../../common';
 import { EXTENSION_ROOT_DIR_FOR_TESTS } from '../../constants';
 import { closeActiveWindows, initialize } from '../../initialize';
 const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed');
@@ -252,3 +252,19 @@ export function assertVSCCellHasErrors(cell: NotebookCell) {
     assert.equal(cell.metadata.runState, vscodeNotebookEnums.NotebookCellRunState.Error);
     return true;
 }
+export function assertVSCCellHasErrorOutput(cell: NotebookCell) {
+    assert.ok(
+        cell.outputs.filter((output) => output.outputKind === vscodeNotebookEnums.CellOutputKind.Error).length,
+        'No error output in cell'
+    );
+    return true;
+}
+
+export async function saveActiveNotebook(disposables: IDisposable[]) {
+    const api = await initialize();
+    const editorProvider = api.serviceContainer.get<INotebookEditorProvider>(INotebookEditorProvider);
+    const savedEvent = createEventHandler(editorProvider.activeEditor!.model!, 'changed', disposables);
+    await commands.executeCommand('workbench.action.files.saveAll');
+
+    await waitForCondition(async () => savedEvent.all.some((e) => e.kind === 'save'), 5_000, 'Not saved');
+}

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

@@ -204,8 +204,8 @@ suite('Data Science - NativeNotebook helpers', () => {
                 [
                     {
                         outputKind: vscodeNotebookEnums.CellOutputKind.Error,
-                        ename: 'Error Name',
-                        evalue: 'Error Value',
+                        ename: '',
+                        evalue: '',
                         traceback: ['stack1', 'stack2', 'stack3']
                     }
                 ]

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

@@ -0,0 +1,167 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+'use strict';
+
+// tslint:disable:no-require-imports no-var-requires
+import { assert, expect } from 'chai';
+import * as path from 'path';
+import * as sinon from 'sinon';
+import { commands, Uri } from 'vscode';
+import { NotebookCell } from '../../../../typings/vscode-proposed';
+import { IVSCodeNotebook } from '../../../client/common/application/types';
+import { IDisposable } from '../../../client/common/types';
+import { sleep } from '../../../client/common/utils/async';
+import { INotebookEditorProvider } from '../../../client/datascience/types';
+import { IExtensionTestApi, waitForCondition } from '../../common';
+import { closeActiveWindows, EXTENSION_ROOT_DIR_FOR_TESTS, initialize } from '../../initialize';
+import {
+    assertHasExecutionCompletedSuccessfully,
+    assertHasExecutionCompletedWithErrors,
+    assertHasTextOutputInVSCode,
+    assertVSCCellHasErrorOutput,
+    assertVSCCellIsIdle,
+    assertVSCCellIsRunning,
+    canRunTests,
+    closeNotebooksAndCleanUpAfterTests,
+    createTemporaryNotebook,
+    deleteAllCellsAndWait,
+    insertPythonCellAndWait,
+    saveActiveNotebook,
+    startJupyter,
+    swallowSavingOfNotebooks
+} from './helper';
+const vscodeNotebookEnums = require('vscode') as typeof import('vscode-proposed');
+
+// tslint:disable: no-any no-invalid-this
+suite('DataScience - VSCode Notebook - (Saving)', function () {
+    this.timeout(15_000);
+    this.retries(0);
+    const templateIPynb = path.join(EXTENSION_ROOT_DIR_FOR_TESTS, 'src', 'test', 'datascience', 'test.ipynb');
+    let api: IExtensionTestApi;
+    let testIPynb: Uri;
+    let editorProvider: INotebookEditorProvider;
+    const disposables: IDisposable[] = [];
+    let vscodeNotebook: IVSCodeNotebook;
+    suiteSetup(async function () {
+        this.timeout(15_000);
+        this.retries(0);
+        api = await initialize();
+        if (!(await canRunTests())) {
+            return this.skip();
+        }
+        await startJupyter();
+        vscodeNotebook = api.serviceContainer.get<IVSCodeNotebook>(IVSCodeNotebook);
+        editorProvider = api.serviceContainer.get<INotebookEditorProvider>(INotebookEditorProvider);
+    });
+    setup(async () => {
+        sinon.restore();
+        // Don't use same file (due to dirty handling, we might save in dirty.)
+        // Cuz we won't save to file, hence extension will backup in dirty file and when u re-open it will open from dirty.
+        testIPynb = Uri.file(await createTemporaryNotebook(templateIPynb, disposables));
+        await editorProvider.open(testIPynb);
+        await deleteAllCellsAndWait();
+    });
+    teardown(async () => {
+        await swallowSavingOfNotebooks();
+        await closeNotebooksAndCleanUpAfterTests(disposables);
+    });
+
+    test('Verify output & metadata when re-opening (slow)', async () => {
+        await insertPythonCellAndWait('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
+        await insertPythonCellAndWait('import time\nfor i in range(10000):\n  print(i)\n  time.sleep(0.1)', 0);
+        await insertPythonCellAndWait('print(a)', 0);
+        await insertPythonCellAndWait('print(1)', 0);
+        let cell1: NotebookCell;
+        let cell2: NotebookCell;
+        let cell3: NotebookCell;
+        let cell4: NotebookCell;
+
+        function initializeCells() {
+            cell1 = vscodeNotebook.activeNotebookEditor?.document.cells![0]!;
+            cell2 = vscodeNotebook.activeNotebookEditor?.document.cells![1]!;
+            cell3 = vscodeNotebook.activeNotebookEditor?.document.cells![2]!;
+            cell4 = vscodeNotebook.activeNotebookEditor?.document.cells![3]!;
+        }
+        initializeCells();
+        await commands.executeCommand('notebook.execute');
+
+        // Wait till 1 & 2 finish & 3rd cell starts executing.
+        await waitForCondition(
+            async () =>
+                assertHasExecutionCompletedSuccessfully(cell1) &&
+                assertHasExecutionCompletedWithErrors(cell2) &&
+                assertVSCCellIsRunning(cell3) &&
+                assertVSCCellIsRunning(cell4),
+            15_000,
+            'Cells not running'
+        );
+
+        await sleep(1); // Wait for some output.
+        await commands.executeCommand('notebook.cancelExecution');
+
+        // Wait till execution count changes and status is error.
+        await waitForCondition(
+            async () => assertHasExecutionCompletedWithErrors(cell3) && assertVSCCellIsIdle(cell4),
+            15_000,
+            'Cells not running'
+        );
+
+        function verifyCelMetadata() {
+            assert.lengthOf(cell1.outputs, 1, 'Incorrect output for cell 1');
+            assert.lengthOf(cell2.outputs, 1, 'Incorrect output for cell 2');
+            assert.lengthOf(cell3.outputs, 2, 'Incorrect output for cell 3'); // stream and interrupt error.
+            assert.lengthOf(cell4.outputs, 0, 'Incorrect output for cell 4');
+
+            assert.equal(
+                cell1.metadata.runState,
+                vscodeNotebookEnums.NotebookCellRunState.Success,
+                'Incorrect state 1'
+            );
+            assert.equal(cell2.metadata.runState, vscodeNotebookEnums.NotebookCellRunState.Error, 'Incorrect state 2');
+            assert.equal(cell3.metadata.runState, vscodeNotebookEnums.NotebookCellRunState.Error, 'Incorrect state 3');
+            assert.equal(cell4.metadata.runState, vscodeNotebookEnums.NotebookCellRunState.Idle, 'Incorrect state 4');
+
+            assertHasTextOutputInVSCode(cell1, '1', 0);
+            assertVSCCellHasErrorOutput(cell2);
+            assertHasTextOutputInVSCode(cell3, '0', 0);
+            assertVSCCellHasErrorOutput(cell3);
+
+            expect(cell1.metadata.executionOrder).to.be.greaterThan(0, 'Execution count should be > 0');
+            expect(cell2.metadata.executionOrder).to.be.greaterThan(
+                cell1.metadata.executionOrder!,
+                'Execution count > cell 1'
+            );
+            expect(cell3.metadata.executionOrder).to.be.greaterThan(
+                cell2.metadata.executionOrder!,
+                'Execution count > cell 2'
+            );
+            assert.isNotOk(cell4.metadata.executionOrder, 'Execution count should be 0|null');
+
+            assert.isEmpty(cell1.metadata.statusMessage, 'Cell 1 status should be empty'); // No errors.
+            assert.isNotEmpty(cell2.metadata.statusMessage, 'Cell 1 status should be empty'); // Errors.
+            assert.isNotEmpty(cell3.metadata.statusMessage, 'Cell 1 status should be empty'); // Errors (interrupted).
+            assert.isEmpty(cell4.metadata.statusMessage, 'Cell 1 status should be empty'); // No errors (didn't run).
+
+            assert.isOk(cell1.metadata.runStartTime, 'Start time should be > 0');
+            assert.isOk(cell1.metadata.lastRunDuration, 'Duration should be > 0');
+            assert.isOk(cell2.metadata.runStartTime, 'Start time should be > 0');
+            assert.isOk(cell2.metadata.lastRunDuration, 'Duration should be > 0');
+            assert.isOk(cell3.metadata.runStartTime, 'Start time should be > 0');
+            assert.isOk(cell3.metadata.lastRunDuration, 'Duration should be > 0');
+            assert.isOk(cell4.metadata.runStartTime, 'Start time should be > 0');
+            assert.isOk(cell4.metadata.lastRunDuration, 'Duration should be > 0');
+        }
+
+        verifyCelMetadata();
+
+        // Save and close this nb.
+        await saveActiveNotebook(disposables);
+        await closeActiveWindows();
+
+        // Reopen the notebook & validate the metadata.
+        await editorProvider.open(testIPynb);
+        initializeCells();
+        verifyCelMetadata();
+    });
+});