@marimo-team/islands 0.25.3-dev28 → 0.25.3-dev284

package/src/core/executions/hooks.ts

@@ -0,0 +1,47 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import { dequal } from "dequal";
+import { useAtomValue } from "jotai";
+import { selectAtom } from "jotai/utils";
+import { useMemo } from "react";
+import { cellExecutionInputsAtom, cellFreshnessAtom } from "./atoms";
+import type { CellFreshness } from "./freshness";
+import { type CellExecutionView, cellExecutionView } from "./select";
+
+/**
+ * A cell's status, timing, and how it last ended.
+ *
+ * Read from the host's executions, with the editor's own marks for a run
+ * it has requested and the kernel has not planned yet, and for a cell the
+ * kernel reports blocked by a disabled ancestor.
+ */
+export function useCellExecution(id: Protocol.CellId): CellExecutionView {
+  // One atom per mounted cell, kept for its subscription and compared by
+  // value so the cell renders only when its own view changes.
+  const cellAtom = useMemo(
+    () =>
+      selectAtom(
+        cellExecutionInputsAtom,
+        (inputs) => cellExecutionView(inputs, id),
+        dequal,
+      ),
+    [id],
+  );
+  return useAtomValue(cellAtom);
+}
+
+/**
+ * Whether a cell's result is current, by the host protocol's rules over
+ * the confirmed document and the executions mirror.
+ *
+ * The confirmed document is what the host has accepted, so an edit reads
+ * as such once the model has flushed it, not on the keystroke.
+ */
+export function useCellFreshness(id: Protocol.CellId): CellFreshness {
+  const cellAtom = useMemo(
+    () => selectAtom(cellFreshnessAtom, (freshness) => freshness(id), dequal),
+    [id],
+  );
+  return useAtomValue(cellAtom);
+}

package/src/core/executions/index.ts

@@ -0,0 +1,33 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/**
+ * Cell status and timing read from the host's executions.
+ *
+ * Every atom here is derived, so it can be read but not set. The state
+ * behind them changes only through `submitCells`, `withdrawCells`, and
+ * `observeCellMessage`.
+ */
+
+export {
+  cellExecutionInputsAtom,
+  cellFreshnessAtom,
+  notebookIsRunningAtom,
+  notebookQueuedOrRunningCountAtom,
+  observeCellMessage,
+  runtimeStatusAtom,
+  submitCells,
+  waitForRuntimeRunning,
+  withdrawCells,
+} from "./atoms";
+export type { CellFreshness, CellFreshnessLookup } from "./freshness";
+export { useCellExecution, useCellFreshness } from "./hooks";
+export {
+  type CellExecutionInputs,
+  type CellExecutionView,
+  type CellOutcome,
+  type CellRunStatus,
+  cellExecutionView,
+  onlyScratchRunning,
+  runningCellIds,
+  scratchStatus,
+} from "./select";

package/src/core/executions/select.ts

