@marimo-team/islands 0.25.3-dev266 → 0.25.3-dev268

package/src/components/editor/cell/RunButton.tsx

@@ -1,7 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { HardDriveDownloadIcon, PlayIcon } from "lucide-react";
 import type { JSX } from "react";
-import type { CellConfig, RuntimeState } from "@/core/network/types";
+import type { CellRunStatus } from "@/core/executions";
+import type { CellConfig } from "@/core/network/types";
 import {
   getConnectionTooltip,
   isAppInteractionDisabled,
@@ -35,20 +36,29 @@ function computeColor({
 
 export const RunButton = (props: {
   edited: boolean;
-  status: RuntimeState;
+  status: CellRunStatus;
+  /** Whether a cell it depends on is disabled. */
+  blocked: boolean;
   needsRun: boolean;
   connectionState: ConnectionState;
   config: CellConfig;
   onClick?: () => void;
 }): JSX.Element => {
-  const { onClick, connectionState, needsRun, status, config, edited } = props;
+  const {
+    onClick,
+    connectionState,
+    needsRun,
+    status,
+    blocked,
+    config,
+    edited,
+  } = props;
 
-  const blockedStatus = status === "disabled-transitively";
   const loading = status === "running" || status === "queued";
   const inactive =
     isAppInteractionDisabled(connectionState) ||
     loading ||
-    (!config.disabled && blockedStatus && !edited);
+    (!config.disabled && blocked && !edited);
   const variant = computeColor({
     connectionState,
     needsRun,
@@ -69,7 +79,7 @@ export const RunButton = (props: {
       </ToolbarItem>
     );
   }
-  if (!config.disabled && blockedStatus && !edited) {
+  if (!config.disabled && blocked && !edited) {
     return (
       <ToolbarItem
         disabled={inactive}

package/src/components/editor/cell/StagedAICell.tsx

@@ -1,24 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { useAtomValue, useStore } from "jotai";
-import { SparklesIcon } from "lucide-react";
-import { Tooltip } from "@/components/ui/tooltip";
+import { atom, useAtomValue, useStore } from "jotai";
+import { useMemo } from "react";
 import {
   type Edit,
+  stagedAICellsAtom,
   useStagedAICell,
   useStagedCells,
 } from "@/core/ai/staged-cells";
 import { chatGenerationInProgressAtom } from "@/core/ai/state";
-import { getCellEditorView } from "@/core/cells/cells";
 import type { CellId } from "@/core/cells/ids";
-import { prepareCellCodeUpdate } from "@/core/codemirror/language/cell-code";
-import { cn } from "@/utils/cn";
-import { Logger } from "@/utils/Logger";
-import {
-  AI_EDIT_ACTION_LABELS,
-  CompletionActionsCellFooter,
-} from "../ai/completion-handlers";
-import { useRunCell } from "./useRunCells";
+import { AI_EDIT_ACTION_LABELS } from "../ai/completion-handlers";
+import type { ActionButton } from "../actions/types";
 
 const actionLabels: Record<Edit["type"], { accept: string; decline: string }> =
   {
@@ -27,142 +20,56 @@ const actionLabels: Record<Edit["type"], { accept: string; decline: string }> =
     delete_cell: { accept: "Delete cell", decline: "Keep cell" },
   };
 
-export const StagedAICellBackground: React.FC<{
-  cellId: CellId;
-  className?: string;
-}> = ({ cellId, className }) => {
-  const stagedCell = useStagedAICell(cellId);
-
-  if (!stagedCell) {
-    return null;
-  }
-
-  const cellClass =
-    stagedCell.type === "delete_cell"
-      ? "mo-ai-deleted-cell"
-      : "mo-ai-generated-cell";
-
-  return <div className={cn(cellClass, className)} />;
-};
-
-export const StagedAICellFooter: React.FC<{ cellId: CellId }> = ({
+export const StagedAIDeletionMarker: React.FC<{ cellId: CellId }> = ({
   cellId,
-}) => {
+}) =>
+  useStagedAICell(cellId)?.type === "delete_cell" ? (
+    <div className="mo-ai-deleted-cell" />
+  ) : null;
+
+export function useStagedAICellActions(
+  cellId: CellId | undefined,
+): ActionButton[] {
   const store = useStore();
-  const stagedAiCell = useStagedAICell(cellId);
-  const isGenerating = useAtomValue(chatGenerationInProgressAtom);
-  const runCell = useRunCell(cellId);
-
-  const { deleteStagedCell, removeStagedCell } = useStagedCells(store);
-
-  if (!stagedAiCell) {
-    return null;
-  }
-
-  const isDeletion = stagedAiCell.type === "delete_cell";
-  const labels = actionLabels[stagedAiCell.type];
-
-  const handleCompletion = (type: "accept" | "reject") => {
-    const completionFunc =
-      type === "accept" ? acceptStagedCell : rejectStagedCell;
-    completionFunc(cellId, stagedAiCell, removeStagedCell, deleteStagedCell);
-  };
-
-  const tooltipContent =
-    stagedAiCell.type === "add_cell"
-      ? "AI-generated cell"
-      : isDeletion
-        ? "AI suggests deleting this cell"
-        : "AI changed this cell";
-
-  return (
-    <div className="mo-ai-cell-footer">
-      <div className="flex items-center gap-1.5 text-xs text-muted-foreground">
-        <Tooltip content={tooltipContent}>
-          <span
-            className="inline-flex items-center gap-1.5"
-            aria-label={tooltipContent}
-          >
-            <SparklesIcon
-              aria-hidden="true"
-              className="size-3.5 text-(--blue-11)"
-            />
-            <span>AI</span>
-          </span>
-        </Tooltip>
-      </div>
-      <div className="flex items-center gap-1.5">
-        <CompletionActionsCellFooter
-          isLoading={isGenerating}
-          onAccept={() => handleCompletion("accept")}
-          onDecline={() => handleCompletion("reject")}
-          size="xs"
-          runCell={isDeletion ? undefined : runCell}
-          acceptLabel={labels.accept}
-          declineLabel={labels.decline}
-        />
-      </div>
-    </div>
+  const edit = useStagedAICell(cellId);
+  // Every cell calls this; only cells with a pending edit care about the lock.
+  const isLocked = useAtomValue(
+    useMemo(
+      () =>
+        atom(
+          (get) =>
+            cellId !== undefined &&
+            get(stagedAICellsAtom).has(cellId) &&
+            get(chatGenerationInProgressAtom),
+        ),
+      [cellId],
+    ),
   );
-};
+  const { acceptStagedCell, rejectStagedCell } = useStagedCells();
 
-/**
- * Accept a staged cell and apply the changes.
- */
-export function acceptStagedCell(
-  cellId: CellId,
-  edit: Edit,
-  removeStagedCell: (cellId: CellId) => void,
-  deleteStagedCell: (cellId: CellId) => void,
-): void {
-  switch (edit.type) {
-    case "delete_cell":
-      // For delete cells, the cell is deleted when the completion is accepted
-      deleteStagedCell(cellId);
-      break;
-    default:
-      removeStagedCell(cellId);
-      break;
+  if (!edit || cellId === undefined) {
+    return [];
   }
-}
-
-/**
- * Reject a staged cell and revert the changes.
- */
-export function rejectStagedCell(
-  cellId: CellId,
-  edit: Edit,
-  removeStagedCell: (cellId: CellId) => void,
-  deleteStagedCell: (cellId: CellId) => void,
-): void {
-  switch (edit.type) {
-    case "update_cell": {
-      // Revert cell code
-      const editorView = getCellEditorView(cellId);
-      if (!editorView) {
-        Logger.error("Editor for this cell not found", { cellId });
-        break;
-      }
 
-      const update = prepareCellCodeUpdate(
-        editorView.state,
-        edit.previousCode,
-        { initialize: true },
-      );
-
-      if (update) {
-        editorView.dispatch(update);
-      }
-      removeStagedCell(cellId);
-      break;
-    }
-    case "add_cell":
-      // Delete the cell since it's newly created
-      deleteStagedCell(cellId);
-      break;
-    case "delete_cell":
-      // Just remove the deletion marker - cell stays in notebook
-      removeStagedCell(cellId);
-      break;
-  }
+  const labels = actionLabels[edit.type];
+  return [
+    {
+      label: labels.accept,
+      disabled: isLocked,
+      handle: () => {
+        if (!store.get(chatGenerationInProgressAtom)) {
+          acceptStagedCell(cellId, edit);
+        }
+      },
+    },
+    {
+      label: labels.decline,
+      disabled: isLocked,
+      handle: () => {
+        if (!store.get(chatGenerationInProgressAtom)) {
+          rejectStagedCell(cellId, edit);
+        }
+      },
+    },
+  ];
 }

package/src/components/editor/cell/__tests__/useRunCells.test.ts

@@ -0,0 +1,45 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it, vi } from "vitest";
+import { MockExecutions } from "@/__mocks__/executions";
+import { MockNotebook } from "@/__mocks__/notebook";
+import { cellId } from "@/__tests__/branded";
+import {
+  cellExecutionInputsAtom,
+  cellExecutionView,
+  notebookQueuedOrRunningCountAtom,
+} from "@/core/executions";
+import { store } from "@/core/state/jotai";
+import { runCells } from "../useRunCells";
+
+const ENABLED = cellId("enabled");
+const DISABLED = cellId("disabled");
+
+describe("runCells", () => {
+  it("sends a disabled cell to the kernel without reading it as queued", async () => {
+    MockExecutions.seed(store, []);
+    const notebook = MockNotebook.notebookState({
+      cellData: {
+        [ENABLED]: { code: "x = 1" },
+        [DISABLED]: { code: "y = 2", config: { disabled: true } },
+      },
+    });
+    const sendRun = vi.fn(async () => null);
+
+    await runCells({
+      cellIds: [ENABLED, DISABLED],
+      sendRun,
+      prepareForRun: vi.fn(),
+      notebook,
+    });
+
+    expect(sendRun).toHaveBeenCalledWith({
+      cellIds: [ENABLED, DISABLED],
+      codes: ["x = 1", "y = 2"],
+    });
+    const inputs = store.get(cellExecutionInputsAtom);
+    expect(cellExecutionView(inputs, ENABLED).status).toBe("queued");
+    expect(cellExecutionView(inputs, DISABLED).status).toBe("idle");
+    expect(store.get(notebookQueuedOrRunningCountAtom)).toBe(1);
+  });
+});

package/src/components/editor/cell/cell-context-menu.tsx

@@ -23,6 +23,7 @@ import { toast } from "@/components/ui/use-toast";
 import { useCellData, useCellRuntime } from "@/core/cells/cells";
 import { CellOutputId } from "@/core/cells/ids";
 import { isOutputEmpty } from "@/core/cells/outputs";
+import { useCellExecution } from "@/core/executions";
 import { goToDefinitionWithLspFallback } from "@/core/codemirror/go-to-definition/utils";
 import { sendToPanelManager } from "@/core/vscode/vscode-bindings";
 import { copyImageToClipboard, copyToClipboard } from "@/utils/copy";
@@ -54,12 +55,14 @@ export const CellActionsContextMenu = ({
 }: Props) => {
   const cellData = useCellData(cellId);
   const cellRuntime = useCellRuntime(cellId);
+  const execution = useCellExecution(cellId);
   const actions = useCellActionButtons({
     cell: {
       cellId: cellId,
       name: cellData.name,
       config: cellData.config,
-      status: cellRuntime.status,
+      status: execution.status,
+      blocked: execution.blocked,
       hasOutput: !isOutputEmpty(cellRuntime.output),
       hasConsoleOutput: cellRuntime.consoleOutputs.length > 0,
       getEditorView,

package/src/components/editor/cell/code/cell-editor.tsx

@@ -19,7 +19,8 @@ import {
   useNotebookDocumentEdit,
 } from "@/core/notebook/document";
 import { usePendingDeleteService } from "@/core/cells/pending-delete-service";
-import type { CellData, CellRuntimeState } from "@/core/cells/types";
+import type { CellData } from "@/core/cells/types";
+import type { CellRunStatus } from "@/core/executions";
 import { setupCodeMirror } from "@/core/codemirror/cm";
 import { acceptCompletionOnEnterAtom } from "@/core/codemirror/completion/accept-on-enter-atom";
 import { editorMountScheduler } from "@/core/codemirror/editor-mount-scheduler";
@@ -54,10 +55,11 @@ import { useSplitCellCallback } from "../useSplitCell";
 import { CodePlaceholder } from "./code-placeholder";
 import { LanguageToggles } from "./language-toggle";
 
-export interface CellEditorProps
-  extends
-    Pick<CellRuntimeState, "status">,
-    Pick<CellData, "id" | "code" | "config"> {
+export interface CellEditorProps extends Pick<
+  CellData,
+  "id" | "code" | "config"
+> {
+  status: CellRunStatus;
   runCell: () => void;
   theme: Theme;
   showPlaceholder: boolean;

package/src/components/editor/cell/collapse.tsx

@@ -9,6 +9,7 @@ import {
   RefreshCcw,
 } from "lucide-react";
 import type React from "react";
+import { useAtomValue } from "jotai";
 import { memo } from "react";
 import { Toolbar } from "@/components/layout/toolbar";
 import { Button } from "@/components/ui/button";
@@ -17,6 +18,7 @@ import { useNotebook } from "@/core/cells/cells";
 import { columnOf, useLayout } from "@/core/cells/layout";
 import type { CellId } from "@/core/cells/ids";
 import { getDescendantsStatus } from "@/core/cells/utils";
+import { cellExecutionInputsAtom } from "@/core/executions";
 import { cn } from "@/utils/cn";
 
 interface Props {
@@ -65,9 +67,11 @@ export const CollapsedCellBanner: React.FC<{
   count: number;
 }> = memo(({ onClick, count, cellId }) => {
   const notebook = useNotebook();
+  const executions = useAtomValue(cellExecutionInputsAtom);
   const layout = useLayout();
   const states = getDescendantsStatus(
     notebook,
+    executions,
     columnOf(layout, cellId).folds.get(cellId) ?? [],
   );
 

package/src/components/editor/cell/useRunCells.ts

@@ -7,10 +7,16 @@ import { getNotebook, type NotebookState } from "@/core/cells/cells";
 import { useCellActions } from "@/core/cells/actions";
 import type { CellId } from "@/core/cells/ids";
 import { enabledCellIds, staleCellIds } from "@/core/cells/utils";
+import {
+  cellExecutionInputsAtom,
+  submitCells,
+  withdrawCells,
+} from "@/core/executions";
 import { getCurrentLanguageAdapter } from "@/core/codemirror/language/commands";
 import { getEditorCodeAsPython } from "@/core/codemirror/language/utils";
 import { useRequestClient } from "@/core/network/requests";
 import type { ExecuteCellsRequest } from "@/core/network/types";
+import { store } from "@/core/state/jotai";
 import { Logger } from "@/utils/Logger";
 
 export const hasRunAnyCellAtom = atom<boolean>(false);
@@ -20,7 +26,9 @@ export const hasRunAnyCellAtom = atom<boolean>(false);
  */
 export function useRunStaleCells() {
   const runCells = useRunCells();
-  const runStaleCells = useEvent(() => runCells(staleCellIds(getNotebook())));
+  const runStaleCells = useEvent(() =>
+    runCells(staleCellIds(getNotebook(), store.get(cellExecutionInputsAtom))),
+  );
   return runStaleCells;
 }
 
@@ -107,9 +115,16 @@ export async function runCells({
     codes.push(code);
     prepareForRun({ cellId });
   }
-
-  // Send the run request to the Kernel
+  // The kernel sets a disabled cell aside without planning an execution,
+  // so a submission for it would never be superseded.
+  submitCells(
+    store,
+    cellIds.filter((cellId) => !cellData[cellId]?.config.disabled),
+  );
+
+  // TODO(MO-8346): Submit through the host's `runtime.run`.
   await sendRun({ cellIds: cellIds, codes: codes }).catch((error) => {
+    withdrawCells(store, cellIds);
     Logger.error(error);
   });
 }