@jupyterlab/notebook 4.6.0 → 4.6.1

package/src/actions.tsx

@@ -603,15 +603,77 @@ export namespace NotebookActions {
       lastIndex = notebook.model.cells.length;
     }
 
-    if (shift > 0) {
-      notebook.moveCell(firstIndex, lastIndex, lastIndex - firstIndex);
-    } else {
-      notebook.moveCell(firstIndex, firstIndex + shift, lastIndex - firstIndex);
-    }
+    const toIndex = shift > 0 ? lastIndex : firstIndex + shift;
+    moveCells(notebook, firstIndex, toIndex, lastIndex - firstIndex);
 
     void Private.handleState(notebook, state, true);
   }
 
+  /**
+   * Move cells while preserving in-flight kernel futures.
+   *
+   * The underlying `jupyter-ydoc` `moveCells` implementation currently
+   * serializes cells to JSON and recreates them via a delete + insert
+   * transaction, which disposes any active kernel futures attached to the
+   * old widgets. This wrapper detaches futures before the move and
+   * reattaches them to the new widgets afterwards, and stores them in the
+   * undo stack so that undoing the move also restores execution state.
+   *
+   * @param notebook - The target notebook.
+   * @param from - Index of the first cell to move.
+   * @param to - Target index (as passed to `notebook.moveCell`).
+   * @param n - Number of cells to move.
+   */
+  export function moveCells(
+    notebook: Notebook,
+    from: number,
+    to: number,
+    n = 1
+  ): void {
+    if (!notebook.model) {
+      return;
+    }
+
+    // Mirror `Notebook.moveCell`'s bounding/no-op logic so that we do not
+    // capture futures or touch the undo stack when the move won't happen.
+    // Otherwise we could attach execution metadata to an unrelated previous
+    // undo item, corrupting subsequent undo behavior.
+    const boundedTo = Math.min(
+      notebook.model.cells.length - 1,
+      Math.max(0, to)
+    );
+    if (boundedTo === from) {
+      return;
+    }
+
+    // moveCells serializes cells to JSON and recreates widgets (delete+insert),
+    // which would dispose any in-flight futures. Capture them first.
+    const storedExecutions: Private.IStoredCellExecution[] = [];
+    notebook.widgets.slice(from, from + n).forEach(child => {
+      if (!(child instanceof CodeCell)) {
+        return;
+      }
+      const stored = Private.captureExecution(child);
+      if (stored) {
+        storedExecutions.push(stored);
+      }
+    });
+
+    notebook.moveCell(from, to, n);
+
+    // Immediately reconnect futures to the newly created widgets.
+    for (const stored of storedExecutions) {
+      Private.restoreExecution(notebook, stored);
+    }
+
+    // Store in the undo stack so that undoing the move can also restore state.
+    if (storedExecutions.length > 0) {
+      const undoManager = (notebook.model.sharedModel as YNotebook).undoManager;
+      const lastItem = undoManager.undoStack[undoManager.undoStack.length - 1];
+      lastItem?.meta.set(Private.CELL_EXECUTION_META_KEY, storedExecutions);
+    }
+  }
+
   /**
    * Move the selected cell(s) down.
    *
@@ -1609,9 +1671,37 @@ export namespace NotebookActions {
     const state = Private.getState(notebook);
     notebook.mode = 'command';
 
+    const undoManager = (notebook.model.sharedModel as YNotebook).undoManager;
+
+    // For cells that will be MOVED by the undo (i.e. they still exist in the
+    // notebook at their current position), pre-capture their futures now before
+    // sharedModel.undo() destroys those widgets. This prevents OutputArea.dispose()
+    // from cancelling the future during the Y.js delete+insert that implements the move.
+    const topItem = undoManager.undoStack[undoManager.undoStack.length - 1];
+    const pendingExecutions = topItem?.meta.get(
+      Private.CELL_EXECUTION_META_KEY
+    ) as Private.IStoredCellExecution[] | undefined;
+    const preCaptured = new Map<string, Private.IStoredCellExecution>();
+    pendingExecutions?.forEach(stored => {
+      const cell = notebook.widgets.find(w => w.model.id === stored.cellId);
+      if (!(cell instanceof CodeCell)) {
+        return; // cell was deleted (not moved) — handled via stored future below
+      }
+      // Cell still present → move undo. Fresh capture protects the future.
+      const fresh = Private.captureExecution(cell) ?? {
+        ...stored,
+        isDone: () => true,
+        buffered: []
+      };
+      // The undo will roll the outputs back to their state at the time of
+      // the move; snapshot the current outputs so anything received since
+      // then can be re-applied after the undo.
+      fresh.outputs = cell.model.outputs.toJSON();
+      preCaptured.set(stored.cellId, fresh);
+    });
+
     // Capture execution context from the stack item being popped.
     let storedExecutions: Private.IStoredCellExecution[] | undefined;
-    const undoManager = (notebook.model.sharedModel as YNotebook).undoManager;
     const onStackItemPopped = ({
       stackItem
     }: {
@@ -1625,40 +1715,13 @@ export namespace NotebookActions {
     notebook.model.sharedModel.undo();
     undoManager.off('stack-item-popped', onStackItemPopped);
 
-    // Restore execution state on resurrected cell widgets.
-    storedExecutions?.forEach(({ cellId, future, isDone, buffered }) => {
-      const cell = notebook.widgets.find(w => w.model.id === cellId);
-      if (!(cell instanceof CodeCell)) {
-        return;
-      }
-      if (isDone()) {
-        // Execution finished or was interrupted before undo - ensure state is idle.
-        cell.model.executionState = 'idle';
-        return;
-      }
-      // Still running - reconnect the future so the cell receives remaining
-      // output and stdin requests (e.g. input()), and tracks the idle transition.
-      cell.outputArea.reattachFuture(future);
-      // Replay any IOPub messages that arrived while the future was detached.
-      // After reattachFuture, future.onIOPub routes to the new output area.
-      for (const msg of buffered) {
-        void future.onIOPub(msg);
-      }
-      // The original execute() call targeted the old cell widget, so it will
-      // not update executionCount/executionState on this resurrected cell.
-      // Drive the idle transition here instead.
-      const cellRef = cell;
-      void future.done.then(
-        reply => {
-          if (!cellRef.isDisposed) {
-            cellRef.model.executionCount = reply.content.execution_count;
-          }
-        },
-        () => {
-          if (!cellRef.isDisposed) {
-            cellRef.model.executionState = 'idle';
-          }
-        }
+    // Restore execution state on resurrected/moved cell widgets.
+    // For move-undo: use freshly pre-captured data (stored data is stale).
+    // For delete-undo: stored data has the futures captured at deletion time.
+    storedExecutions?.forEach(stored => {
+      Private.restoreExecution(
+        notebook,
+        preCaptured.get(stored.cellId) ?? stored
       );
     });
 
@@ -2568,6 +2631,15 @@ namespace Private {
   /** Key used to store cell execution state in Y.js undo stack item metadata. */
   export const CELL_EXECUTION_META_KEY = Symbol('cellExecutionState');
 