@@ -0,0 +1,175 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** Pure views over the replica's executions and results, by cell. */
+
+import type * as Protocol from "@marimo-team/client/protocol";
+import type { Execution, Executions, Results } from "@/core/notebook/replica";
+import { type Milliseconds, Time } from "@/utils/time";
+
+/** Where a cell is in its run. */
+export type CellRunStatus = "idle" | "queued" | "running";
+
+/** How a cell's last completed execution ended. */
+export type CellOutcome = "value" | "error" | "interrupted" | "skipped";
+
+export interface CellExecutionView {
+  readonly status: CellRunStatus;
+  /** Whether a cell it depends on is disabled, so it cannot run. */
+  readonly blocked: boolean;
+  /** Null until the cell completes once, and while it runs again. */
+  readonly outcome: CellOutcome | null;
+  /** When the current run started, or the last one, since the epoch. */
+  readonly startedAt: Milliseconds | null;
+  /** How long the last completed execution took, null while the cell runs again. */
+  readonly duration: Milliseconds | null;
+  /** Whether the running execution has shown a value so far. TODO(MO-8404) */
+  readonly rendered: boolean;
+}
+
+/** What a cell's view is derived from. */
+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.
+   */
+  readonly submissions: ReadonlyMap<
+    Protocol.CellId,
+    Protocol.ExecutionId | null
+  >;
+  /** Cells the kernel reports as blocked by a disabled ancestor. */
+  readonly blocked: ReadonlySet<Protocol.CellId>;
+}
+
+export function outcomeOf(record: Protocol.Execution): CellOutcome | null {
+  const outcome = record.outcome;
+  if (outcome === null) {
+    return null;
+  }
+  return outcome.kind;
+}
+
+function millis(iso: string | null): Milliseconds | null {
+  return iso === null ? null : Time.parse(iso).toMilliseconds();
+}
+
+function latestOf(
+  executions: Executions,
+  id: Protocol.CellId,
+): Execution | null {
+  const executionId = executions.latest.get(id);
+  return executionId === undefined
+    ? null
+    : (executions.executions.get(executionId) ?? null);
+}
+
+/**
+ * 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.
+ */
+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;
+}
+
+export function cellExecutionView(
+  inputs: CellExecutionInputs,
+  id: Protocol.CellId,
+): 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;
+  if (latest !== null && latest.record.status !== "completed") {
+    status = latest.record.status;
+    startedAt = millis(latest.record.started_at) ?? startedAt;
+    rendered = latest.record.outcome !== null;
+  } else if (submitted) {
+    status = "queued";
+  }
+  const settled = status === "idle";
+  const duration =
+    settled &&
+    result !== null &&
+    result.started_at !== null &&
+    result.completed_at !== null
+      ? Time.parse(result.completed_at).since(Time.parse(result.started_at))
+      : null;
+  return {
+    status,
+    blocked: inputs.blocked.has(id),
+    outcome: settled && result !== null ? outcomeOf(result) : null,
+    startedAt,
+    duration,
+    rendered,
+  };
+}
+
+/** Executions the host has not completed, cells and scratch alike. */
+function unfinished(executions: Executions): Execution[] {
+  return [...executions.executions.values()].filter(
+    (execution) => execution.record.status !== "completed",
+  );
+}
+
+export function anyRunning(executions: Executions): boolean {
+  return unfinished(executions).some(
+    (execution) => execution.record.status === "running",
+  );
+}
+
+/** Cells and scratch work queued or running, plus submissions the host has not shown yet. */
+export function queuedOrRunningCount(inputs: CellExecutionInputs): number {
+  const pending = unfinished(inputs.executions);
+  const shown = new Set(
+    pending.flatMap((execution) =>
+      execution.record.cell_id === null ? [] : [execution.record.cell_id],
+    ),
+  );
+  let awaiting = 0;
+  for (const [cellId, executionId] of inputs.submissions) {
+    if (
+      !shown.has(cellId) &&
+      executionId === latestExecutionId(inputs, cellId)
+    ) {
+      awaiting += 1;
+    }
+  }
+  return pending.length + awaiting;
+}
+
+export function onlyScratchRunning(executions: Executions): boolean {
+  return unfinished(executions).every(
+    (execution) =>
+      execution.record.status !== "running" ||
+      execution.record.cell_id === null,
+  );
+}
+
+/** Where the scratchpad's work is, from its unfinished execution if any. */
+export function scratchStatus(executions: Executions): CellRunStatus {
+  const scratch = unfinished(executions).find(
+    (execution) => execution.record.cell_id === null,
+  );
+  if (scratch === undefined || scratch.record.status === "completed") {
+    return "idle";
+  }
+  return scratch.record.status;
+}
+
+/** Cells with a running execution, in the order the host reported them. */
+export function runningCellIds(executions: Executions): Protocol.CellId[] {
+  return unfinished(executions).flatMap((execution) =>
+    execution.record.status === "running" && execution.record.cell_id !== null
+      ? [execution.record.cell_id]
+      : [],
+  );
+}

package/src/core/export/__tests__/hooks.test.ts

