Control session and kernel initialization individually per Cell (#276)

* feat: control isKernelAvaiable for each cell (#273)

* feat: improve event and add flag to allCellsReady (#273)

* feat: control isKernelAvaiable for each cell (#273)

---------

Co-authored-by: Marcos Alves <marcos@aevotech.com.br>

Author: MarcosVn

packages/react/src/components/cell/Cell.tsx

@@ -6,7 +6,6 @@
 
 import { useState, useEffect } from 'react';
 import { CodeCell, MarkdownCell } from '@jupyterlab/cells';
-import { KernelMessage } from '@jupyterlab/services';
 import { Box } from '@primer/react';
 import CellAdapter from './CellAdapter';
 import Lumino from '../lumino/Lumino';
@@ -61,24 +60,28 @@ export const Cell = (props: ICellProps) => {
 
     adapter.sessionContext.initialize().then(() => {
       if (!autoStart || !adapter.cell.model) {
-        return
+        return;
       }
-
+    
       // Perform auto-start for code or markdown cells
       if (adapter.cell instanceof CodeCell) {
-        const execute = CodeCell.execute(
+        CodeCell.execute(
           adapter.cell,
           adapter.sessionContext
         );
-        execute.then((msg: void | KernelMessage.IExecuteReplyMsg) => {
-          cellStore.setKernelAvailable(true);
-        });
       }
 
       if (adapter.cell instanceof MarkdownCell) {
         adapter.cell.rendered = true;
       }
     });
+
+    adapter.sessionContext.kernelChanged.connect(() => {
+      void adapter.sessionContext.session?.kernel?.info.then(info => {
+        // Set that session/kernel is ready for this cell when the kernel is guaranteed to be connected 
+        cellStore.setIsKernelSessionAvailable(id, true);
+      })
+    });
   }
 
   useEffect(() => {

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

@@ -268,7 +268,9 @@ export class CellAdapter {
         if (this._type === 'code') {
           const lang = info.language_info;
           const mimeType = mimeService.getMimeTypeByLanguage(lang);
-          this._cell.model.mimeType = mimeType;
+          if (this._cell.model) {
+            this._cell.model.mimeType = mimeType;
+          } 
         }
       });
     });

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

@@ -12,30 +12,45 @@ export interface ICellState {
   source?: string;
   outputsCount?: number;
   adapter?: CellAdapter;
+  isKernelSessionAvailable?: boolean; // Individual, for cell
 }
 
 export interface ICellsState {
   cells: Map<string, ICellState>;
-  kernelAvailable: boolean; // Currently shared between all cells
+  areAllKernelSessionsReady: boolean; // Control the state for all cells
 }
 
 export type CellState = ICellsState & {
   setCells: (cells: Map<string, ICellState>) => void;
   setSource: (id: string, source: string) => void;
   setOutputsCount: (id: string, outputsCount: number) => void;
-  setKernelAvailable: (kernelAvailable: boolean) => void;
+  setIsKernelSessionAvailable: (id: string, kernelAvailable: boolean) => void;
   setAdapter: (id: string, adapter?: CellAdapter) => void;
   getAdapter: (id: string) => CellAdapter | undefined;
   getSource: (id: string) => string | undefined;
   getOutputsCount: (id: string) => number | undefined;
+  getIsKernelSessionAvailable: (id: string) => boolean | undefined;
   execute: (id?: string) => void;
 };
 
+/**
+ * Iterate over all cells map and check if all cells/sessions are ready
+ */
+const areAllKernelSessionsAvailable = (cells: Map<string, ICellState>): boolean => {
+  for (const cell of cells.values()) {
+    if (!cell.isKernelSessionAvailable) {
+      return false;
+    }
+  }
+  return true;
+};
+
 export const cellStore = createStore<CellState>((set, get) => ({
   cells: new Map<string, ICellState>(),
   source: '',
   outputsCount: 0,
-  kernelAvailable: false,
+  isKernelSessionAvailable: false,
+  areAllKernelSessionsReady: false,
   adapter: undefined,
   setCells: (cells: Map<string, ICellState>) => set((cell: CellState) => ({ cells })),
 
@@ -59,7 +74,17 @@ export const cellStore = createStore<CellState>((set, get) => ({
     }
     set((state: CellState) => ({ cells }))
   },
-  setKernelAvailable: (kernelAvailable: boolean) => set((state: CellState) => ({ kernelAvailable })),
+  setIsKernelSessionAvailable: (id: string, isKernelSessionAvailable: boolean) => {
+    const cells = get().cells;
+    const cell = cells.get(id);
+    if (cell) {
+      cell.isKernelSessionAvailable = isKernelSessionAvailable;
+    } else {
+      cells.set(id, {isKernelSessionAvailable});
+    }
+    const areAllKernelSessionsReady = areAllKernelSessionsAvailable(cells);
+    set((cell: CellState) => ({ cells, areAllKernelSessionsReady }));
+  },
   setAdapter: (id: string, adapter?: CellAdapter) => {
     const cells = get().cells;
     const cell = cells.get(id);
@@ -79,6 +104,9 @@ export const cellStore = createStore<CellState>((set, get) => ({
   getOutputsCount: (id: string): number | undefined => {
     return get().cells.get(id)?.outputsCount;
   },
+  getIsKernelSessionAvailable: (id: string): boolean | undefined => {
+    return get().cells.get(id)?.isKernelSessionAvailable;
+  },
   execute: (id: string) => { 
     const cells = get().cells;
     const cell = cells.get(id);

packages/react/src/examples/All.tsx

@@ -58,7 +58,7 @@ const CellPreview = (props: ICellToolProps) => {
   return (
     <>
       <>source: {cellStore.getSource(props.id)}</>
-      <>kernel available: {String(cellStore.kernelAvailable)}</>
+      <>kernel available: {String(cellStore.getIsKernelSessionAvailable(props.id))}</>
     </>
   );
 };