+  /**
+   * A kernel message that arrived while the future was detached, tagged with
+   * its channel so it can be dispatched to the right handler on replay.
+   */
+  export type IBufferedMessage =
+    | { channel: 'iopub'; msg: KernelMessage.IIOPubMessage }
+    | { channel: 'stdin'; msg: KernelMessage.IStdinMessage }
+    | { channel: 'reply'; msg: KernelMessage.IExecuteReplyMsg };
+
   export interface IStoredCellExecution {
     cellId: string;
     future: Kernel.IShellFuture<
@@ -2575,13 +2647,28 @@ namespace Private {
       KernelMessage.IExecuteReplyMsg
     >;
     isDone: () => boolean;
-    /** IOPub messages that arrived while the future was detached. */
-    buffered: KernelMessage.IIOPubMessage[];
+    /**
+     * Kernel messages (IOPub, stdin and reply) that arrived while the future
+     * was detached, in arrival order, so they can be replayed on reattach.
+     */
+    buffered: IBufferedMessage[];
+    /**
+     * Snapshot of the cell outputs to restore after an undo.
+     *
+     * The Y.js undo of a move (a delete + insert transaction) resurrects the
+     * cell as it was when the move happened, rolling back any output received
+     * since. Re-applying this snapshot after the undo prevents that loss.
+     */
+    outputs?: nbformat.IOutput[];
   }
 
   /**
-   * Detach the kernel future from a code cell, buffering any IOPub messages
-   * that arrive while it is detached so they can be replayed on reattach.
+   * Detach the kernel future from a code cell, buffering any messages that
+   * arrive while it is detached so they can be replayed on reattach.
+   *
+   * `detachFuture` clears the IOPub, stdin and reply handlers, so all three
+   * channels are buffered here; otherwise a message arriving during the
+   * detached window (e.g. an `input()` request on stdin) would be dropped.
    *
    * Returns null if the cell has no active future.
    */
@@ -2596,13 +2683,88 @@ namespace Private {
     void future.done.finally(() => {
       done = true;
     });
-    const buffered: KernelMessage.IIOPubMessage[] = [];
+    const buffered: IBufferedMessage[] = [];
     future.onIOPub = msg => {
-      buffered.push(msg);
+      buffered.push({ channel: 'iopub', msg });
+    };
+    future.onStdin = msg => {
+      buffered.push({ channel: 'stdin', msg });
+    };
+    future.onReply = msg => {
+      buffered.push({ channel: 'reply', msg });
     };
     return { cellId: cell.model.id, future, isDone: () => done, buffered };
   }
 
+  /**
+   * Reconnect a captured execution to the cell widget that now holds the model.
+   *
+   * Handles both the "still running" and "already finished" cases.
+   */
+  export function restoreExecution(
+    notebook: Notebook,
+    { cellId, future, isDone, buffered, outputs }: IStoredCellExecution
+  ): void {
+    const cell = notebook.widgets.find(w => w.model.id === cellId);
+    if (!(cell instanceof CodeCell)) {
+      return;
+    }
+    if (outputs && !JSONExt.deepEqual(outputs, cell.model.outputs.toJSON())) {
+      // Re-apply the output snapshot taken just before the undo: the Y.js
+      // undo rolled the outputs back to their state at the time of the
+      // undone action. Going through the output area model keeps the
+      // in-memory model and the shared model in sync (output changes are
+      // not tracked by the undo manager, so this does not pollute history).
+      cell.model.outputs.fromJSON(outputs);
+    }
+    // Reattach the future (without clearing existing outputs) and replay any
+    // messages buffered while it was detached, in arrival order. This is done
+    // whether or not the execution has already finished, so that final outputs
+    // (or a pending stdin request) that arrived while detached are not lost
+    // (e.g. if the kernel completed in the brief window during the undo).
+    cell.outputArea.reattachFuture(future);
+    for (const buf of buffered) {
+      switch (buf.channel) {
+        case 'iopub':
+          void future.onIOPub(buf.msg);
+          break;
+        case 'stdin':
+          void future.onStdin(buf.msg);
+          break;
+        case 'reply':
+          void future.onReply(buf.msg);
+          break;
+      }
+    }
+    if (isDone()) {
+      // The execution already finished (e.g. it completed or was interrupted
+      // during the undo). The resurrected cell may carry a stale 'running'
+      // state from its restored snapshot, so reset it synchronously rather
+      // than relying solely on the asynchronous `future.done` handler below.
+      cell.model.executionState = 'idle';
+    } else {
+      // Restore the running state on the recreated cell widget.
+      cell.model.executionState = 'running';
+    }
+    const cellRef = cell;
+    void future.done.then(
+      reply => {
+        if (!cellRef.isDisposed) {
+          // The future is authoritative for the prompt number; a snapshot
+          // taken before completion would be stale (still null). Setting a
+          // non-null execution count also flips the state back to 'idle'.
+          cellRef.model.executionCount = reply.content.execution_count;
+          cellRef.model.executionState = 'idle';
+        }
+      },
+      () => {
+        if (!cellRef.isDisposed) {
+          cellRef.model.executionState = 'idle';
+        }
+      }
+    );
+  }
+
   /**
    * Notebook cell executor
    */

package/src/widget.ts

@@ -3374,8 +3374,8 @@ export class Notebook extends StaticNotebook {
         return;
       }
 
-      // Move the cells one by one
-      this.moveCell(fromIndex, toIndex, toMove.length);
+      // Move the selected block of cells, preserving in-flight executions.
+      NotebookActions.moveCells(this, fromIndex, toIndex, toMove.length);
     } else {
       // Handle the case where we are copying cells between
       // notebooks.