[v2] Control when a Cell is executing #279 (#285)

* feat: add state control for cell execution (#279)

* feat: add example for cell execution control (#279)

* fix: reject promise error cases in onReply (#279)

* feat: improve kernelexecutor error reject (#279)

---------

Co-authored-by: Marcos Alves <[email protected]>

packages/react/src/components/cell/CellAdapter.ts

@@ -55,6 +55,7 @@ import ClassicWidgetManager from '../../jupyter/ipywidgets/classic/manager';
 import Kernel from '../../jupyter/kernel/Kernel';
 import getMarked from '../notebook/marked/marked';
 import CellCommands from './CellCommands';
+import { cellsStore } from './CellState';
 
 interface BoxOptions {
   showToolbar?: boolean;
@@ -393,12 +394,14 @@ export class CellAdapter {
     cell: CodeCell,
     metadata?: JSONObject
   ): Promise<KernelMessage.IExecuteReplyMsg | void> {
+    cellsStore.getState().setIsExecuting(this._id, true);
     const model = cell.model;
     const code = model.sharedModel.getSource();
     if (!code.trim() || !this.kernel) {
       model.sharedModel.transact(() => {
         model.clearExecution();
       }, false);
+      cellsStore.getState().setIsExecuting(this._id, false);
       return new Promise(() => {});
     }
     const cellId = { cellId: model.sharedModel.getId() };
@@ -479,8 +482,10 @@ export class CellAdapter {
           finished || new Date().toISOString();
         model.setMetadata('execution', timingInfo);
       }
+      cellsStore.getState().setIsExecuting(this._id, false);
       return executeReplyMessage;
     } catch (e) {
+      cellsStore.getState().setIsExecuting(this._id, false);
       // If we started executing, and the cell is still indicating this execution, clear the prompt.
       if (future && !cell.isDisposed && cell.outputArea.future === future) {
         cell.setPrompt('');

packages/react/src/components/cell/CellState.ts

@@ -8,16 +8,20 @@ import { createStore } from 'zustand/vanilla';
 import { useStore } from 'zustand';
 import { CellAdapter } from './CellAdapter';
 
+// Individual, for cell
 export interface ICellState {
   source?: string;
   outputsCount?: number;
   adapter?: CellAdapter;
-  isKernelSessionAvailable?: boolean; // Individual, for cell
+  isKernelSessionAvailable?: boolean;
+  isExecuting?: boolean;
 }
 
+// For all cells
 export interface ICellsState {
   cells: Map<string, ICellState>;
-  areAllKernelSessionsReady: boolean; // Control the state for all cells
+  areAllKernelSessionsReady: boolean;
+  isAnyCellExecuting: boolean;
 }
 
 export type CellsState = ICellsState & {
@@ -31,6 +35,7 @@ export type CellsState = ICellsState & {
   getOutputsCount: (id: string) => number | undefined;
   isKernelSessionAvailable: (id: string) => boolean | undefined;
   execute: (id?: string) => void;
+  setIsExecuting: (id: string, isExecuting: boolean) => void;
 };
 
 /**
@@ -45,12 +50,26 @@ const areAllKernelSessionsAvailable = (cells: Map<string, ICellState>): boolean
   return true;
 };
 
+/**
+ * Check if any cell is currently executing
+ */
+export const isAnyCellRunning = (cells: Map<string, ICellState>): boolean => {
+  for (const cell of cells.values()) {
+    if (cell.isExecuting) {
+      return true;
+    }
+  }
+  return false;
+};
+
+
 export const cellsStore = createStore<CellsState>((set, get) => ({
   cells: new Map<string, ICellState>(),
   source: '',
   outputsCount: 0,
-  areAllKernelSessionsReady: false,
   adapter: undefined,
+  areAllKernelSessionsReady: false, // prop refers to all cells
+  isAnyCellExecuting: false, // prop refers to all cells,
   setCells: (cells: Map<string, ICellState>) => set((cell: CellsState) => ({ cells })),
   setSource: (id: string, source: string) => {
     const cells = get().cells;
@@ -114,6 +133,20 @@ export const cellsStore = createStore<CellsState>((set, get) => ({
       get().cells.forEach((cell) => cell.adapter?.execute());
     }
   },
+  setIsExecuting: (id: string, isExecuting: boolean) => {
+    const cells = get().cells;
+    const cell = cells.get(id);
+    if (cell) {
+      cell.isExecuting = isExecuting;
+    } else {
+      get().cells.forEach((cell) => cell.adapter?.execute())
+      cells.set(id, { isExecuting });
+    }
+
+    // Also update isAnyCellRunning state (for all cells)
+    const isAnyCellExecuting = isAnyCellRunning(cells);
+    set((state: CellsState) => ({ cells, isAnyCellExecuting }));
+  },
 }));
 
 export function useCellsStore(): CellsState;

packages/react/src/examples/CellExecuteControl.tsx

@@ -0,0 +1,68 @@
+/*
+ * Copyright (c) 2021-2023 Datalayer, Inc.
+ *
+ * MIT License
+ */
+import React, {useEffect} from 'react';
+import { createRoot } from 'react-dom/client';
+
+import {Button, Box} from '@primer/react';
+
+import Jupyter from '../jupyter/Jupyter';
+import Cell from '../components/cell/Cell';
+import { cellsStore, ICellsState } from '../components/cell/CellState';
+
+
+const div = document.createElement('div');
+document.body.appendChild(div);
+const root = createRoot(div);
+
+const btnStyle = {
+    marginLeft: '50px',
+    marginTop: '20px'
+}
+
+const CELL_CODE = `import time\ntime.sleep(3)`
+
+const CellExecuteControl = () => {
+  const [executionDisable, setExecutionDisable] = React.useState(false);
+
+  useEffect(() => {
+    const handleChange = (newState: ICellsState) => {
+      setExecutionDisable(newState.isAnyCellExecuting);
+    };
+
+    const unsubscribe = cellsStore.subscribe(handleChange);
+    return () => {
+      unsubscribe();
+    };
+  }, []);
+
+  const onExecuteClick = () => {
+    cellsStore.getState().execute();
+  }
+
+  return (
+    <Jupyter>
+      <Box style={{marginTop: '20px'}}>
+        <Cell 
+          id='1' 
+          type='code'
+          source={CELL_CODE} 
+          autoStart={false} 
+          showToolbar={false} />
+        <Cell 
+          id='2' 
+          type='code'
+          autoStart={false} 
+          showToolbar={false} />
+        <Button 
+          onClick={onExecuteClick} 
+          disabled={executionDisable} 
+          style={btnStyle}>Execute all</Button>
+      </Box>
+    </Jupyter>
+  )
+};
+
+root.render(<CellExecuteControl/>);

packages/react/src/jupyter/kernel/KernelExecutor.ts

@@ -258,9 +258,7 @@ export class KernelExecutor {
           break;
         case 'error':
           {
-            const { ename, evalue, traceback } = (
-              content as any as KernelMessage.IErrorMsg
-            ).content;
+            const { ename, evalue, traceback } = content as KernelMessage.IReplyErrorContent;
             this._executed.reject(
               `${ename}: ${evalue}\n${(traceback ?? []).join('\n')}`
             );