@marimo-team/islands 0.25.3-dev3 → 0.25.3-dev301

package/src/core/notebook/replica/__tests__/results.test.ts

@@ -0,0 +1,193 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import * as Protocol from "@marimo-team/client/protocol";
+import { results } from "../resources/results";
+import { createNotebookReducer } from "@marimo-team/client";
+import { fixtureEvents } from "./fixtures";
+
+const reducer = createNotebookReducer({ results });
+const snapshot = fixtureEvents("notebook-events.sse");
+const runCells = fixtureEvents("run-cells.sse");
+const time = "2026-10-01T09:30:00Z";
+
+function replay(
+  events: readonly Protocol.EventOutput[],
+  from = reducer.initial,
+) {
+  return events.reduce(reducer.reduce, from);
+}
+
+const resource = (state: ReturnType<typeof replay>) => state.resources.results;
+
+/** Build one event as the stream would deliver it. */
+const event = (message: unknown) => Protocol.Event.parse(message);
+
+const removed = (kind: string, id: string, cursor: string) =>
+  event({ event: `${kind}.removed`, id: cursor, data: { id, time } });
+
+/** How the host retires a cell execution: its items first, then itself. */
+const retire = (executionId: string, variableId: string) => [
+  removed("variable", variableId, "111"),
+  removed("execution", executionId, "112"),
+];
+
+const restored = {
+  id: "e-restored",
+  run_id: null,
+  runtime_id: null,
+  caused_by: null,
+  cell_id: "c3",
+  code: null,
+  code_fingerprint: "0f71456db2ea83f4",
+  depends_on: {},
+  status: "completed",
+  outcome: { kind: "value", display: null },
+  started_at: "2026-09-30T09:20:01Z",
+  completed_at: "2026-09-30T09:20:02Z",
+};
+
+describe("results resource", () => {
+  it("has an entry per document cell, filled by a restored result", () => {
+    const ready = snapshot.findIndex((item) => item.event === "ready");
+    const opening = resource(replay(snapshot.slice(0, ready + 1)));
+    expect([...opening]).toEqual([
+      ["c1", null],
+      ["c3", null],
+      ["k7f2", null],
+    ]);
+
+    const opened = resource(
+      replay([
+        ...snapshot.slice(0, ready),
+        event({ event: "execution", id: "94", data: { execution: restored } }),
+        ...snapshot.slice(ready),
+      ]),
+    );
+    expect(opened.get("c3")?.id).toBe("e-restored");
+    expect(opened.get("c3")?.run_id).toBeNull();
+  });
+
+  it("records a cell's result when its execution completes", () => {
+    const start = replay(snapshot);
+    const running = resource(replay(runCells.slice(0, 5), start));
+    expect(running.get("c3")).toBeNull();
+    expect(running.get("k7f2")).toBeNull();
+
+    const done = resource(replay(runCells, start));
+    expect(done.get("c3")?.id).toBe("e12");
+    expect(done.get("c3")?.status).toBe("completed");
+    expect(done.get("k7f2")?.depends_on).toEqual({
+      x: { execution: "e12", update: 0 },
+    });
+    expect(done.get("c1")).toBeNull();
+  });
+
+  it("keeps a result while the host replaces it", () => {
+    const done = replay(runCells, replay(snapshot));
+    const result = resource(done).get("c3")!;
+    const retired = replay(retire("e12", "e12-x"), done);
+    expect(resource(retired).get("c3")).toBe(result);
+
+    const again = {
+      ...result,
+      id: "e14",
+      run_id: "run-again",
+      status: "queued",
+      outcome: null,
+      started_at: null,
+      completed_at: null,
+    };
+    const queued = reducer.reduce(
+      retired,
+      event({ event: "execution", id: "113", data: { execution: again } }),
+    );
+    expect(resource(queued).get("c3")).toBe(result);
+
+    const completed = reducer.reduce(
+      queued,
+      event({
+        event: "execution",
+        id: "114",
+        data: {
+          execution: {
+            ...again,
+            status: "completed",
+            outcome: { kind: "value", display: null },
+            started_at: time,
+            completed_at: time,
+          },
+        },
+      }),
+    );
+    expect(resource(completed).get("c3")?.id).toBe("e14");
+  });
+
+  it("replaces a result with an execution interrupted before it started", () => {
+    const done = replay(runCells, replay(snapshot));
+    const result = resource(done).get("c3")!;
+    const interrupted = {
+      ...result,
+      id: "e14",
+      run_id: "run-again",
+      status: "completed",
+      outcome: { kind: "interrupted" },
+      started_at: null,
+      completed_at: time,
+    };
+    const cancelled = replay(
+      [
+        ...retire("e12", "e12-x"),
+        event({
+          event: "execution",
+          id: "113",
+          data: {
+            execution: {
+              ...interrupted,
+              status: "queued",
+              outcome: null,
+              completed_at: null,
+            },
+          },
+        }),
+        event({
+          event: "execution",
+          id: "114",
+          data: { execution: interrupted },
+        }),
+      ],
+      done,
+    );
+    expect(resource(cancelled).get("c3")?.outcome).toEqual({
+      kind: "interrupted",
+    });
+  });
+
+  it("drops a cell deleted while its execution runs", () => {
+    const start = replay(runCells.slice(0, 5), replay(snapshot));
+    const deleted = reducer.reduce(start, removed("cell", "c3", "105"));
+    expect(resource(deleted).has("c3")).toBe(false);
+
+    const finished = replay(runCells.slice(5), deleted);
+    expect(resource(finished).has("c3")).toBe(false);
+    expect(resource(finished).get("k7f2")?.id).toBe("e13");
+  });
+
+  it("ignores cells it has not seen and clears on reset", () => {
+    const start = replay(snapshot);
+    const unknown = reducer.reduce(
+      start,
+      event({
+        event: "execution",
+        id: "98",
+        data: { execution: { ...restored, id: "e-zz", cell_id: "zz" } },
+      }),
+    );
+    expect(unknown.resources).toBe(start.resources);
+    const reset = reducer.reduce(replay(runCells, start), {
+      event: "reset",
+      data: {},
+    });
+    expect(resource(reset)).toBe(results.initial);
+  });
+});

