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

package/src/core/executions/__tests__/run.test.ts

@@ -0,0 +1,390 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import { atom } from "jotai";
+import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
+import { MockExecutions } from "@/__mocks__/executions";
+import { cellId } from "@/__tests__/branded";
+import type { DocumentModel, EditableDocument } from "@/core/notebook/document";
+import {
+  cell,
+  CONFIG,
+  document as documentOf,
+  withCells,
+} from "@/core/notebook/document/__tests__/fixtures";
+import {
+  ConnectionClosedReason,
+  ConnectionState,
+} from "@/core/notebook/replica";
+import {
+  type FakeRun,
+  fakeBackend,
+} from "@/core/notebook/replica/__tests__/fixtures";
+import { activeDocumentModelAtom } from "@/core/session/active-notebook";
+import { store } from "@/core/state/jotai";
+import { cellExecutionInputsAtom } from "../atoms";
+import { runCells } from "../run";
+import { cellExecutionView } from "../select";
+
+const toast = vi.fn();
+vi.mock("@/components/ui/use-toast", () => ({
+  toast: (...args: unknown[]) => toast(...args),
+}));
+vi.mock("@/core/runtime/config", () => ({
+  startRuntimeIfLazy: async () => undefined,
+}));
+
+const C1 = cellId("c1");
+const C2 = cellId("c2");
+
+function accepted(cells: Protocol.CellId[]): Protocol.Run {
+  return {
+    id: "run-1",
+    notebook_id: "n1",
+    runtime_id: "r1",
+    cells,
+    status: "queued",
+    outcome: null,
+    started_at: null,
+    completed_at: null,
+  };
+}
+
+function editable(disabled: Protocol.CellId[] = []): EditableDocument {
+  const confirmed = documentOf();
+  return {
+    order: confirmed.revision.order,
+    cells: new Map(
+      [...confirmed.cells.values()].map((cell) => [
+        cell.id,
+        {
+          ...cell,
+          config: { ...cell.config, disabled: disabled.includes(cell.id) },
+        },
+      ]),
+    ),
+  };
+}
+
+/** A model whose edits sync with `sync`, confirming the fixture document by default. */
+function model(
+  sync: () => Promise<Protocol.Document> = async () => documentOf().revision,
+  document: EditableDocument | null = editable(),
+): DocumentModel {
+  const model = {
+    retry: vi.fn(sync),
+    documentAtom: atom(document),
+  } as unknown as DocumentModel;
+  store.set(activeDocumentModelAtom, model);
+  return model;
+}
+
+function status(id: Protocol.CellId) {
+  return cellExecutionView(store.get(cellExecutionInputsAtom), id).status;
+}
+
+describe("runCells", () => {
+  let run: ReturnType<typeof vi.fn<FakeRun>>;
+
+  beforeEach(() => {
+    toast.mockClear();
+    run = vi.fn<FakeRun>(async (_notebookId, request) => ({
+      ok: true,
+      status: 202,
+      data: accepted(request.cells?.map((cell) => cell.id) ?? []),
+    }));
+    const replica = MockExecutions.seed(
+      store,
+      [],
+      MockExecutions.runtime("running"),
+      documentOf(),
+    );
+    store.set(replica.backendAtom, fakeBackend({ run }));
+  });
+
+  afterEach(() => {
+    store.set(activeDocumentModelAtom, null);
+  });
+
+  it("queues the cells at once, then runs them by id after the edits sync", async () => {
+    const { retry } = model();
+
+    const pending = runCells(store, [C1, C2]);
+    expect(status(C1)).toBe("queued");
+    expect(status(C2)).toBe("queued");
+    await pending;
+
+    // Retrying covers a synchronization that stalled before the run.
+    expect(retry).toHaveBeenCalledOnce();
+    expect(run).toHaveBeenCalledWith("n1", {
+      cells: [{ id: C1 }, { id: C2 }],
+    });
+    expect(vi.mocked(retry).mock.invocationCallOrder[0]).toBeLessThan(
+      run.mock.invocationCallOrder[0],
+    );
+    // Queued until the host plans the run or ends it.
+    expect(status(C1)).toBe("queued");
+    expect(toast).not.toHaveBeenCalled();
+  });
+
+  it("sends a disabled cell without reading it as queued", async () => {
+    const replica = MockExecutions.seed(
+      store,
+      [],
+      MockExecutions.runtime("running"),
+      withCells(documentOf(), {
+        c2: cell("c2", "y = 2", { config: { ...CONFIG, disabled: true } }),
+      }),
+    );
+    store.set(replica.backendAtom, fakeBackend({ run }));
+    model(undefined, editable([C2]));
+
+    await runCells(store, [C1, C2]);
+
+    expect(run).toHaveBeenCalledWith("n1", {
+      cells: [{ id: C1 }, { id: C2 }],
+    });
+    expect(status(C1)).toBe("queued");
+    expect(status(C2)).toBe("idle");
+  });
+
+  it("waits for the runtime to run before submitting", async () => {
+    const replica = MockExecutions.seed(
+      store,
+      [],
+      MockExecutions.runtime("starting"),
+      documentOf(),
+    );
+    store.set(replica.backendAtom, fakeBackend({ run }));
+    model();
+
+    const pending = runCells(store, [C1]);
+    await Promise.resolve();
+    expect(status(C1)).toBe("queued");
+    expect(run).not.toHaveBeenCalled();
+
+    store.set(replica.runtimeAtom, MockExecutions.runtime("running"));
+    await pending;
+    expect(run).toHaveBeenCalledOnce();
+  });
+
+  it("records the request once the stream has a notebook, as after a reconnect", async () => {
+    const replica = MockExecutions.seed(store, [], null, documentOf());
+    store.set(replica.backendAtom, fakeBackend({ run }));
+    store.set(replica.notebookAtom, null);
+    model();
+
+    const pending = runCells(store, [C1]);
+    await Promise.resolve();
+    expect(run).not.toHaveBeenCalled();
+
+    // The real replica reports both at once, from one snapshot.
+    store.set(replica.runtimeAtom, MockExecutions.runtime("running"));
+    store.set(replica.notebookAtom, MockExecutions.notebook());
+    await pending;
+
+    expect(run).toHaveBeenCalledWith("n1", { cells: [{ id: C1 }] });
+    expect(status(C1)).toBe("queued");
+    expect(toast).not.toHaveBeenCalled();
+  });
+
+  it("withdraws the submission and reports when the edits cannot sync", async () => {
+    model(async () => {
+      throw new Error("offline");
+    });
+
+    await runCells(store, [C1]);
+
+    expect(run).not.toHaveBeenCalled();
+    expect(status(C1)).toBe("idle");
+    expect(toast).toHaveBeenCalledWith(
+      expect.objectContaining({ title: "Failed to run", variant: "danger" }),
+    );
+  });
+
+  it("withdraws the submission and reports when the runtime has failed", async () => {
+    const replica = MockExecutions.seed(
+      store,
+      [],
+      { ...MockExecutions.runtime("failed"), message: "The kernel exited" },
+      documentOf(),
+    );
+    store.set(replica.backendAtom, fakeBackend({ run }));
+    model();
+
+    await runCells(store, [C1]);
+
+    expect(run).not.toHaveBeenCalled();
+    expect(status(C1)).toBe("idle");
+    expect(toast).toHaveBeenCalledWith(
+      expect.objectContaining({ description: "The kernel exited" }),
+    );
+  });
+
+  it("withdraws the submission and reports when the notebook has no runtime", async () => {
+    const replica = MockExecutions.seed(store, [], null, documentOf());
+    store.set(replica.backendAtom, fakeBackend({ run }));
+    model();
+
+    await runCells(store, [C1]);
+
+    expect(run).not.toHaveBeenCalled();
+    expect(status(C1)).toBe("idle");
+    expect(toast).toHaveBeenCalledWith(
+      expect.objectContaining({ description: "The notebook has no runtime" }),
+    );
+  });
+
+  it("stops waiting for the runtime once the connection closes", async () => {
+    const replica = MockExecutions.seed(
+      store,
+      [],
+      MockExecutions.runtime("starting"),
+      documentOf(),
+    );
+    store.set(replica.backendAtom, fakeBackend({ run }));
+    model();
+
+    const pending = runCells(store, [C1]);
+    await Promise.resolve();
+    expect(status(C1)).toBe("queued");
+
+    store.set(replica.connectionAtom, {
+      state: ConnectionState.CLOSED,
+      code: ConnectionClosedReason.KERNEL_DISCONNECTED,
+      reason: "kernel not found",
+    });
+    await pending;
+
+    expect(run).not.toHaveBeenCalled();
+    expect(status(C1)).toBe("idle");
+    expect(toast).toHaveBeenCalledWith(
+      expect.objectContaining({ description: "The notebook is disconnected" }),
+    );
+  });
+
+  it("withdraws the submission and reports when the host refuses the run", async () => {
+    run.mockResolvedValue({
+      ok: false,
+      status: 409,
+      problem: {
+        type: "tag:marimo.io,2026:host/conflict",
+        title: "Conflict",
+        status: 409,
+        detail: "the runtime is starting.",
+      },
+    });
+    model();
+
+    await runCells(store, [C1]);
+
+    expect(status(C1)).toBe("idle");
+    expect(toast).toHaveBeenCalledWith(
+      expect.objectContaining({ description: "the runtime is starting." }),
+    );
+  });
+
+  it("withdraws quietly when the connection is read-only", async () => {
+    run.mockResolvedValue({
+      ok: false,
+      status: 403,
+      problem: {
+        type: "tag:marimo.io,2026:host/forbidden",
+        title: "Forbidden",
+        status: 403,
+      },
+    });
+    model();
+
+    await runCells(store, [C1]);
+
+    expect(status(C1)).toBe("idle");
+    expect(toast).not.toHaveBeenCalled();
+  });
+
+  it("does not run when the notebook closed while the edits synced", async () => {
+    model(async () => {
+      store.set(activeDocumentModelAtom, null);
+      return documentOf().revision;
+    });
+
+    await runCells(store, [C1]);
+
+    expect(run).not.toHaveBeenCalled();
+    expect(toast).not.toHaveBeenCalled();
+  });
+
+  it("does not run in a notebook that replaced the one the request was made in", async () => {
+    // Islands keep the replica and model while the worker's session changes.
+    const replica = MockExecutions.seed(
+      store,
+      [],
+      MockExecutions.runtime("running"),
+      documentOf(),
+    );
+    store.set(replica.backendAtom, fakeBackend({ run }));
+    model(async () => {
+      store.set(replica.notebookAtom, MockExecutions.notebook("n2"));
+      return documentOf("n2").revision;
+    });
+
+    const pending = runCells(store, [C1]);
+    expect(status(C1)).toBe("queued");
+    await pending;
+
+    expect(run).not.toHaveBeenCalled();
+    expect(status(C1)).toBe("idle");
+    expect(toast).not.toHaveBeenCalled();
+  });
+
+  it("leaves the replacing notebook's request when a late refusal lands", async () => {
+    const replica = MockExecutions.seed(
+      store,
+      [],
+      MockExecutions.runtime("running"),
+      documentOf(),
+    );
+    store.set(replica.backendAtom, fakeBackend({ run }));
+    model(
+      async () =>
+        documentOf(store.get(replica.notebookAtom)?.id ?? "n1").revision,
+    );
+    let refuse = (): void => undefined;
+    run.mockImplementationOnce(
+      () =>
+        new Promise((resolve) => {
+          refuse = () =>
+            resolve({
+              ok: false,
+              status: 409,
+              problem: {
+                type: "tag:marimo.io,2026:host/conflict",
+                title: "Conflict",
+                status: 409,
+              },
+            });
+        }),
+    );
+
+    const first = runCells(store, [C1]);
+    await vi.waitFor(() => expect(run).toHaveBeenCalledOnce());
+    store.set(replica.notebookAtom, MockExecutions.notebook("n2"));
+    const second = runCells(store, [C1]);
+    await vi.waitFor(() => expect(run).toHaveBeenCalledTimes(2));
+    expect(status(C1)).toBe("queued");
+
+    refuse();
+    await first;
+    await second;
+
+    expect(run).toHaveBeenLastCalledWith("n2", { cells: [{ id: C1 }] });
+    expect(status(C1)).toBe("queued");
+  });
+
+  it("does nothing without an open notebook", async () => {
+    await runCells(store, [C1]);
+
+    expect(run).not.toHaveBeenCalled();
+    expect(status(C1)).toBe("idle");
+  });
+});

