@marimo-team/frontend 0.25.2 → 0.25.3-dev12

package/src/core/edit-app.tsx

@@ -41,7 +41,7 @@ import { viewStateAtom } from "./mode";
 import { useRequestClient } from "./network/requests";
 import { useFilename } from "./saving/filename";
 import { lastSavedNotebookAtom } from "./saving/state";
-import { useMarimoKernelConnection } from "./websocket/useMarimoKernelConnection";
+import { useNotebookSession } from "./websocket/useNotebookSession";
 
 interface AppProps {
   /**
@@ -83,7 +83,7 @@ export const EditApp: React.FC<AppProps> = ({
     };
   }, []);
 
-  const { connection, reconnect } = useMarimoKernelConnection({
+  const { connection, reconnect } = useNotebookSession({
     autoInstantiate: userConfig.runtime.auto_instantiate,
     setCells: (cells, layout) => {
       setCells(cells);

package/src/core/host/AGENTS.md

@@ -0,0 +1,6 @@
+# Host core
+
+Production code may import host modules, `jotai/vanilla`, `zod`, and
+`@/utils/Logger`. Enforced by `marimo/host-imports`. Tests are exempt.
+
+Define interfaces here and inject application behavior through adapters.

package/src/core/host/__tests__/events.test.ts

@@ -0,0 +1,46 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { z } from "zod";
+import { parseNotebookEvent } from "../events";
+
+describe("notebook messages", () => {
+  it("normalizes session and host envelopes for native marimo messages", () => {
+    const notification = {
+      op: "variables",
+      variables: [{ name: "x", value: Number.NaN }],
+    };
+    const host = { event: "marimo", data: { marimo: notification } };
+    expect(parseNotebookEvent({ data: notification })).toEqual(host);
+    expect(parseNotebookEvent({ ...host, id: "42" })).toEqual({
+      ...host,
+      id: "42",
+    });
+  });
+
+  it("retains unknown event cursors without interpreting their payloads", () => {
+    expect(
+      parseNotebookEvent({ event: "future", id: "42", data: null }),
+    ).toEqual({ event: "unknown", name: "future", id: "42" });
+  });
+
+  it("decodes known events outside the client's subscription filters", () => {
+    const event = {
+      event: "run.removed",
+      id: "42",
+      data: { id: "r1", time: "2026-10-19T12:00:00Z" },
+    };
+    expect(parseNotebookEvent(event)).toEqual(event);
+  });
+
+  it.each([
+    { event: "cell", id: "42", data: { cell: { id: "c1" } } },
+    { event: "run.removed", id: "42", data: {} },
+    { event: "ready", data: {} },
+    { event: "marimo", data: {} },
+    {},
+    null,
+  ])("rejects malformed known messages: %j", (message) => {
+    expect(() => parseNotebookEvent(message)).toThrow(z.ZodError);
+  });
+});

package/src/core/host/__tests__/notebook-client.test.ts

@@ -0,0 +1,191 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+// @vitest-environment jsdom
+
+import { describe, expect, it, vi } from "vitest";
+import { createStore } from "jotai";
+import { ConnectionState } from "../../websocket/types";
+import { Deferred } from "@/utils/Deferred";
+import { jsonParseWithSpecialChar } from "@/utils/json/json-parser";
+import { classifyCloseEvent } from "../../websocket/connection-state";
+import { createNotebookClient } from "../notebook-client";
+
+function setup() {
+  const events = new EventTarget();
+  const healthy = new Deferred<void>();
+  const transport = {
+    readyState: WebSocket.CLOSED,
+    reconnect: vi.fn(),
+    close: vi.fn(() => events.dispatchEvent(new CloseEvent("close"))),
+    send: vi.fn(),
+    addEventListener: events.addEventListener.bind(events),
+    removeEventListener: events.removeEventListener.bind(events),
+  };
+  const onNotification = vi.fn();
+  const checkHealth = vi.fn(async () => true);
+  const store = createStore();
+  const controller = createNotebookClient(store);
+  const client = controller.client;
+  const connect = (signal: AbortSignal) =>
+    controller.connect({
+      createTransport: () => transport,
+      url: () => "ws://localhost/ws",
+      messageFormat: "host",
+      signal,
+      waitUntilReady: () => healthy.promise,
+      checkHealth,
+      onNotification: (notification) => onNotification({ data: notification }),
+      onMessageError: vi.fn(),
+      decodeJson: jsonParseWithSpecialChar,
+      classifyClose: classifyCloseEvent,
+    });
+  return {
+    events,
+    healthy,
+    transport,
+    client,
+    connect,
+    onNotification,
+    store,
+    checkHealth,
+  };
+}
+
+describe("NotebookClient lifetime", () => {
+  it.each(["{broken", JSON.stringify({ event: "ready", data: {} })])(
+    "stops consuming a malformed host stream: %s",
+    (data) => {
+      const { connect, events, client, store, transport, onNotification } =
+        setup();
+      const lifetime = new AbortController();
+      connect(lifetime.signal);
+      events.dispatchEvent(new Event("open"));
+      events.dispatchEvent(new MessageEvent("message", { data }));
+      events.dispatchEvent(
+        new MessageEvent("message", {
+          data: JSON.stringify({
+            event: "marimo",
+            data: { marimo: { op: "reconnected" } },
+          }),
+        }),
+      );
+      expect(store.get(client.connectionAtom)).toEqual({
+        state: ConnectionState.CLOSED,
+        code: "KERNEL_DISCONNECTED",
+        reason: "Failed to decode notebook message",
+      });
+      expect(transport.close).toHaveBeenCalledOnce();
+      expect(onNotification).not.toHaveBeenCalled();
+      lifetime.abort();
+    },
+  );
+
+  it.each(["MARIMO_UNAUTHORIZED", "MARIMO_TRANSPORT_EXHAUSTED"])(
+    "stops automatic recovery after %s but permits an explicit reconnect",
+    async (reason) => {
+      const {
+        connect,
+        client,
+        store,
+        events,
+        transport,
+        healthy,
+        checkHealth,
+      } = setup();
+      const lifetime = new AbortController();
+      connect(lifetime.signal);
+      healthy.resolve();
+      await healthy.promise;
+      events.dispatchEvent(new CloseEvent("close", { reason }));
+      events.dispatchEvent(new Event("error"));
+      expect(store.get(client.connectionAtom).state).toBe(
+        ConnectionState.CLOSED,
+      );
+      expect(transport.reconnect).toHaveBeenCalledOnce();
+      await client.reconnect();
+      expect(checkHealth).toHaveBeenCalledOnce();
+      expect(transport.reconnect).toHaveBeenCalledTimes(2);
+      lifetime.abort();
+    },
+  );
+
+  it("only resets the automatic retry budget after an open event", () => {
+    const { connect, events, transport } = setup();
+    const lifetime = new AbortController();
+    connect(lifetime.signal);
+    events.dispatchEvent(new CloseEvent("close"));
+    events.dispatchEvent(new Event("error"));
+    events.dispatchEvent(new CloseEvent("close"));
+    expect(transport.reconnect).toHaveBeenCalledOnce();
+    events.dispatchEvent(new Event("open"));
+    events.dispatchEvent(new CloseEvent("close"));
+    expect(transport.reconnect).toHaveBeenCalledTimes(2);
+    lifetime.abort();
+  });
+
+  it("does not revive an aborted lifetime when a manual health check finishes", async () => {
+    const { connect, client, transport, checkHealth } = setup();
+    const health = new Deferred<boolean>();
+    checkHealth.mockReturnValue(health.promise);
+    const lifetime = new AbortController();
+    connect(lifetime.signal);
+    const attempt = client.reconnect();
+    lifetime.abort();
+    health.resolve(true);
+    await attempt;
+    expect(transport.reconnect).not.toHaveBeenCalled();
+  });
+
+  it("does not connect after abort, even when the health check finishes", async () => {
+    const { client, connect, healthy, transport } = setup();
+    const controller = new AbortController();
+    connect(controller.signal);
+    controller.abort();
+    healthy.resolve();
+    await healthy.promise;
+    client.reconnect();
+    expect(transport.reconnect).not.toHaveBeenCalled();
+    expect(transport.close).toHaveBeenCalledOnce();
+  });
+
+  it("detaches listeners before closing so cleanup cannot trigger a session retry", () => {
+    const { connect, client, store, events, transport, onNotification } =
+      setup();
+    const controller = new AbortController();
+    connect(controller.signal);
+    controller.abort();
+    events.dispatchEvent(
+      new MessageEvent("message", {
+        data: JSON.stringify({ data: { op: "reconnected" } }),
+      }),
+    );
+    expect(transport.reconnect).not.toHaveBeenCalled();
+    expect(store.get(client.connectionAtom)).toEqual({
+      state: ConnectionState.NOT_STARTED,
+    });
+    expect(onNotification).not.toHaveBeenCalled();
+  });
+
+  it("can reconnect after effect cleanup without reviving the earlier client attempt", async () => {
+    const { connect, healthy, transport, events, onNotification } = setup();
+    const first = new AbortController();
+    connect(first.signal);
+    first.abort();
+    const second = new AbortController();
+    connect(second.signal);
+    healthy.resolve();
+    await healthy.promise;
+    expect(transport.reconnect).toHaveBeenCalledOnce();
+    events.dispatchEvent(
+      new MessageEvent("message", {
+        data: JSON.stringify({
+          event: "marimo",
+          data: { marimo: { op: "reconnected" } },
+        }),
+      }),
+    );
+    expect(onNotification).toHaveBeenCalledExactlyOnceWith({
+      data: { op: "reconnected" },
+    });
+    second.abort();
+  });
+});

package/src/core/host/__tests__/protocol.test.ts

@@ -0,0 +1,214 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import NodeFs from "node:fs";
+import NodeUrl from "node:url";
+import { describe, expect, expectTypeOf, it } from "vitest";
+import * as Protocol from "../protocol";
+
+const fixtures = new NodeUrl.URL(
+  "../../../../../protocols/host/fixtures/",
+  import.meta.url,
+);
+
+function readFixture(name: string): unknown {
+  return JSON.parse(
+    NodeFs.readFileSync(new NodeUrl.URL(name, fixtures), "utf8"),
+  );
+}
+
+const records = {
+  "record.json": Protocol.HostRecord,
+  "cell.json": Protocol.Cell,
+  "analysis.json": Protocol.Analysis,
+  "diagnostic.json": Protocol.Error,
+  "edit.json": Protocol.EditRequest,
+  "edit-conflict-request.json": Protocol.EditRequest,
+  "edit-response.json": Protocol.EditResponse,
+  "execution-restored.json": Protocol.Execution,
+  "run-cells-request.json": Protocol.RunRequest,
+};
+
+describe("host protocol fixtures", () => {
+  it.each(Object.entries(records))("parses %s unchanged", (name, schema) => {
+    const input = readFixture(name);
+    const original = structuredClone(input);
+    expect(schema.parse(input)).toEqual(original);
+    expect(input).toEqual(original);
+  });
+
+  it.each(NodeFs.readdirSync(new NodeUrl.URL("problems/", fixtures)))(
+    "parses problem %s",
+    (name) => {
+      const input = readFixture(`problems/${name}`);
+      expect(Protocol.Problem.parse(input)).toEqual(input);
+    },
+  );
+
+  it.each(NodeFs.readdirSync(fixtures).filter((name) => name.endsWith(".sse")))(
+    "parses events in %s unchanged",
+    (name) => {
+      const stream = NodeFs.readFileSync(
+        new NodeUrl.URL(name, fixtures),
+        "utf8",
+      );
+      let count = 0;
+      for (const frame of stream.split(/\r?\n\r?\n/)) {
+        const event = /^event: (.+)$/m.exec(frame)?.[1];
+        if (!event && !/^data:/m.test(frame)) {
+          continue;
+        }
+        const data = frame
+          .split(/\r?\n/)
+          .filter((line) => line.startsWith("data:"))
+          .map((line) => line.slice(5).trimStart())
+          .join("\n");
+        const id = /^id: (.+)$/m.exec(frame)?.[1];
+        const payload: unknown = JSON.parse(data);
+        const input = { event, data: payload, ...(id ? { id } : {}) };
+        const original = structuredClone(input);
+        expect(Protocol.Event.parse(input)).toEqual(original);
+        expect(input).toEqual(original);
+        count += 1;
+      }
+      expect(count).toBeGreaterThan(0);
+    },
+  );
+});
+
+describe("host protocol bindings", () => {
+  it("preserves native marimo payloads without interpreting them", () => {
+    expectTypeOf<Protocol.MarimoEvent["data"]["marimo"]>().toBeUnknown();
+    const input = {
+      event: "marimo",
+      id: "42",
+      data: {
+        marimo: {
+          future_payload: { value: Number.NaN },
+        },
+      },
+    };
+    expect(Protocol.Event.parse(input)).toEqual(input);
+  });
+
+  it("accepts additive fields and preserves arbitrary notebook settings", () => {
+    const cell = Protocol.Cell.parse(readFixture("cell.json"));
+    const input = {
+      ...cell,
+      future: { enabled: true },
+      config: { ...cell.config, future: 42 },
+    };
+    const original = structuredClone(input);
+    expect(Protocol.Cell.parse(input)).toEqual(cell);
+    expect(input).toEqual(original);
+
+    const { document } = Protocol.EditResponse.parse(
+      readFixture("edit-response.json"),
+    );
+    const configured = {
+      ...document,
+      app_config: { extension: { values: [1, null, false] } },
+    };
+    expect(Protocol.Document.parse(configured)).toEqual(configured);
+  });
+
+  it("applies the defaults declared for cell inputs", () => {
+    expect(
+      Protocol.CellInput.parse({ id: "c3", code: "x = 1", config: {} }),
+    ).toEqual({
+      id: "c3",
+      code: "x = 1",
+      name: "_",
+      config: {
+        column: null,
+        disabled: false,
+        hide_code: false,
+        expand_output: false,
+      },
+    });
+  });
+
+  it.each(["run_id", "runtime_id", "caused_by"] as const)(
+    "requires the nullable execution field %s",
+    (field) => {
+      const execution = Protocol.Execution.parse(
+        readFixture("execution-restored.json"),
+      );
+      const missing = Object.fromEntries(
+        Object.entries(execution).filter(([key]) => key !== field),
+      );
+      expect(Protocol.Execution.safeParse(missing).success).toBe(false);
+      expect(
+        Protocol.Execution.safeParse({ ...execution, [field]: null }).success,
+      ).toBe(true);
+    },
+  );
+
+  it("preserves snapshot metadata omissions and validates live timestamps", () => {
+    const cell = Protocol.Cell.parse(readFixture("cell.json"));
+    expect(Protocol.CellMessage.parse({ cell })).toEqual({ cell });
+    const message = { cell, time: "2026-01-01T12:00:00+05:30" };
+    expect(Protocol.CellMessage.parse(message)).toEqual(message);
+    for (const time of [
+      null,
+      "yesterday",
+      "2026-02-30T12:00:00Z",
+      "2026-01-01T12:00:00",
+    ]) {
+      expect(Protocol.CellMessage.safeParse({ cell, time }).success).toBe(
+        false,
+      );
+    }
+  });
+
+  it("rejects invalid cell fields without coercing their values", () => {
+    const cell = Protocol.Cell.parse(readFixture("cell.json"));
+    for (const patch of [
+      { id: "" },
+      { code: 42 },
+      { etag: 42 },
+      { version: "2" },
+      { version: 0 },
+      { version: 1.5 },
+      { config: { ...cell.config, disabled: "false" } },
+    ]) {
+      expect(Protocol.Cell.safeParse({ ...cell, ...patch }).success).toBe(
+        false,
+      );
+    }
+    const result = Protocol.Cell.safeParse({ ...cell, code: 42 });
+    expect(result.success).toBe(false);
+    if (result.success) {
+      throw new Error("Expected an invalid cell");
+    }
+    expect(result.error.issues).toEqual([
+      expect.objectContaining({ path: ["code"], code: "invalid_type" }),
+    ]);
+  });
+
+  it("requires the fields of the selected change kind", () => {
+    const change = { kind: "delete_cell", id: "c3", base_etag: "*" };
+    expect(Protocol.Change.parse(change)).toEqual(change);
+    expect(
+      Protocol.Change.safeParse({ kind: "delete_cell", id: "c3" }).success,
+    ).toBe(false);
+    expect(
+      Protocol.Change.safeParse({ ...change, kind: "unknown" }).success,
+    ).toBe(false);
+  });
+
+  it("requires custom payloads while allowing any JSON value", () => {
+    for (const data of [
+      null,
+      false,
+      42,
+      "text",
+      { nested: [1, null, false] },
+    ]) {
+      const custom = { type: "example.event", data };
+      expect(Protocol.Custom.parse(custom)).toEqual(custom);
+    }
+    expect(Protocol.Custom.safeParse({ type: "example.event" }).success).toBe(
+      false,
+    );
+  });
+});

package/src/core/host/connection.ts

@@ -0,0 +1,81 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+export const ConnectionState = {
+  NOT_STARTED: "NOT_STARTED",
+  CONNECTING: "CONNECTING",
+  OPEN: "OPEN",
+  CLOSING: "CLOSING",
+  CLOSED: "CLOSED",
+} as const;
+
+export type ConnectionState =
+  (typeof ConnectionState)[keyof typeof ConnectionState];
+
+export const ConnectionClosedReason = {
+  KERNEL_DISCONNECTED: "KERNEL_DISCONNECTED",
+  KERNEL_STARTUP_ERROR: "KERNEL_STARTUP_ERROR",
+} as const;
+
+export type ConnectionClosedReason =
+  (typeof ConnectionClosedReason)[keyof typeof ConnectionClosedReason];
+
+export type ConnectionStatus =
+  | {
+      state: typeof ConnectionState.CLOSED;
+      code: ConnectionClosedReason;
+      reason: string;
+    }
+  | {
+      state: typeof ConnectionState.CONNECTING;
+    }
+  | {
+      state:
+        | typeof ConnectionState.OPEN
+        | typeof ConnectionState.CLOSING
+        | typeof ConnectionState.NOT_STARTED;
+    };
+
+/** The recovery policy for a closed connection. */
+export interface CloseDecision {
+  kind: "terminal" | "gave-up" | "retry";
+  status: ConnectionStatus;
+}
+
+export type ConnectionEvent = "open" | "close" | "message" | "error";
+
+export interface IConnectionTransportMap {
+  open: WebSocketEventMap["open"];
+  message: WebSocketEventMap["message"];
+  close: WebSocketEventMap["close"];
+  error: WebSocketEventMap["error"];
+}
+
+export type ConnectionTransportCallback<T extends ConnectionEvent> = (
+  data: IConnectionTransportMap[T],
+) => void;
+
+/** Transport readiness, using the WebSocket readyState values. */
+export const TransportReadyState = {
+  CONNECTING: 0,
+  OPEN: 1,
+  CLOSING: 2,
+  CLOSED: 3,
+} as const;
+
+export type TransportReadyState =
+  (typeof TransportReadyState)[keyof typeof TransportReadyState];
+
+/** A reconnectable source of notebook messages. */
+export interface NotebookTransport {
+  reconnect(code?: number | undefined, reason?: string | undefined): void;
+  close(): void;
+  addEventListener<T extends ConnectionEvent>(
+    event: T,
+    callback: ConnectionTransportCallback<T>,
+  ): void;
+  removeEventListener<T extends ConnectionEvent>(
+    event: T,
+    callback: ConnectionTransportCallback<T>,
+  ): void;
+  readyState: TransportReadyState;
+}

package/src/core/host/events.ts

@@ -0,0 +1,60 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { z } from "zod";
+import * as Protocol from "./protocol";
+
+/**
+ * Notebook stream events, including native marimo notifications.
+ *
+ * Known events follow the generated protocol. Unknown events retain their name
+ * and cursor so clients can skip them without losing their place in the stream.
+ */
+
+/** Event kinds requested by the client; removals and controls are implicit. */
+export const NOTEBOOK_EVENT_TYPES = [
+  "notebook",
+  "document",
+  "cell",
+  "marimo",
+] as const satisfies readonly (Protocol.ObjectType | "marimo")[];
+
+const envelope = z.object({
+  event: z.string().min(1),
+  data: z.unknown(),
+  id: z.string().min(1).optional(),
+});
+
+const knownNames = new Set<string>(
+  Protocol.Event.options.flatMap((schema) => schema.shape.event.options),
+);
+
+/** A decoded event, or the name and cursor of an unrecognized event. */
+export type NotebookEvent =
+  | Protocol.EventOutput
+  | {
+      readonly event: "unknown";
+      readonly name: string;
+      readonly id?: string;
+    };
+
+const sessionEnvelope = z.object({ data: Protocol.MarimoNotification });
+
+/**
+ * Decode a host event or a legacy `{ data: notification }` session message.
+ *
+ * Throws if the envelope or a known event is malformed. Native notification
+ * payloads are passed through for the session handler to interpret.
+ */
+export function parseNotebookEvent(value: unknown): NotebookEvent {
+  if (typeof value === "object" && value !== null && !("event" in value)) {
+    return {
+      event: "marimo",
+      data: { marimo: sessionEnvelope.parse(value).data },
+    };
+  }
+  const parsed = envelope.parse(value);
+  if (knownNames.has(parsed.event)) {
+    return Protocol.Event.parse(parsed);
+  }
+  return { event: "unknown", name: parsed.event, id: parsed.id };
+}