package/src/core/notebook/replica/__tests__/runtime.test.ts

@@ -0,0 +1,78 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { runtime } from "../resources/runtime";
+import { createNotebookReducer } from "@marimo-team/client";
+import { fixtureEvents } from "./fixtures";
+
+const reducer = createNotebookReducer({ runtime });
+const snapshot = fixtureEvents("notebook-events.sse");
+const confirmed = () => snapshot.reduce(reducer.reduce, reducer.initial);
+const current = (state: ReturnType<typeof confirmed>) =>
+  state.resources.runtime;
+const time = "2026-10-01T09:30:00Z";
+
+describe("runtime resource", () => {
+  it("keeps the latest record and clears when that runtime is removed", () => {
+    const state = confirmed();
+    const running = current(state);
+    expect(running?.id).toBe("r1");
+    expect(running?.status).toBe("running");
+
+    const terminating = reducer.reduce(state, {
+      event: "runtime",
+      id: "98",
+      data: { runtime: { ...running!, status: "terminating" } },
+    });
+    expect(current(terminating)?.status).toBe("terminating");
+    expect(current(state)?.status).toBe("running");
+
+    const other = reducer.reduce(terminating, {
+      event: "runtime.removed",
+      id: "99",
+      data: { id: "r0", time },
+    });
+    expect(current(other)).toBe(current(terminating));
+
+    const removed = reducer.reduce(other, {
+      event: "runtime.removed",
+      id: "100",
+      data: { id: "r1", time },
+    });
+    expect(current(removed)).toBeNull();
+    const repeated = reducer.reduce(removed, {
+      event: "runtime.removed",
+      id: "101",
+      data: { id: "r1", time },
+    });
+    expect(repeated.resources).toBe(removed.resources);
+  });
+
+  it("follows a restart to the replacement runtime", () => {
+    const state = confirmed();
+    const removed = reducer.reduce(state, {
+      event: "runtime.removed",
+      id: "98",
+      data: { id: "r1", time },
+    });
+    const replacement = {
+      ...current(state)!,
+      id: "r2",
+      status: "starting" as const,
+      message: "Starting the kernel",
+      marimo_version: null,
+      started_at: "2026-10-01T09:31:00Z",
+    };
+    const starting = reducer.reduce(removed, {
+      event: "runtime",
+      id: "99",
+      data: { runtime: replacement },
+    });
+    expect(current(starting)).toEqual(replacement);
+  });
+
+  it("clears on reset", () => {
+    const reset = reducer.reduce(confirmed(), { event: "reset", data: {} });
+    expect(current(reset)).toBeNull();
+  });
+});