package/src/core/executions/__tests__/select.test.ts

@@ -2,6 +2,7 @@
 
 import type * as Protocol from "@marimo-team/client/protocol";
 import { describe, expect, it } from "vitest";
+import { CONFIG } from "@/core/notebook/document/__tests__/fixtures";
 import type { Execution, Executions, Results } from "@/core/notebook/replica";
 import {
   anyRunning,
@@ -39,19 +40,43 @@ function execution(overrides: Partial<Protocol.Execution> & { id: string }) {
   return { record: record(overrides), items: new Map(), variables: new Map() };
 }
 
+function run(
+  cells: string[],
+  status: Protocol.Run["status"],
+  id = "run-1",
+): Protocol.Run {
+  return {
+    id,
+    notebook_id: "n1",
+    runtime_id: "r1",
+    cells,
+    status,
+    outcome: status === "completed" ? "succeeded" : null,
+    started_at: null,
+    completed_at: null,
+  };
+}
+
+const CELLS = ["c1", "c2", "c3", "c4", "c5"];
+
+/** Inputs whose document has `CELLS`, with `disabled` kept from running. */
 function inputs({
   executions = [],
+  runs = [],
   results = new Map(),
-  submissions = new Map(),
+  requested = new Set(),
+  disabled = [],
   blocked = new Set(),
 }: {
   executions?: Execution[];
+  runs?: Protocol.Run[];
   results?: Results;
-  submissions?: ReadonlyMap<string, string | null>;
+  requested?: ReadonlySet<string>;
+  disabled?: readonly string[];
   blocked?: ReadonlySet<string>;
 }): CellExecutionInputs {
   const stream: Executions = {
-    runs: new Map(),
+    runs: new Map(runs.map((r) => [r.id, r])),
     executions: new Map(executions.map((e) => [e.record.id, e])),
     latest: new Map(
       executions.flatMap((e) =>
@@ -59,7 +84,13 @@ function inputs({
       ),
     ),
   };
-  return { executions: stream, results, submissions, blocked };
+  const cells = new Map(
+    CELLS.map((id) => [
+      id,
+      { config: { ...CONFIG, disabled: disabled.includes(id) } },
+    ]),
+  );
+  return { executions: stream, results, cells, requested, blocked };
 }
 
 describe("cellExecutionView", () => {
@@ -165,40 +196,87 @@ describe("cellExecutionView", () => {
     expect(ended({ kind: "interrupted" }).outcome).toBe("interrupted");
   });
 
-  it("reads as queued from a submission until the host shows a newer execution", () => {
-    const before = inputs({
+  it("reads as queued from a request the stream has not shown a run for", () => {
+    const requested = inputs({
       executions: [execution({ id: "e1" })],
       results: new Map([["c1", record({ id: "e1" })]]),
-      submissions: new Map([["c1", "e1"]]),
+      requested: new Set(["c1"]),
     });
-    expect(cellExecutionView(before, "c1").status).toBe("queued");
-    expect(cellExecutionView(before, "c1").outcome).toBeNull();
+    expect(cellExecutionView(requested, "c1").status).toBe("queued");
+    expect(cellExecutionView(requested, "c1").outcome).toBeNull();
 
-    // The stream let e1 go while the request was in flight; its result remains.
-    const forgotten = inputs({
-      results: new Map([["c1", record({ id: "e1" })]]),
-      submissions: new Map([["c1", "e1"]]),
+    const fresh = inputs({ requested: new Set(["c1"]) });
+    expect(cellExecutionView(fresh, "c1").status).toBe("queued");
+  });
+
+  it("reads as queued from an unfinished run that names the cell until it is planned", () => {
+    const result = new Map([["c1", record({ id: "e1" })]]);
+    const accepted = inputs({
+      runs: [run(["c1"], "queued", "run-2")],
+      executions: [execution({ id: "e1" })],
+      results: result,
     });
-    expect(cellExecutionView(forgotten, "c1").status).toBe("queued");
+    expect(cellExecutionView(accepted, "c1").status).toBe("queued");
+    expect(cellExecutionView(accepted, "c1").outcome).toBeNull();
 
-    const shown = inputs({
+    // The cell's execution in the run takes over once the kernel plans it.
+    const planned = inputs({
+      runs: [run(["c1"], "running", "run-2")],
       executions: [
-        execution({ id: "e2", status: "running", completed_at: null }),
+        execution({
+          id: "e2",
+          run_id: "run-2",
+          status: "running",
+          completed_at: null,
+        }),
       ],
-      results: new Map([["c1", record({ id: "e1" })]]),
-      submissions: new Map([["c1", "e1"]]),
+      results: result,
     });
-    expect(cellExecutionView(shown, "c1").status).toBe("running");
+    expect(cellExecutionView(planned, "c1").status).toBe("running");
 
     const done = inputs({
-      executions: [execution({ id: "e2" })],
-      results: new Map([["c1", record({ id: "e2" })]]),
-      submissions: new Map([["c1", "e1"]]),
+      runs: [run(["c1", "c2"], "running", "run-2")],
+      executions: [execution({ id: "e2", run_id: "run-2" })],
+      results: new Map([["c1", record({ id: "e2", run_id: "run-2" })]]),
     });
     expect(cellExecutionView(done, "c1").status).toBe("idle");
+    expect(cellExecutionView(done, "c1").outcome).toBe("value");
+    expect(cellExecutionView(done, "c2").status).toBe("queued");
 
-    const fresh = inputs({ submissions: new Map([["c1", null]]) });
-    expect(cellExecutionView(fresh, "c1").status).toBe("queued");
+    // The stream let the execution go while its run goes on with other cells.
+    const retained = inputs({
+      runs: [run(["c1", "c2"], "running", "run-2")],
+      results: new Map([["c1", record({ id: "e2", run_id: "run-2" })]]),
+    });
+    expect(cellExecutionView(retained, "c1").status).toBe("idle");
+    expect(cellExecutionView(retained, "c2").status).toBe("queued");
+
+    // A later run names the cell again while the earlier one's result stands.
+    const again = inputs({
+      runs: [run(["c1"], "running", "run-2"), run(["c1"], "queued", "run-3")],
+      executions: [execution({ id: "e2", run_id: "run-2" })],
+      results: new Map([["c1", record({ id: "e2", run_id: "run-2" })]]),
+    });
+    expect(cellExecutionView(again, "c1").status).toBe("queued");
+
+    const over = inputs({ runs: [run(["c1"], "completed")] });
+    expect(cellExecutionView(over, "c1").status).toBe("idle");
+  });
+
+  it("does not read as queued a named cell a run cannot reach", () => {
+    const runs = [run(["c1", "c9"], "queued")];
+    expect(
+      cellExecutionView(inputs({ runs, disabled: ["c1"] }), "c1").status,
+    ).toBe("idle");
+    expect(
+      cellExecutionView(inputs({ runs, blocked: new Set(["c1"]) }), "c1")
+        .status,
+    ).toBe("idle");
+    // The document no longer has c9.
+    expect(cellExecutionView(inputs({ runs }), "c9").status).toBe("idle");
+    expect(
+      cellExecutionView(inputs({ requested: new Set(["c9"]) }), "c9").status,
+    ).toBe("idle");
   });
 
   it("shows a blocked cell as disabled transitively, with its last outcome", () => {
@@ -234,15 +312,12 @@ describe("notebook-level views", () => {
         completed_at: null,
       }),
     ],
+    runs: [run(["c1", "c2", "c5"], "running"), run(["c5"], "queued", "run-2")],
     results: new Map([["c5", record({ id: "e5", cell_id: "c5" })]]),
-    submissions: new Map([
-      ["c2", "e2"],
-      ["c4", null],
-      ["c5", "e5"],
-    ]),
+    requested: new Set(["c2", "c4"]),
   });
 
-  it("counts unfinished executions and submissions the host has not shown", () => {
+  it("counts unfinished executions and cells awaiting a run the kernel has not planned", () => {
     expect(anyRunning(stream.executions)).toBe(true);
     expect(queuedOrRunningCount(stream)).toBe(5);
     expect(runningCellIds(stream.executions)).toEqual(["c1"]);