@@ -67,11 +67,6 @@ describe("useEnrichCellOutputs", () => {
         cellId,
         {
           output: cell.output || null,
-          status: cell.status || "idle",
-          interrupted: false,
-          errored: false,
-          runStartTimestamp: null,
-          runElapsedTimeMs: null,
           stallTime: null as unknown as number,
           ...cell,
         } as CellRuntimeState,

package/src/core/export/hooks.ts

@@ -16,7 +16,7 @@ import { connectionAtom } from "../network/connection";
 import { useRequestClient } from "../network/requests";
 import type { UpdateCellOutputsRequest } from "../network/types";
 import { VirtualFileTracker } from "../static/virtual-file-tracker";
-import { ConnectionState } from "../websocket/types";
+import { ConnectionState } from "../session/types";
 
 const DELAY = 5000; // 5 seconds;
 

package/src/core/export/layout.ts

@@ -1,9 +1,9 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { flushDocumentChanges } from "@/core/cells/document-changes";
+import { flushDocument } from "@/core/cells/document";
 import { getSerializedLayout } from "@/core/layout/layout";
 
 export async function getExportLayout() {
-  await flushDocumentChanges();
+  await flushDocument();
   return getSerializedLayout();
 }

package/src/core/hotkeys/actions.ts

@@ -2,6 +2,7 @@
 import { atom, useAtomValue, useSetAtom } from "jotai";
 import useEvent from "react-use-event-hook";
 import type { HotkeyAction } from "@/core/hotkeys/hotkeys";
+import { Objects } from "@/utils/objects";
 
 /**
  * Map of registered keyboard shortcuts and their callbacks.
@@ -29,10 +30,7 @@ export function useSetRegisteredAction() {
       },
     ),
     unregisterAction: useEvent((shortcut: HotkeyAction) => {
-      set((actions) => {
-        const { [shortcut]: unused, ...rest } = actions;
-        return rest;
-      });
+      set((actions) => Objects.omit(actions, [shortcut]));
     }),
   };
 }

package/src/core/hotkeys/hotkeys.ts

@@ -221,7 +221,8 @@ const DEFAULT_HOT_KEY = {
     key: "Mod-h",
   },
   "cell.aiCompletion": {
-    name: "AI completion",
+    name: "Refactor cell",
+    additionalKeywords: ["AI", "AI completion"],
     group: "Editing",
     key: "Mod-Shift-e",
   },

package/src/core/islands/__tests__/bootstrap.test.ts

@@ -1,22 +1,29 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
+import type { NotebookTransport } from "@marimo-team/client";
 import { createStore } from "jotai";
 import { afterEach, describe, expect, it, vi } from "vitest";
+import { cellId } from "@/__tests__/branded";
+import { notebookAtom } from "@/core/cells/cells";
+import {
+  fakeBackend,
+  fixtureMessages,
+} from "@/core/notebook/replica/__tests__/fixtures";
 import { RuntimeState } from "@/core/kernel/RuntimeState";
 import { initializeIslands } from "../bootstrap";
 import type { IslandsPyodideBridge } from "../bridge";
 import { islandsHydratedAtom, islandsPendingInitialRunsAtom } from "../state";
-import type { IslandsKernelMessage } from "../worker/worker";
-
-function completedRun(sessionGeneration: number): IslandsKernelMessage {
-  return {
-    sessionGeneration,
-    message: JSON.stringify({
-      op: "completed-run",
-      data: { op: "completed-run", run_id: null },
-    }) as IslandsKernelMessage["message"],
-  };
-}
+import type { NotificationMessage } from "@/core/kernel/messages";
+
+type Consumer = (
+  notification: NotificationMessage,
+  sessionGeneration: number,
+) => void;
+
+const completedRun = {
+  op: "completed-run",
+  run_id: null,
+} as NotificationMessage;
 
 describe("initializeIslands", () => {
   afterEach(vi.restoreAllMocks);
@@ -26,12 +33,13 @@ describe("initializeIslands", () => {
     root.innerHTML =
       '<marimo-island data-app-id="app-1" data-reactive="true"></marimo-island>';
 
-    let consumeMessage: ((message: IslandsKernelMessage) => void) | undefined;
+    let consumeMessage: Consumer | undefined;
     const bridge = {
-      consumeMessages: (consumer: (message: IslandsKernelMessage) => void) => {
+      consumeNotifications: (consumer: Consumer) => {
         consumeMessage = consumer;
       },
       sendComponentValues: vi.fn(),
+      notebookBackend: fakeBackend(),
     } as unknown as IslandsPyodideBridge;
     const store = createStore();
     vi.spyOn(RuntimeState.INSTANCE, "start").mockImplementation(
@@ -48,14 +56,58 @@ describe("initializeIslands", () => {
     expect(store.get(islandsHydratedAtom)).toBe(false);
     store.set(islandsPendingInitialRunsAtom, new Set([1, 2]));
 
-    consumeMessage?.(completedRun(1));
+    consumeMessage?.(completedRun, 1);
     expect(store.get(islandsHydratedAtom)).toBe(false);
 
-    consumeMessage?.(completedRun(1));
-    consumeMessage?.(completedRun(0));
+    consumeMessage?.(completedRun, 1);
+    consumeMessage?.(completedRun, 0);
     expect(store.get(islandsHydratedAtom)).toBe(false);
 
-    consumeMessage?.(completedRun(2));
+    consumeMessage?.(completedRun, 2);
     expect(store.get(islandsHydratedAtom)).toBe(true);
   });
 });
+
+describe("initializeIslands document", () => {
+  afterEach(vi.restoreAllMocks);
+
+  it("fills the notebook's cells from the managed session's host document", async () => {
+    const root = document.createElement("div");
+    root.innerHTML =
+      '<marimo-island data-app-id="app-1" data-reactive="true"></marimo-island>';
+    let handlers: Parameters<NotebookTransport>[1] | undefined;
+    const transport: NotebookTransport = (_attempt, opened) =>
+      new Promise(() => {
+        handlers = opened;
+        opened.open();
+      });
+    const bridge = {
+      consumeNotifications: vi.fn(),
+      sendComponentValues: vi.fn(),
+      notebookBackend: fakeBackend({ transport }),
+    } as unknown as IslandsPyodideBridge;
+    const store = createStore();
+    vi.spyOn(RuntimeState.INSTANCE, "start").mockImplementation(
+      () => undefined,
+    );
+
+    await initializeIslands({
+      bridge,
+      store,
+      root,
+      autoInitializePlugins: false,
+    });
+    await vi.waitFor(() => expect(handlers).toBeDefined());
+    for (const message of fixtureMessages("notebook-events.sse")) {
+      handlers?.message(message);
+    }
+
+    const notebook = store.get(notebookAtom);
+    expect(notebook.cellIds).toEqual(["c1", "c3", "k7f2"]);
+    expect(notebook.cellData[cellId("c3")].code).toBe("x = 21");
+    expect(notebook.cellRuntime[cellId("c3")]).toBeDefined();
+
+    handlers?.message({ event: "reset", data: {} });
+    expect(store.get(notebookAtom).cellIds).toEqual(["c1", "c3", "k7f2"]);
+  });
+});

package/src/core/islands/__tests__/bridge.test.ts

@@ -10,7 +10,6 @@ import {
 } from "@/__tests__/branded";
 import { notebookAtom } from "@/core/cells/cells";
 import { islandsPendingInitialRunsAtom } from "@/core/islands/state";
-import type { IslandsKernelMessage } from "../worker/worker";
 
 type Base64String = components["schemas"]["Base64String"];
 interface TestIslandApp {
@@ -89,9 +88,6 @@ vi.mock("@/core/wasm/rpc", () => ({
         startSession: mockStartSessionRequest,
         stopSession: mockStopSessionRequest,
       },
-      send: {
-        consumerReady: vi.fn(),
-      },
     },
     addMessageListener: vi.fn(
       (name: string, listener: (payload: never) => void) => {
@@ -177,17 +173,22 @@ describe("IslandsPyodideBridge", () => {
     });
   }
 
-  function sendKernelMessage(sessionGeneration: number) {
-    const listener = mockMessageListeners.get("kernelMessage");
+  function sendNotebookEvent(sessionGeneration: number, frame: unknown) {
+    const listener = mockMessageListeners.get("notebookEvent");
     if (!listener) {
-      throw new Error("Missing kernel message listener");
+      throw new Error("Missing notebook event listener");
     }
-    const message = {
-      message: "" as IslandsKernelMessage["message"],
-      sessionGeneration,
-    };
-    listener(message as never);
-    return message;
+    listener({ event: JSON.stringify(frame), sessionGeneration } as never);
+  }
+
+  /** Send a notification and return the call the consumer should receive. */
+  function sendKernelMessage(sessionGeneration: number) {
+    const notification = { op: "interrupted" };
+    sendNotebookEvent(sessionGeneration, {
+      event: "marimo",
+      data: { marimo: notification },
+    });
+    return [notification, sessionGeneration];
   }
 
   async function initializeSingleApp() {
@@ -300,7 +301,7 @@ describe("IslandsPyodideBridge", () => {
   describe("app lifecycle", () => {
     it("ignores messages after a session yields ownership", async () => {
       const consumeMessage = vi.fn();
-      bridge.consumeMessages(consumeMessage);
+      bridge.consumeNotifications(consumeMessage);
       mockSingleApp();
       await bridge.initializeApps();
 
@@ -319,10 +320,7 @@ describe("IslandsPyodideBridge", () => {
       sendKernelMessage(1);
       const currentMessage = sendKernelMessage(2);
 
-      expect(consumeMessage.mock.calls).toEqual([
-        [firstMessage],
-        [currentMessage],
-      ]);
+      expect(consumeMessage.mock.calls).toEqual([firstMessage, currentMessage]);
 
       finishReplacement();
       await replacement;
@@ -330,7 +328,7 @@ describe("IslandsPyodideBridge", () => {
 
     it("restores message ownership when replacement fails", async () => {
       const consumeMessage = vi.fn();
-      bridge.consumeMessages(consumeMessage);
+      bridge.consumeNotifications(consumeMessage);
       mockSingleApp();
       await bridge.initializeApps();
       mockReplaceSessionRequest.mockRejectedValueOnce(
@@ -344,12 +342,12 @@ describe("IslandsPyodideBridge", () => {
       const previousMessage = sendKernelMessage(1);
       sendKernelMessage(2);
 
-      expect(consumeMessage.mock.calls).toEqual([[previousMessage]]);
+      expect(consumeMessage.mock.calls).toEqual([previousMessage]);
     });
 
     it("ignores messages while their session is stopping", async () => {
       const consumeMessage = vi.fn();
-      bridge.consumeMessages(consumeMessage);
+      bridge.consumeNotifications(consumeMessage);
       mockSingleApp();
       await bridge.initializeApps();
       let finishStop!: () => void;
@@ -373,7 +371,7 @@ describe("IslandsPyodideBridge", () => {
 
     it("forwards messages from every app in a multi-app document", async () => {
       const consumeMessage = vi.fn();
-      bridge.consumeMessages(consumeMessage);
+      bridge.consumeNotifications(consumeMessage);
       mockParseMarimoIslandApps.mockReturnValue([app(), app("app-2")]);
       mockCreateMarimoFile
         .mockReturnValueOnce("generated app 1")
@@ -385,9 +383,58 @@ describe("IslandsPyodideBridge", () => {
       const secondMessage = sendKernelMessage(2);
       sendKernelMessage(3);
 
-      expect(consumeMessage.mock.calls).toEqual([
-        [firstMessage],
-        [secondMessage],
+      expect(consumeMessage.mock.calls).toEqual([firstMessage, secondMessage]);
+    });
+
+    it("delivers only the notebook stream's marimo notifications", async () => {
+      const consumeMessage = vi.fn();
+      bridge.consumeNotifications(consumeMessage);
+      await initializeSingleApp();
+      sendNotebookEvent(1, { event: "ready", id: "3", data: {} });
+      const notification = sendKernelMessage(1);
+
+      expect(consumeMessage.mock.calls).toEqual([notification]);
+    });
+
+    it("forwards the managed session's host events to the notebook stream", async () => {
+      await initializeSingleApp();
+      const handlers = { open: vi.fn(), message: vi.fn() };
+      void bridge.notebookBackend.transport(
+        { signal: new AbortController().signal, cursor: null },
+        handlers,
+      );
+      expect(handlers.open).toHaveBeenCalledOnce();
+
+      const cell = { event: "cell", id: "1", data: { cell: { id: "c1" } } };
+      sendNotebookEvent(1, cell);
+      sendNotebookEvent(2, { event: "ready", id: "2", data: {} });
+      sendKernelMessage(1);
+
+      expect(handlers.message.mock.calls).toEqual([[cell]]);
+      expect(() => bridge.notebookBackend.notebook("n1")).toThrow(
+        "Islands do not edit",
+      );
+    });
+
+    it("resets the notebook stream when the managed session is replaced or stopped", async () => {
+      await initializeSingleApp();
+      const handlers = { open: vi.fn(), message: vi.fn() };
+      void bridge.notebookBackend.transport(
+        { signal: new AbortController().signal, cursor: null },
+        handlers,
+      );
+
+      mockSingleApp("generated app 2");
+      await bridge.initializeApps();
+      const ready = { event: "ready", id: "2", data: {} };
+      sendNotebookEvent(1, { event: "ready", id: "1", data: {} });
+      sendNotebookEvent(2, ready);
+      await bridge.stopSession("app-1");
+
+      expect(handlers.message.mock.calls).toEqual([
+        [{ event: "reset", data: {} }],
+        [ready],
+        [{ event: "reset", data: {} }],
       ]);
     });
 
@@ -723,8 +770,19 @@ describe("IslandsPyodideBridge", () => {
         functionName: "put_control_request",
         payload: {
           type: "execute-cells",
-          cellIds: ["cell-1", "cell-2"],
-          codes: ["print('hello')", "print('world')"],
+          request: undefined,
+          executionRequests: [
+            {
+              type: "execute-cell",
+              cellId: "cell-1",
+              code: "print('hello')",
+            },
+            {
+              type: "execute-cell",
+              cellId: "cell-2",
+              code: "print('world')",
+            },
+          ],
         },
         sessionGeneration: 1,
       });