package/src/core/notebook/replica/backend.ts

@@ -0,0 +1,19 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** A notebook's events and commands, from a host or a stand-in for one. */
+
+import type { NotebookClient, NotebookTransport } from "@marimo-team/client";
+import type { NotebookId } from "@marimo-team/client/protocol";
+
+/** Events and commands for a notebook, independent of React or a replica. */
+export interface NotebookBackend {
+  readonly transport: NotebookTransport;
+  /**
+   * Commands for a notebook the transport reported.
+   *
+   * Method presence is not permission. The notebook advertises its operations.
+   */
+  notebook(id: NotebookId): Pick<NotebookClient, "edit" | "save">;
+  /** Whether the host is reachable before an explicit reconnect. Absent means always. */
+  checkHealth?(): Promise<boolean>;
+}

package/src/core/{websocket/types.ts → notebook/replica/connection.ts}

@@ -1,11 +1,5 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import type { NotificationMessageData } from "../kernel/messages";
-
-export type ConnectionPhase =
-  | NotificationMessageData<"startup-progress">["phase"]
-  | "reconnecting";
-
 export const ConnectionState = {
   NOT_STARTED: "NOT_STARTED",
   CONNECTING: "CONNECTING",
@@ -17,27 +11,22 @@ export const ConnectionState = {
 export type ConnectionState =
   (typeof ConnectionState)[keyof typeof ConnectionState];
 
-export const WebSocketClosedReason = {
+export const ConnectionClosedReason = {
   KERNEL_DISCONNECTED: "KERNEL_DISCONNECTED",
   KERNEL_STARTUP_ERROR: "KERNEL_STARTUP_ERROR",
 } as const;
 
-export type WebSocketClosedReason =
-  (typeof WebSocketClosedReason)[keyof typeof WebSocketClosedReason];
+export type ConnectionClosedReason =
+  (typeof ConnectionClosedReason)[keyof typeof ConnectionClosedReason];
 
 export type ConnectionStatus =
   | {
       state: typeof ConnectionState.CLOSED;
-      code: WebSocketClosedReason;
-      /**
-       * Human-readable reason for closing the connection.
-       */
+      code: ConnectionClosedReason;
       reason: string;
-      phase?: ConnectionPhase;
     }
   | {
       state: typeof ConnectionState.CONNECTING;
-      phase?: ConnectionPhase;
     }
   | {
       state:
@@ -45,3 +34,9 @@ export type ConnectionStatus =
         | typeof ConnectionState.CLOSING
         | typeof ConnectionState.NOT_STARTED;
     };
+
+/** The recovery policy for a closed connection. */
+export interface CloseDecision {
+  kind: "terminal" | "retry";
+  status: ConnectionStatus;
+}

package/src/core/notebook/replica/index.ts

@@ -0,0 +1,20 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** Backend-confirmed notebook state, independent of React and transports. */
+
+export type { NotebookBackend } from "./backend";
+export {
+  type CloseDecision,
+  ConnectionClosedReason,
+  ConnectionState,
+  type ConnectionStatus,
+} from "./connection";
+export { createNotebookReplica, type NotebookReplica } from "./replica";
+export type { Document } from "./resources/document";
+export {
+  type Execution,
+  type Executions,
+  type Item,
+  itemsInOrder,
+} from "./resources/executions";
+export type { Results } from "./resources/results";

package/src/core/notebook/replica/replica.ts

@@ -0,0 +1,336 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/** What a backend's stream has confirmed about a notebook, kept in atoms. */
+
+import { Logger } from "@/utils/Logger";
+import {
+  createNotebookReducer,
+  type NotebookEvent,
+  NotebookEvents,
+  StreamStopped,
+} from "@marimo-team/client";
+import type {
+  Attachment,
+  NotebookOutput,
+  Runtime,
+} from "@marimo-team/client/protocol";
+import { type Atom, atom, type createStore } from "jotai/vanilla";
+import { z } from "zod";
+import type { NotebookBackend } from "./backend";
+import {
+  ConnectionClosedReason,
+  ConnectionState,
+  type ConnectionStatus,
+} from "./connection";
+import { attachments } from "./resources/attachments";
+import { type Document, document } from "./resources/document";
+import { type Executions, executions } from "./resources/executions";
+import { type Results, results } from "./resources/results";
+import { runtime } from "./resources/runtime";
+
+/**
+ * The kinds the editor reads, with marimo's own notifications.
+ *
+ * Console and display items stay out until the editor reads outputs from
+ * the replica. The native notifications still carry them.
+ */
+const TYPES = [
+  "notebook",
+  "cell",
+  "document",
+  "attachment",
+  "runtime",
+  "run",
+  "execution",
+  "input",
+  "variable",
+  "marimo",
+] as const;
+
+type EditorEvent = NotebookEvent<(typeof TYPES)[number]>;
+
+/** A native marimo notification, passed through without interpretation. */
+const MarimoMessage = z.object({ marimo: z.unknown() });
+
+/** A notebook as its backend's stream reports it, kept current while connected. */
+export interface NotebookReplica {
+  /** Connection status, independent of application session readiness. */
+  readonly connectionAtom: Atom<ConnectionStatus>;
+  /** Identity and capabilities, available while the stream is connected and ready. */
+  readonly notebookAtom: Atom<NotebookOutput | null>;
+  /**
+   * The most recently confirmed document.
+   *
+   * Null until a snapshot is complete, after reset or removal, and on
+   * connections without host document events. A temporary disconnection
+   * retains the last document. Partial updates are not exposed.
+   */
+  readonly documentAtom: Atom<Document | null>;
+  /**
+   * The backend this replica follows, or null while not connected.
+   *
+   * Its commands reach the host that confirmed `documentAtom`. A temporary
+   * disconnection keeps it. Connecting again or disconnecting replaces it.
+   */
+  readonly backendAtom: Atom<NotebookBackend | null>;
+  /** Live attachments, empty until the stream is ready or while disconnected. */
+  readonly attachmentsAtom: Atom<readonly Attachment[]>;
+  /** The notebook's runtime, null while it has none or the stream is not ready. */
+  readonly runtimeAtom: Atom<Runtime | null>;
+  /**
+   * Runs and executions as the stream reports them.
+   *
+   * Visible as events arrive, including during a snapshot, so a consumer
+   * joins with `documentAtom`, which is null until the snapshot is ready.
+   * Kept across a temporary disconnection, like the document, and cleared
+   * on reset or removal.
+   */
+  readonly executionsAtom: Atom<Executions>;
+  /** Each document cell's last completed execution record, kept while it runs again. */
+  readonly resultsAtom: Atom<Results>;
+  /**
+   * Follow a backend, replacing any current connection.
+   *
+   * The returned function ends this connection and clears the replica.
+   * It does nothing once another connection has replaced this one.
+   */
+  connect(
+    backend: NotebookBackend,
+    options: {
+      /** Handle a native marimo notification. Its payload has not been validated. */
+      onNotification: (notification: unknown) => void;
+    },
+  ): () => void;
+  /** Restore the current connection after its stream closed, if the host is healthy. */
+  reconnect(): Promise<void>;
+}
+
+const disconnected = {
+  state: ConnectionState.CLOSED,
+  code: ConnectionClosedReason.KERNEL_DISCONNECTED,
+  reason: "kernel not found",
+} satisfies ConnectionStatus;
+
+const malformed = {
+  ...disconnected,
+  reason: "Failed to decode notebook message",
+} satisfies ConnectionStatus;
+
+function isStatus(value: unknown): value is ConnectionStatus {
+  return typeof value === "object" && value !== null && "state" in value;
+}
+
+/**
+ * The status to show after the stream stopped.
+ *
+ * A transport says which status its stop means. An event the protocol
+ * rejects is malformed. Anything else means the host was not reached.
+ */
+function statusAfter(error: unknown): ConnectionStatus {
+  if (error instanceof StreamStopped && isStatus(error.cause)) {
+    return error.cause;
+  }
+  if (error instanceof StreamStopped && error.cause instanceof z.ZodError) {
+    return malformed;
+  }
+  return disconnected;
+}
+
+/** Create a replica without connecting it. */
+export function createNotebookReplica(
+  store: ReturnType<typeof createStore>,
+): NotebookReplica {
+  // `useNotebookSession` copies this into the editor's shared `connectionAtom`,
+  // which gates features until MO-8051 derives availability separately.
+  const connectionAtom = atom<ConnectionStatus>({
+    state: ConnectionState.NOT_STARTED,
+  });
+  // One atom holds the whole stream, which resource atoms derive from.
+  const reducer = createNotebookReducer({
+    document,
+    attachments,
+    runtime,
+    executions,
+    results,
+  });
+  const streamAtom = atom(reducer.initial);
+  // A resumed stream is live, but presence waits for this connection's ready.
+  const readyAtom = atom(false);
+  const presentAtom = atom(
+    (get) =>
+      get(readyAtom) && get(connectionAtom).state === ConnectionState.OPEN,
+  );
+  const notebookAtom = atom((get) =>
+    get(presentAtom) ? get(streamAtom).notebook : null,
+  );
+  const documentAtom = atom(
+    (get) => get(streamAtom).resources.document.current,
+  );
+  const attachmentsAtom = atom((get) =>
+    get(presentAtom)
+      ? get(streamAtom).resources.attachments
+      : attachments.initial,
+  );
+  const runtimeAtom = atom((get) =>
+    get(presentAtom) ? get(streamAtom).resources.runtime : null,
+  );
+  const executionsAtom = atom((get) => get(streamAtom).resources.executions);
+  const resultsAtom = atom((get) => get(streamAtom).resources.results);
+  const backendAtom = atom<NotebookBackend | null>(null);
+  /** The connection this replica follows. Connecting again ends it. */
+  let current: { end(): void; reconnect(): Promise<void> } | undefined;
+
+  function setStatus(status: ConnectionStatus) {
+    store.set(readyAtom, false);
+    store.set(connectionAtom, status);
+  }
+
+  /** Drop what the stream reported and show `status`. */
+  function clear(status: ConnectionStatus) {
+    store.set(streamAtom, reducer.initial);
+    setStatus(status);
+  }
+
+  return {
+    connectionAtom,
+    notebookAtom,
+    documentAtom,
+    backendAtom,
+    attachmentsAtom,
+    runtimeAtom,
+    executionsAtom,
+    resultsAtom,
+    connect,
+    reconnect: async () => current?.reconnect(),
+  };
+
+  function connect(
+    connected: NotebookBackend,
+    { onNotification }: { onNotification: (notification: unknown) => void },
+  ) {
+    current?.end();
+    const lifetime = new AbortController();
+    const { signal } = lifetime;
+    const checkHealth = connected.checkHealth ?? (async () => true);
+    store.set(backendAtom, connected);
+    const { transport } = connected;
+    let events: NotebookEvents<(typeof TYPES)[number]> | undefined;
+    let checkingHealth = false;
+    // Whether the replica still holds the state the last cursor describes.
+    let resumable = true;
+
+    /** Mirror one stream into the atoms until it closes. */
+    function follow(cursor?: string) {
+      if (cursor === undefined) {
+        store.set(streamAtom, reducer.initial);
+      }
+      resumable = true;
+      const current = new NotebookEvents(transport, {
+        types: TYPES,
+        cursor,
+        signal,
+        // TODO(MO-8051): Gate editing and running on the notebook's operations
+        // and request results instead. The stream's status would then only
+        // mark the view as stale.
+        onStatus: (status) => {
+          if (status === "connecting") {
+            setStatus({ state: ConnectionState.CONNECTING });
+          } else if (status === "connected") {
+            setStatus({ state: ConnectionState.OPEN });
+          }
+        },
+      });
+      current.subscribe(
+        (event) => apply(current, event),
+        (error) => {
+          if (!signal.aborted && current === events) {
+            Logger.warn("Notebook stream stopped", error);
+            const status = statusAfter(error);
+            if (status === malformed) {
+              stopMalformed(current);
+            } else {
+              setStatus(status);
+            }
+          }
+        },
+      );
+      events = current;
+    }
+
+    function stopMalformed(source: NotebookEvents<string>) {
+      source.close();
+      // The cursor can already include an event the replica rejected.
+      resumable = false;
+      clear(malformed);
+    }
+
+    function apply(source: NotebookEvents<string>, event: EditorEvent) {
+      if (event.event === "marimo") {
+        const message = MarimoMessage.safeParse(event.data);
+        if (message.success) {
+          onNotification(message.data.marimo);
+        } else {
+          Logger.error("Failed to decode notebook message", event);
+          stopMalformed(source);
+        }
+        return;
+      }
+      try {
+        store.set(streamAtom, reducer.reduce(store.get(streamAtom), event));
+      } catch (error) {
+        // An event the replica cannot apply, such as another notebook's.
+        Logger.error("Failed to apply notebook event", error);
+        stopMalformed(source);
+        return;
+      }
+      if (event.event === "ready" || event.event === "reset") {
+        store.set(readyAtom, event.event === "ready");
+      }
+    }
+
+    const reconnect = async () => {
+      if (signal.aborted || checkingHealth || events?.status !== "closed") {
+        return;
+      }
+      checkingHealth = true;
+      setStatus({ state: ConnectionState.CONNECTING });
+      try {
+        const healthy = await checkHealth();
+        if (signal.aborted) {
+          return;
+        }
+        if (!healthy) {
+          setStatus(disconnected);
+          return;
+        }
+        follow(resumable ? (events.cursor ?? undefined) : undefined);
+      } catch (error) {
+        if (!signal.aborted) {
+          Logger.error("Notebook health check failed", error);
+          setStatus(disconnected);
+        }
+      } finally {
+        checkingHealth = false;
+      }
+    };
+
+    const connection = {
+      end: () => {
+        if (signal.aborted) {
+          return;
+        }
+        lifetime.abort();
+        if (current === connection) {
+          current = undefined;
+          store.set(backendAtom, null);
+        }
+        clear({ state: ConnectionState.NOT_STARTED });
+      },
+      reconnect,
+    };
+    current = connection;
+    setStatus({ state: ConnectionState.CONNECTING });
+    follow();
+    return connection.end;
+  }
+}