Re-work cell execution end logic

Author: znorman-harris

libs/vscode/notebooks/src/lib/controller/idl-notebook-controller.class.ts

@@ -42,7 +42,10 @@ import { existsSync, mkdirSync, rmSync, writeFileSync } from 'fs';
 import { dirname, join } from 'path';
 import * as vscode from 'vscode';
 
-import { ICurrentCell } from './idl-notebook-controller.interface';
+import {
+  ICurrentCell,
+  IEndCellExecutionActions,
+} from './idl-notebook-controller.interface';
 import { ProcessIDLNotebookEmbeddedItems } from './process-idl-notebook-embedded-items';
 
 /**
@@ -157,7 +160,7 @@ export class IDLNotebookController {
         content: [`Stopped because: "${reason}"`, stack],
       });
 
-      await this._endCellExecution(false);
+      await this._endCellExecution(false, { decorateStack: true });
     });
 
     // listen for debug output
@@ -322,7 +325,10 @@ export class IDLNotebookController {
    * We also do post-processing and, if we succeeded, we try to retrieve any
    * graphics.
    */
-  private async _endCellExecution(success: boolean, postExecute = true) {
+  private async _endCellExecution(
+    success: boolean,
+    actions: Partial<IEndCellExecutionActions> = {}
+  ) {
     /**
      * Get current cell
      */
@@ -348,6 +354,28 @@ export class IDLNotebookController {
     this._runtime.once(IDL_EVENT_LOOKUP.END, onDidCloseWhileBusy);
     this._runtime.once(IDL_EVENT_LOOKUP.CRASHED, onDidCloseWhileBusy);
 
+    // check if we need to decorate our call stack
+    if (this.isStarted() && actions.decorateStack) {
+      // get the current scope
+      const stack = (await this._runtime.getCurrentStack()).reverse();
+
+      // get the fsPath for the current cell
+      const fsPath = IDLFileHelper.notebookCellUriToFSPath(
+        cell.cell.document.uri
+      );
+
+      /** Check if any of our scope items are in this notebook cell */
+      const ourFile = stack.filter((item) => item.file === fsPath);
+
+      // see if we have a location we stopped on
+      if (ourFile.length > 0) {
+        IDL_DECORATIONS_MANAGER.addStackTraceDecorations(
+          cell.cell.document.uri,
+          ourFile.map((item) => item.line - 1) // in notebooks, we need zero-based instead of one
+        );
+      }
+    }
+
     /**
      * Wrap in try/catch so we don't have to worry about unhandled promise exceptions
      */
@@ -364,7 +392,7 @@ export class IDLNotebookController {
          *
          * Only run this if launched. If not launched we hang forever, for some reason
          */
-        if (success && postExecute && this.isStarted()) {
+        if (success && actions.postExecute && this.isStarted()) {
           await this.postCellExecution(success, cell);
         }
       }
@@ -380,25 +408,6 @@ export class IDLNotebookController {
 
     // always return from current scope
     if (this.isStarted()) {
-      // get the current scope
-      const stack = (await this._runtime.getCurrentStack()).reverse();
-
-      // get the fsPath for the current cell
-      const fsPath = IDLFileHelper.notebookCellUriToFSPath(
-        cell.cell.document.uri
-      );
-
-      /** Check if any of our scope items are in this notebook cell */
-      const ourFile = stack.filter((item) => item.file === fsPath);
-
-      // see if we have a location we stopped on
-      if (ourFile.length > 0) {
-        IDL_DECORATIONS_MANAGER.addStackTraceDecorations(
-          cell.cell.document.uri,
-          ourFile.map((item) => item.line - 1) // in notebooks, we need zero-based instead of one
-        );
-      }
-
       /**
        * Commands to run after executing a cell
        *
@@ -675,7 +684,7 @@ export class IDLNotebookController {
     );
 
     // make sure cells are done executing
-    this._endCellExecution(false, false);
+    this._endCellExecution(false, { postExecute: false });
 
     // return a prom
     return launchPromise;

libs/vscode/notebooks/src/lib/controller/idl-notebook-controller.interface.ts

@@ -15,3 +15,22 @@ export interface ICurrentCell {
   /** Did we succeed or not */
   success: boolean;
 }
+
+/**
+ * Actions we take at the end of a notebook cell being executed
+ */
+export interface IEndCellExecutionActions {
+  /**
+   * Do we try to add any graphics
+   */
+  postExecute?: boolean;
+  decorateStack?: boolean;
+}
+
+/**
+ * What do we do at the end of notebook cell execution?
+ */
+export const DEFAULT_END_CELL_EXECUTION_ACTIONS: IEndCellExecutionActions = {
+  postExecute: true,
+  decorateStack: false,
+};