@marimo-team/islands 0.25.3-dev308 → 0.25.3-dev310

package/src/core/executions/atoms.ts

@@ -5,19 +5,18 @@
 import type * as Protocol from "@marimo-team/client/protocol";
 import { atom, type createStore } from "jotai";
 import type { CellMessage } from "@/core/kernel/messages";
-import type {
-  Executions,
-  NotebookReplica,
-  Results,
+import {
+  ConnectionState,
+  type Executions,
+  type NotebookReplica,
+  type Results,
 } from "@/core/notebook/replica";
 import { activeNotebookReplicaAtom } from "@/core/session/active-notebook";
-import { waitFor } from "@/core/state/jotai";
 import { type CellFreshnessInputs, freshnessOf } from "./freshness";
 import type { CellOutputInputs } from "./output";
 import {
   anyRunning,
   type CellExecutionInputs,
-  latestExecutionId,
   queuedOrRunningCount,
 } from "./select";
 
@@ -30,40 +29,52 @@ const NO_EXECUTIONS: Executions = {
 };
 const NO_RESULTS: Results = new Map();
 const NO_CELLS: ReadonlyMap<Protocol.CellId, Protocol.Cell> = new Map();
-const NO_SUBMISSIONS: ReadonlyMap<
-  Protocol.CellId,
-  Protocol.ExecutionId | null
-> = new Map();
+const NO_REQUESTS: ReadonlySet<RunRequest> = new Set();
+const NO_REQUESTED: ReadonlySet<Protocol.CellId> = new Set();
 const NO_BLOCKED: ReadonlySet<Protocol.CellId> = new Set();
 const NO_CLEARED: ReadonlyMap<Protocol.CellId, Protocol.ExecutionId> =
   new Map();
 
+/** The notebook a replica is showing, which the editor's marks belong to. */
+interface ShownNotebook {
+  readonly replica: NotebookReplica;
+  readonly id: Protocol.NotebookId;
+}
+
 /**
  * What the editor marks about cells of one replica's notebook.
  *
- * The marks are keyed by cell id, which another notebook, or the same one
+ * The marks refer to cells by id, which another notebook, or the same one
  * after a reconnect, reuses with different executions behind it. They
  * therefore count only for the replica and host notebook they belong to.
  */
 type Marks<T> = {
-  readonly notebook: {
-    readonly replica: NotebookReplica;
-    readonly id: Protocol.NotebookId;
-  };
+  readonly notebook: ShownNotebook;
   readonly value: T;
 } | null;
 
 /**
- * Cells a run was requested for, with the execution each had at the time.
+ * A run this browser has asked the host for.
+ *
+ * The host answers some time after the request is made, and by then the
+ * notebook may have been replaced or the same cells asked for again. So
+ * the functions that handle the answer are given the request back, and
+ * they act only on it, in the notebook it was made in.
+ */
+export interface RunRequest {
+  readonly notebook: ShownNotebook;
+  readonly cells: readonly Protocol.CellId[];
+}
+
+/**
+ * Requests whose run has not appeared on the stream yet.
  *
- * The host shows a queued execution once the kernel plans the run. Until
- * then the cell reads as queued from this record. The host's newer
- * execution supersedes it; a rejected request, a cell the kernel will not
- * run, or the cell leaving the document withdraws it. Written only by
- * `submitCells`, `withdrawCells`, and `observeCellMessage`.
+ * They let a requested cell read as queued right away. Once the host
+ * accepts the run, the stream shows it with the same cells named on it,
+ * and the request is no longer needed. A request the host refuses, or
+ * that the browser gives up on, is withdrawn instead.
  */
-const submissionsAtom =
-  atom<Marks<ReadonlyMap<Protocol.CellId, Protocol.ExecutionId | null>>>(null);
+const requestsAtom = atom<Marks<ReadonlySet<RunRequest>>>(null);
 
 /**
  * Cells the kernel reports as blocked by a disabled ancestor. Written only
@@ -86,28 +97,58 @@ function current<T>(
     : empty;
 }
 
-function mark<T>(store: Store, value: T): Marks<T> {
+function shownNotebook(store: Store): ShownNotebook | null {
   const replica = store.get(activeNotebookReplicaAtom);
   const notebook = replica === null ? null : store.get(replica.notebookAtom);
   return replica === null || notebook === null
     ? null
-    : { notebook: { replica, id: notebook.id }, value };
+    : { replica, id: notebook.id };
 }
 
+function mark<T>(store: Store, value: T): Marks<T> {
+  const notebook = shownNotebook(store);
+  return notebook === null ? null : { notebook, value };
+}
+
+/**
+ * Whether `request` still belongs to the notebook being shown.
+ *
+ * Between a reset and the next snapshot the replica has no notebook for a
+ * moment. That says nothing about where the request belongs, so it does
+ * not count against it. Another replica, or another notebook on the same
+ * replica, does.
+ */
+export function isCurrent(store: Store, request: RunRequest): boolean {
+  const replica = store.get(activeNotebookReplicaAtom);
+  if (replica !== request.notebook.replica) {
+    return false;
+  }
+  const notebook = store.get(replica.notebookAtom);
+  return notebook === null || notebook.id === request.notebook.id;
+}
+
+/** Cells named by pending requests made in the notebook now shown. */
+const requestedCellsAtom = atom((get) => {
+  const replica = get(activeNotebookReplicaAtom);
+  const notebookId =
+    replica === null ? null : (get(replica.notebookAtom)?.id ?? null);
+  const requests = current(get(requestsAtom), replica, notebookId, NO_REQUESTS);
+  return requests.size === 0
+    ? NO_REQUESTED
+    : new Set([...requests].flatMap((request) => request.cells));
+});
+
 /** What every cell's view derives from. Pair with `cellExecutionView`. */
 export const cellExecutionInputsAtom = atom((get): CellExecutionInputs => {
   const replica = get(activeNotebookReplicaAtom);
   const notebookId =
     replica === null ? null : (get(replica.notebookAtom)?.id ?? null);
+  const document = replica === null ? null : get(replica.documentAtom);
   return {
     executions: replica === null ? NO_EXECUTIONS : get(replica.executionsAtom),
     results: replica === null ? NO_RESULTS : get(replica.resultsAtom),
-    submissions: current(
-      get(submissionsAtom),
-      replica,
-      notebookId,
-      NO_SUBMISSIONS,
-    ),
+    cells: document?.cells ?? NO_CELLS,
+    requested: get(requestedCellsAtom),
     blocked: current(get(blockedCellsAtom), replica, notebookId, NO_BLOCKED),
   };
 });
@@ -178,49 +219,126 @@ export const notebookQueuedOrRunningCountAtom = atom((get) =>
   queuedOrRunningCount(get(cellExecutionInputsAtom)),
 );
 
-/** Resolve once the host reports the runtime running. */
-export async function waitForRuntimeRunning(): Promise<void> {
-  await waitFor(runtimeStatusAtom, (status) => status === "running");
+type RuntimeReadiness = "running" | "waiting" | { refused: string };
+
+/**
+ * Whether a run can go to the host now, should wait, or has nothing to wait for.
+ *
+ * A runtime that is starting, or a stream that has not reported yet, will
+ * settle one way or the other. A failed or absent runtime, or a connection
+ * that has closed, will not, so a run waiting on it is refused instead.
+ */
+const runtimeReadinessAtom = atom((get): RuntimeReadiness => {
+  const replica = get(activeNotebookReplicaAtom);
+  if (replica === null) {
+    return { refused: "No notebook is connected" };
+  }
+  if (get(replica.connectionAtom).state === ConnectionState.CLOSED) {
+    return { refused: "The notebook is disconnected" };
+  }
+  // Until the stream's snapshot is complete there is no notebook to speak
+  // of, and nothing it reports about the runtime counts yet.
+  if (get(replica.notebookAtom) === null) {
+    return "waiting";
+  }
+  const runtime = get(replica.runtimeAtom);
+  if (runtime === null) {
+    return { refused: "The notebook has no runtime" };
+  }
+  switch (runtime.status) {
+    case "running":
+      return "running";
+    case "failed":
+      return { refused: runtime.message ?? "The runtime failed" };
+    default:
+      return "waiting";
+  }
+});
+
+/** Resolve once the host reports the runtime running; reject once it cannot. */
+export function waitForRuntimeRunning(store: Store): Promise<void> {
+  return new Promise((resolve, reject) => {
+    const unsubscribe = store.sub(runtimeReadinessAtom, check);
+    function check() {
+      const readiness = store.get(runtimeReadinessAtom);
+      if (readiness === "waiting") {
+        return;
+      }
+      unsubscribe();
+      if (readiness === "running") {
+        resolve();
+      } else {
+        reject(new Error(readiness.refused));
+      }
+    }
+    check();
+  });
 }
 
 /**
- * Record a run request for `cellIds` so they read as queued at once.
+ * Record that this browser is asking the host to run `cellIds`, so that
+ * they read as queued before the host has answered.
  *
- * Cells blocked by a disabled ancestor stay blocked: the kernel will not
- * queue them. A disabled cell is not queued either; the caller leaves it
- * out, since its configuration lives in the notebook state.
+ * A disabled cell, or one the kernel reports as blocked, may be among them
+ * so the kernel can set it aside, but the views do not show it as queued.
+ * Returns null while no notebook is shown, since there is nothing to
+ * record the request against.
  */
-export function submitCells(
+export function requestRun(
   store: Store,
   cellIds: readonly Protocol.CellId[],
-): void {
-  const inputs = store.get(cellExecutionInputsAtom);
-  const next = new Map(inputs.submissions);
-  for (const cellId of cellIds) {
-    if (!inputs.blocked.has(cellId)) {
-      next.set(cellId, latestExecutionId(inputs, cellId));
-    }
+): RunRequest | null {
+  const notebook = shownNotebook(store);
+  if (notebook === null) {
+    return null;
   }
-  setSubmissions(store, next);
+  const request: RunRequest = { notebook, cells: [...cellIds] };
+  const marks = store.get(requestsAtom);
+  const kept =
+    marks !== null &&
+    marks.notebook.replica === notebook.replica &&
+    marks.notebook.id === notebook.id
+      ? marks.value
+      : NO_REQUESTS;
+  store.set(requestsAtom, { notebook, value: new Set(kept).add(request) });
+  return request;
 }
 
-function setSubmissions(
-  store: Store,
-  value: ReadonlyMap<Protocol.CellId, Protocol.ExecutionId | null>,
-): void {
-  store.set(submissionsAtom, mark(store, value));
+/** Forget a request the host refused or the browser gave up on. */
+export function withdrawRequest(store: Store, request: RunRequest): void {
+  const marks = store.get(requestsAtom);
+  if (marks === null || !marks.value.has(request)) {
+    return;
+  }
+  const value = new Set(marks.value);
+  value.delete(request);
+  store.set(requestsAtom, { notebook: marks.notebook, value });
 }
 
-/** Forget a run request the host rejected. */
-export function withdrawCells(
+/**
+ * Keep a request until the run the host accepted for it appears on the
+ * stream, which can happen after the response has arrived. From then on
+ * the run itself says which cells are queued. If the notebook is no
+ * longer shown by then, the request is simply dropped.
+ */
+export function acceptRun(
   store: Store,
-  cellIds: readonly Protocol.CellId[],
+  request: RunRequest,
+  run: Protocol.Run,
 ): void {
-  const next = new Map(store.get(cellExecutionInputsAtom).submissions);
-  for (const cellId of cellIds) {
-    next.delete(cellId);
+  const { replica } = request.notebook;
+  const unsubscribe = store.sub(replica.executionsAtom, check);
+  function check() {
+    if (
+      isCurrent(store, request) &&
+      !store.get(replica.executionsAtom).runs.has(run.id)
+    ) {
+      return;
+    }
+    unsubscribe();
+    withdrawRequest(store, request);
   }
-  setSubmissions(store, next);
+  check();
 }
 
 function setCellBlocked(
@@ -242,22 +360,19 @@ function setCellBlocked(
 }
 
 /**
- * Keep the blocked set and the submissions in step with a cell-op.
- *
- * A status names whether the cell is blocked. A blocked cell and one whose
- * output is an error will not run, so a submission that never became an
- * execution is withdrawn, as it was when the editor set status itself.
+ * Learn from a cell-op whether the kernel considers the cell blocked by a
+ * disabled ancestor. A blocked cell will not run, so a run that names it
+ * does not show it as queued.
  *
  * TODO(MO-8407): Delete with the cell-op case. The blocked read goes with
- * MO-8463, the error read with MO-8346.
+ * MO-8463.
  */
 export function observeCellMessage(store: Store, message: CellMessage): void {
-  const cellId = message.cell_id;
-  const blocked = message.status === "disabled-transitively";
   if (message.status != null) {
-    setCellBlocked(store, cellId, blocked);
-  }
-  if (blocked || message.output?.mimetype === "application/vnd.marimo+error") {
-    withdrawCells(store, [cellId]);
+    setCellBlocked(
+      store,
+      message.cell_id,
+      message.status === "disabled-transitively",
+    );
   }
 }

package/src/core/executions/index.ts

@@ -1,11 +1,12 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 /**
- * Cell status and timing read from the host's executions.
+ * Cell status, timing, and output, as the editor reads them.
  *
- * Every atom here is derived, so it can be read but not set. The state
- * behind them changes only through `submitCells`, `withdrawCells`, and
- * `observeCellMessage`.
+ * The host's notebook stream is the source of truth. Its runs, executions,
+ * and results say what every cell is doing. The atoms here derive views
+ * from that stream, adding only what the host cannot know yet, such as a
+ * run this browser has asked for but not heard back about.
  */
 
 export {
@@ -17,11 +18,11 @@ export {
   notebookQueuedOrRunningCountAtom,
   observeCellMessage,
   runtimeStatusAtom,
-  submitCells,
   waitForRuntimeRunning,
-  withdrawCells,
 } from "./atoms";
 export type { CellFreshness, CellFreshnessLookup } from "./freshness";
+export { replyInput } from "./input";
+export { runCells, syncEdits, waitForRuntime } from "./run";
 export {
   useCellExecution,
   useCellFreshness,

package/src/core/executions/input.ts

@@ -0,0 +1,55 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** Answer a prompt the runtime is waiting on. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import type { createStore } from "jotai";
+import { toast } from "@/components/ui/use-toast";
+import { activeNotebookReplicaAtom } from "@/core/session/active-notebook";
+import { prettyError } from "@/utils/errors";
+import { Logger } from "@/utils/Logger";
+
+type Store = ReturnType<typeof createStore>;
+
+/**
+ * Send `value` to the host as the answer to the prompt `inputId`.
+ *
+ * The reply names its prompt, so the host marks that prompt answered for
+ * every client and refuses a reply to one that has already been answered
+ * or cancelled. A refusal is reported, since the user typed the answer and
+ * would otherwise wait on a prompt that is not going to take it.
+ */
+export async function replyInput(
+  store: Store,
+  inputId: Protocol.ItemId,
+  value: string,
+): Promise<void> {
+  const replica = store.get(activeNotebookReplicaAtom);
+  const notebook = replica === null ? null : store.get(replica.notebookAtom);
+  const backend = replica === null ? null : store.get(replica.backendAtom);
+  if (notebook === null || backend === null) {
+    Logger.warn("No notebook is connected to answer the prompt in");
+    return;
+  }
+  try {
+    const result = await backend
+      .notebook(notebook.id)
+      .input(inputId, { value });
+    if (result.ok) {
+      return;
+    }
+    Logger.error("The host refused the input reply", result.problem);
+    toast({
+      title: "Failed to send input",
+      description: result.problem.detail ?? result.problem.title,
+      variant: "danger",
+    });
+  } catch (error) {
+    Logger.error("Failed to send input", error);
+    toast({
+      title: "Failed to send input",
+      description: prettyError(error),
+      variant: "danger",
+    });
+  }
+}

package/src/core/executions/run.ts

@@ -0,0 +1,120 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** Submit cell runs to the host. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import type { createStore } from "jotai";
+import type { DocumentModel } from "@/core/notebook/document";
+import { toast } from "@/components/ui/use-toast";
+import { startRuntimeIfLazy } from "@/core/runtime/config";
+import { activeDocumentModelAtom } from "@/core/session/active-notebook";
+import { prettyError } from "@/utils/errors";
+import { Logger } from "@/utils/Logger";
+import {
+  acceptRun,
+  isCurrent,
+  requestRun,
+  waitForRuntimeRunning,
+  withdrawRequest,
+} from "./atoms";
+
+type Store = ReturnType<typeof createStore>;
+
+/**
+ * Start a lazy runtime and wait until the host reports it running.
+ *
+ * The host refuses a run while its runtime is not running, so a request
+ * made during startup waits here instead of failing. It rejects once the
+ * runtime cannot come up, so the request is not left waiting.
+ */
+export async function waitForRuntime(store: Store): Promise<void> {
+  await startRuntimeIfLazy();
+  await waitForRuntimeRunning(store);
+}
+
+/**
+ * Get this browser's edits to the host before a command that names cells.
+ *
+ * A run asks for the edits explicitly, so a synchronization that stalled
+ * earlier is retried here rather than left to the footer's retry.
+ */
+export function syncEdits(model: DocumentModel): Promise<Protocol.Document> {
+  return model.retry();
+}
+
+/**
+ * Run `cellIds` through the host once this browser's edits have reached it.
+ *
+ * The cells read as queued right away. When the host accepts the run, the
+ * stream shows it with the same cells named on it and takes over from the
+ * request. A request that never reaches the host is withdrawn, and the
+ * user is told why. The request belongs to the notebook it was made in, so
+ * if another notebook is shown by the time the edits have synced, it is
+ * dropped rather than sent there.
+ */
+export async function runCells(
+  store: Store,
+  cellIds: readonly Protocol.CellId[],
+): Promise<void> {
+  if (cellIds.length === 0) {
+    return;
+  }
+  const model = store.get(activeDocumentModelAtom);
+  if (model === null) {
+    Logger.warn("No notebook is open to run cells in");
+    return;
+  }
+  // The stream may have no notebook for a moment, as during a reconnect.
+  // The request is then recorded once the runtime, and with it the
+  // notebook, has been seen.
+  let request = requestRun(store, cellIds);
+  try {
+    await waitForRuntime(store);
+    request ??= requestRun(store, cellIds);
+    if (request === null) {
+      throw new Error("No notebook is shown");
+    }
+    const { notebook_id } = await syncEdits(model);
+    if (
+      store.get(activeDocumentModelAtom) !== model ||
+      !isCurrent(store, request) ||
+      notebook_id !== request.notebook.id
+    ) {
+      withdrawRequest(store, request);
+      return;
+    }
+    const backend = store.get(request.notebook.replica.backendAtom);
+    if (backend === null) {
+      throw new Error("The notebook is not connected");
+    }
+    const result = await backend
+      .notebook(notebook_id)
+      .run({ cells: cellIds.map((id) => ({ id })) });
+    if (result.ok) {
+      acceptRun(store, request, result.data);
+      return;
+    }
+    withdrawRequest(store, request);
+    // The connection is read-only for this action, which the UI shows.
+    if (result.status === 403) {
+      Logger.warn("Run refused, connection is read-only", result.problem);
+      return;
+    }
+    Logger.error("The host refused the run", result.problem);
+    toast({
+      title: "Failed to run",
+      description: result.problem.detail ?? result.problem.title,
+      variant: "danger",
+    });
+  } catch (error) {
+    if (request !== null) {
+      withdrawRequest(store, request);
+    }
+    Logger.error("Failed to run cells", error);
+    toast({
+      title: "Failed to run",
+      description: prettyError(error),
+      variant: "danger",
+    });
+  }
+}

package/src/core/executions/select.ts

@@ -31,14 +31,15 @@ export interface CellExecutionInputs {
   readonly executions: Executions;
   readonly results: Results;
   /**
-   * The execution each submitted cell had when its run was requested.
-   *
-   * A cell reads as queued until the host shows a newer execution for it.
+   * The cells of the confirmed document. A run can only reach a cell the
+   * document has, and the kernel sets aside one the document has disabled.
    */
-  readonly submissions: ReadonlyMap<
-    Protocol.CellId,
-    Protocol.ExecutionId | null
-  >;
+  readonly cells: ReadonlyMap<Protocol.CellId, Pick<Protocol.Cell, "config">>;
+  /**
+   * Cells this browser has asked to run, while the stream has not yet
+   * shown the run.
+   */
+  readonly requested: ReadonlySet<Protocol.CellId>;
   /** Cells the kernel reports as blocked by a disabled ancestor. */
   readonly blocked: ReadonlySet<Protocol.CellId>;
 }
@@ -66,14 +67,46 @@ function latestOf(
 }
 
 /**
- * The execution a submission is measured against: the cell's latest on the
- * stream, or its retained result once the stream has let that execution go.
+ * Whether a run can reach `id`, meaning the document has the cell and
+ * nothing keeps it from running.
  */
-export function latestExecutionId(
-  inputs: Pick<CellExecutionInputs, "executions" | "results">,
-  id: Protocol.CellId,
-): Protocol.ExecutionId | null {
-  return inputs.executions.latest.get(id) ?? inputs.results.get(id)?.id ?? null;
+function runnable(inputs: CellExecutionInputs, id: Protocol.CellId): boolean {
+  const cell = inputs.cells.get(id);
+  return cell !== undefined && !cell.config.disabled && !inputs.blocked.has(id);
+}
+
+/**
+ * Whether `id` is waiting on a run that has not planned an execution for
+ * it yet.
+ *
+ * When the host accepts a run it names the cells that were requested, and
+ * this happens before the kernel plans their executions. Until the cell's
+ * execution in that run appears, the run itself is what makes the cell
+ * queued. A request this browser has made works the same way until its
+ * run shows up on the stream. The host keeps one execution per cell, so
+ * the cell's latest execution, or its retained result once the stream has
+ * let that execution go, tells which run has already planned it.
+ */
+function awaited(inputs: CellExecutionInputs, id: Protocol.CellId): boolean {
+  if (!runnable(inputs, id)) {
+    return false;
+  }
+  if (inputs.requested.has(id)) {
+    return true;
+  }
+  const planned =
+    latestOf(inputs.executions, id)?.record.run_id ??
+    inputs.results.get(id)?.run_id;
+  for (const run of inputs.executions.runs.values()) {
+    if (
+      run.status !== "completed" &&
+      run.id !== planned &&
+      run.cells.includes(id)
+    ) {
+      return true;
+    }
+  }
+  return false;
 }
 
 export function cellExecutionView(
@@ -82,9 +115,6 @@ export function cellExecutionView(
 ): CellExecutionView {
   const latest = latestOf(inputs.executions, id);
   const result = inputs.results.get(id) ?? null;
-  const submitted =
-    inputs.submissions.has(id) &&
-    inputs.submissions.get(id) === latestExecutionId(inputs, id);
   let status: CellRunStatus = "idle";
   let startedAt = result === null ? null : millis(result.started_at);
   let rendered = false;
@@ -92,7 +122,7 @@ export function cellExecutionView(
     status = latest.record.status;
     startedAt = millis(latest.record.started_at) ?? startedAt;
     rendered = latest.record.outcome !== null;
-  } else if (submitted) {
+  } else if (awaited(inputs, id)) {
     status = "queued";
   }
   const settled = status === "idle";
@@ -126,7 +156,7 @@ export function anyRunning(executions: Executions): boolean {
   );
 }
 
-/** Cells and scratch work queued or running, plus submissions the host has not shown yet. */
+/** Cells and scratch work queued or running, plus cells awaiting a run the kernel has not planned. */
 export function queuedOrRunningCount(inputs: CellExecutionInputs): number {
   const pending = unfinished(inputs.executions);
   const shown = new Set(
@@ -134,12 +164,17 @@ export function queuedOrRunningCount(inputs: CellExecutionInputs): number {
       execution.record.cell_id === null ? [] : [execution.record.cell_id],
     ),
   );
+  const named = new Set<Protocol.CellId>(inputs.requested);
+  for (const run of inputs.executions.runs.values()) {
+    if (run.status !== "completed") {
+      for (const cellId of run.cells) {
+        named.add(cellId);
+      }
+    }
+  }
   let awaiting = 0;
-  for (const [cellId, executionId] of inputs.submissions) {
-    if (
-      !shown.has(cellId) &&
-      executionId === latestExecutionId(inputs, cellId)
-    ) {
+  for (const cellId of named) {
+    if (!shown.has(cellId) && awaited(inputs, cellId)) {
       awaiting += 1;
     }
   }