From cab8c8acfa7e4474ca55aa9842c814a069598056 Mon Sep 17 00:00:00 2001
From: Marcos Alves <marcos@aevotech.com.br>
Date: Wed, 31 Jul 2024 20:25:26 -0300
Subject: [PATCH] feat: add state control for cell execution (#279)

---
 .../react/src/components/cell/CellAdapter.ts  |  5 +++
 .../react/src/components/cell/CellState.ts    | 39 +++++++++++++++++--
 2 files changed, 41 insertions(+), 3 deletions(-)

diff --git a/packages/react/src/components/cell/CellAdapter.ts b/packages/react/src/components/cell/CellAdapter.ts
index 81c4a12b..6de6b818 100755
--- a/packages/react/src/components/cell/CellAdapter.ts
+++ b/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,6 +482,7 @@ export class CellAdapter {
           finished || new Date().toISOString();
         model.setMetadata('execution', timingInfo);
       }
+      cellsStore.getState().setIsExecuting(this._id, false);
       return executeReplyMessage;
     } catch (e) {
       // If we started executing, and the cell is still indicating this execution, clear the prompt.
@@ -494,6 +498,7 @@ export class CellAdapter {
           model.setMetadata('execution', timingInfo);
         }
       }
+      cellsStore.getState().setIsExecuting(this._id, false);
       throw e;
     }
   }
diff --git a/packages/react/src/components/cell/CellState.ts b/packages/react/src/components/cell/CellState.ts
index 459cebcd..3702a441 100644
--- a/packages/react/src/components/cell/CellState.ts
+++ b/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;