@marimo-team/islands 0.25.3-dev34 → 0.25.3-dev349

package/src/hooks/useScript.ts

@@ -1,90 +1,87 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-import React from "react";
+import { useCallback, useSyncExternalStore } from "react";
 
 type ScriptStatus = "loading" | "ready" | "error" | "unknown";
 interface ScriptOptions {
   removeOnUnmount?: boolean;
 }
 
+interface ScriptEntry {
+  element: HTMLScriptElement;
+  listeners: Set<() => void>;
+  /** Whether useScript created the element, and so may remove it. */
+  owned: boolean;
+}
+
 const defaultOptions = { removeOnUnmount: false };
 
+// One entry per src, so reading the status on render is a map lookup
+const scripts = new Map<string, ScriptEntry>();
+
+function findScript(src: string): HTMLScriptElement | undefined {
+  // Compare attributes rather than building a selector, which `src` could break
+  return [...document.scripts].find(
+    (script) => script.getAttribute("src") === src,
+  );
+}
+
+function acquireScript(src: string): ScriptEntry {
+  const cached = scripts.get(src);
+  if (cached?.element.isConnected) {
+    return cached;
+  }
+  const existing = findScript(src);
+  const element = existing ?? document.createElement("script");
+  const entry: ScriptEntry = {
+    element,
+    listeners: new Set(),
+    owned: !existing,
+  };
+  const setStatus = (status: ScriptStatus) => {
+    element.dataset.status = status;
+    entry.listeners.forEach((listener) => listener());
+  };
+  // Listen for the script's lifetime rather than per subscriber, so the
+  // status still settles if the first subscriber unmounts mid-load.
+  element.addEventListener("load", () => setStatus("ready"), { once: true });
+  element.addEventListener("error", () => setStatus("error"), { once: true });
+  if (!existing) {
+    element.src = src;
+    element.async = true;
+    element.dataset.status = "loading";
+    document.body.append(element);
+  }
+  scripts.set(src, entry);
+  return entry;
+}
+
 export function useScript(
   src: string,
   options: ScriptOptions = defaultOptions,
 ) {
-  const [status, setStatus] = React.useState<ScriptStatus>(() => {
-    const existingScript = document.querySelector<HTMLScriptElement>(
-      `script[src="${src}"]`,
-    );
-    return (existingScript?.dataset.status as ScriptStatus) || "unknown";
-  });
-
-  React.useEffect(() => {
-    let script = document.querySelector<HTMLScriptElement>(
-      `script[src="${src}"]`,
-    );
-
-    const updateStatus = (newStatus: ScriptStatus) => {
-      if (script) {
-        script.dataset.status = newStatus;
-        setStatus(newStatus);
-      }
-    };
-
-    const handleScriptLoad = () => updateStatus("ready");
-    const handleScriptError = () => updateStatus("error");
-
-    // observes DOM changes and updates the status
-    const observeScript = (script: HTMLScriptElement) => {
-      const observer = new MutationObserver(() => {
-        const newStatus = script.dataset.status;
-        if (newStatus) {
-          setStatus(newStatus as ScriptStatus);
-        }
-      });
-      observer.observe(script, {
-        attributes: true,
-        attributeFilter: ["data-status"],
-      });
-      return observer;
-    };
-
-    if (!script) {
-      script = document.createElement("script");
-      script.src = src;
-      script.async = true;
-      updateStatus("loading");
-      document.body.append(script);
-
-      script.addEventListener("load", handleScriptLoad);
-      script.addEventListener("error", handleScriptError);
-
-      const observer = observeScript(script);
+  const { removeOnUnmount } = options;
+  const getSnapshot = useCallback((): ScriptStatus => {
+    const cached = scripts.get(src)?.element;
+    const element = cached?.isConnected ? cached : findScript(src);
+    return (element?.dataset.status as ScriptStatus) || "unknown";
+  }, [src]);
 
+  const subscribe = useCallback(
+    (onChange: () => void) => {
+      const entry = acquireScript(src);
+      entry.listeners.add(onChange);
       return () => {
-        observer.disconnect();
-        if (script) {
-          script.removeEventListener("load", handleScriptLoad);
-          script.removeEventListener("error", handleScriptError);
-
-          if (options.removeOnUnmount) {
-            script.remove();
+        entry.listeners.delete(onChange);
+        if (removeOnUnmount && entry.owned && entry.listeners.size === 0) {
+          entry.element.remove();
+          if (scripts.get(src) === entry) {
+            scripts.delete(src);
           }
         }
       };
-    }
-
-    const domStatus = script.dataset.status;
-    if (domStatus) {
-      if (domStatus === "ready" || domStatus === "error") {
-        return;
-      }
-      const observer = observeScript(script);
-      return () => observer.disconnect();
-    }
-
-    setStatus("unknown");
-  }, [src, options.removeOnUnmount]);
+    },
+    [src, removeOnUnmount],
+  );
 
-  return status;
+  return useSyncExternalStore(subscribe, getSnapshot);
 }

package/src/hooks/useStateForKey.ts

@@ -0,0 +1,93 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { type SetStateAction, useCallback, useReducer } from "react";
+
+interface KeyedState<T> {
+  key: unknown;
+  value: T;
+}
+
+type KeyedAction<T> =
+  | { type: "reset"; state: KeyedState<T> }
+  | { type: "update"; key: unknown; update: SetStateAction<T> };
+
+function reduceKeyedState<T>(
+  previous: KeyedState<T>,
+  action: KeyedAction<T>,
+): KeyedState<T> {
+  if (action.type === "reset") {
+    return action.state;
+  }
+  if (!Object.is(previous.key, action.key)) {
+    return previous;
+  }
+  const value =
+    typeof action.update === "function"
+      ? (action.update as (previous: T) => T)(previous.value)
+      : action.update;
+  return Object.is(previous.value, value)
+    ? previous
+    : { key: action.key, value };
+}
+
+/**
+ * Local state that resets when its owner key changes.
+ *
+ * The key uses Object.is. NaN matches NaN, 0 differs from -0, and objects match
+ * by reference. Keep the key stable between resets. An inline object such as
+ * `{ id }` creates a new key on each render and causes a render loop.
+ *
+ * On a key change, this hook calls initialize() and schedules a state update
+ * during render. React retries the calling component before rendering its
+ * children. Children and effects receive the reset value without an intervening
+ * commit of the old value. This resets only this hook's state. The component
+ * stays mounted, and its other hooks keep their state.
+ *
+ * initialize must be pure. React can call it more than once in Strict Mode or
+ * when retrying an interrupted render. A new initializer alone does not reset
+ * state. The next key change uses the initializer from that render. After
+ * committed key changes A -> B -> A, A starts fresh instead of restoring its
+ * earlier value. If React abandons a pending render for B, the visible state
+ * for A remains intact, including edits made while B was pending.
+ *
+ * The setter accepts values or pure updater functions, like useState. To store
+ * a function, use setValue(() => callback). Its identity is stable for the same
+ * key and changes with the key. Updates from a different key are ignored,
+ * including their updater functions. After A -> B -> A, an old setter for A
+ * can update A again. This hook does not cancel requests. For request results,
+ * use a distinct key for each lifetime so late results stay invalid.
+ *
+ * React can still call the component for ignored or same-value updates.
+ * Render counts do not indicate how many changes React commits to the screen.
+ *
+ * Use useInternalStateWithSync when the external value itself supplies the
+ * reset value. Use this hook when a separate key controls the reset.
+ *
+ * @example
+ * const [expanded, setExpanded] = useStateForKey(tree, () => new Set<string>());
+ *
+ * @see https://react.dev/reference/react/useState#storing-information-from-previous-renders
+ */
+export function useStateForKey<T>(key: unknown, initialize: () => T) {
+  // useState can eagerly evaluate an updater against a suspended render's
+  // state. A reducer checks the key when React processes each render's queue,
+  // so the still-visible owner can accept edits during a transition.
+  const [stored, dispatch] = useReducer(reduceKeyedState<T>, key, (key) => ({
+    key,
+    value: initialize(),
+  }));
+  const current = Object.is(stored.key, key)
+    ? stored
+    : { key, value: initialize() };
+  if (current !== stored) {
+    dispatch({ type: "reset", state: current });
+  }
+
+  const setValue = useCallback(
+    (update: SetStateAction<T>) => {
+      dispatch({ type: "update", key, update });
+    },
+    [key],
+  );
+
+  return [current.value, setValue] as const;
+}

package/src/mount.tsx

@@ -5,8 +5,8 @@ import { Provider } from "jotai";
 import { createRoot } from "react-dom/client";
 import { z } from "zod";
 import {
-  appConfigAtom,
   configOverridesAtom,
+  pageAppConfigAtom,
   userConfigAtom,
 } from "@/core/config/config";
 import { KnownQueryParams } from "@/core/constants";
@@ -21,7 +21,7 @@ import {
 import { Logger } from "@/utils/Logger";
 import { ErrorBoundary } from "./components/editor/boundary/ErrorBoundary";
 import { notebookAtom } from "./core/cells/cells";
-import { notebookStateFromSession } from "./core/cells/session";
+import { readExport } from "./core/cells/session";
 import {
   parseAppConfig,
   parseConfigOverrides,
@@ -40,11 +40,11 @@ import {
   viewStateAtom,
 } from "./core/mode";
 import { cleanupAuthQueryParams } from "./core/network/auth";
-import { connectionAtom } from "./core/network/connection";
 import { requestClientAtom } from "./core/network/requests";
 import { resolveRequestClient } from "./core/network/resolve";
 import {
   DEFAULT_RUNTIME_CONFIG,
+  getRuntimeManager,
   runtimeConfigAtom,
 } from "./core/runtime/config";
 import {
@@ -62,7 +62,7 @@ import {
 import { maybeRegisterVSCodeBindings } from "./core/vscode/vscode-bindings";
 import type { FileStore } from "./core/wasm/store";
 import { notebookFileStore } from "./core/wasm/store";
-import { ConnectionState } from "./core/websocket/types";
+import { embeddedSnapshotAtom } from "./core/session/embedded";
 import {
   handleWidgetMessage,
   WIDGET_REGISTRY,
@@ -274,8 +274,6 @@ const mountOptionsSchema = z.object({
       .transform((val) => val as api.Notebook["NotebookV1"]),
   ]),
 
-  notebookId: z.string().optional(),
-
   /**
    * Runtime configs
    */
@@ -283,10 +281,7 @@ const mountOptionsSchema = z.object({
     .array(
       z.looseObject({
         url: z.string(),
-        // Lazy by default, but can be overridden by the runtime config
-        lazy: z.boolean().default(true),
         authToken: z.string().nullish(),
-        notebookId: z.string().optional(),
       }),
     )
     .nullish()
@@ -349,7 +344,8 @@ function initStore(options: unknown) {
     parseConfigOverrides(parsedOptions.data.configOverrides),
   );
   store.set(userConfigAtom, parseUserConfig(parsedOptions.data.config));
-  store.set(appConfigAtom, parseAppConfig(parsedOptions.data.appConfig));
+  store.set(pageAppConfigAtom, parsedOptions.data.appConfig);
+  const appConfig = parseAppConfig(parsedOptions.data.appConfig);
   store.set(pairConfigAtom, parsedOptions.data.pairConfig);
 
   // Runtime config
@@ -360,27 +356,30 @@ function initStore(options: unknown) {
       ...firstRuntimeConfig,
       serverToken: parsedOptions.data.serverToken,
     });
-    // If the remote runtime is not lazy, start it in CONNECTING
-    if (!firstRuntimeConfig.lazy && !isStaticNotebook()) {
-      store.set(connectionAtom, { state: ConnectionState.CONNECTING });
-    }
   } else {
     store.set(runtimeConfigAtom, {
       ...DEFAULT_RUNTIME_CONFIG,
-      notebookId: parsedOptions.data.notebookId,
-      lazy: false,
       serverToken: parsedOptions.data.serverToken,
     });
   }
+  // The page's runtime is known from here on, so look for it now.
+  void getRuntimeManager().init();
 
-  // Session/notebook
-  const notebook = notebookStateFromSession(
-    parsedOptions.data.session,
-    parsedOptions.data.notebook,
-  );
+  // The notebook the page was served with shows until the host answers.
+  const { snapshot, cellRuntime } = readExport({
+    session: parsedOptions.data.session,
+    notebook: parsedOptions.data.notebook,
+    filename: parsedOptions.data.filename,
+    appConfig,
+  });
+  store.set(embeddedSnapshotAtom, snapshot);
   store.set(layoutStateAtom, initialLayoutState(parsedOptions.data.layout));
-  if (notebook) {
-    store.set(notebookAtom, notebook);
+  if (snapshot !== null) {
+    // The scratch cell keeps its own runtime; the snapshot has none for it.
+    store.set(notebookAtom, (state) => ({
+      ...state,
+      cellRuntime: { ...state.cellRuntime, ...cellRuntime },
+    }));
   }
 }
 

package/src/plugins/__tests__/plugin-object-schemas.test.ts

@@ -0,0 +1,104 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { describe, expect, it } from "vitest";
+import { DataEditorPlugin } from "../impl/DataEditorPlugin";
+import { DataTablePlugin } from "../impl/DataTablePlugin";
+import { DataFramePlugin } from "../impl/data-frames/DataFramePlugin";
+import { PlotlyPlugin } from "../impl/plotly/PlotlyPlugin";
+import { VegaPlugin } from "../impl/vega/VegaPlugin";
+
+describe("plugin object schemas", () => {
+  const rows = [{ value: 1, extra: { nested: true } }];
+  const styles = { "0": { value: { color: "red", customProperty: 1 } } };
+  const tableData = {
+    initialValue: [],
+    label: null,
+    data: rows,
+    totalRows: 1,
+    totalColumns: 1,
+    rowHeaders: [],
+    cellStyles: styles,
+  };
+
+  it("preserves table rows and arbitrary cell style fields", () => {
+    expect(DataTablePlugin.validator.parse(tableData)).toMatchObject({
+      data: rows,
+      cellStyles: styles,
+      pageSize: 10,
+    });
+  });
+
+  it("preserves data editor row fields", () => {
+    expect(
+      DataEditorPlugin.validator.parse({
+        initialValue: { edits: [] },
+        label: null,
+        data: rows,
+        editableColumns: "all",
+      }),
+    ).toMatchObject({ data: rows, columnNames: [] });
+  });
+
+  it("preserves table RPC payload fields", () => {
+    const functions = DataTablePlugin.functions!;
+    const searchResult = {
+      data: rows,
+      total_rows: 1,
+      cell_styles: styles,
+      cell_hover_texts: null,
+    };
+    expect(functions.search.output.parse(searchResult)).toEqual(searchResult);
+    expect(functions.get_row_ids!.input.parse({ custom: true })).toEqual({
+      custom: true,
+    });
+    expect(functions.get_data_url!.input.parse({ custom: true })).toEqual({
+      custom: true,
+    });
+    expect(
+      functions.get_data_url!.output.parse({ data_url: rows, format: "json" }),
+    ).toEqual({ data_url: rows, format: "json" });
+  });
+
+  it("preserves dataframe search row fields", () => {
+    const result = { data: rows, total_rows: 1 };
+    expect(DataFramePlugin.functions!.search.output.parse(result)).toEqual(
+      result,
+    );
+  });
+
+  it("preserves Plotly figure and config fields", () => {
+    const data = {
+      figure: { data: [], layout: { title: "Chart" } },
+      config: { responsive: true },
+    };
+    expect(new PlotlyPlugin().validator.parse(data)).toEqual(data);
+  });
+
+  it("preserves Vega spec and embed options fields and defaults", () => {
+    const spec = { mark: "bar", encoding: { x: { field: "value" } } };
+    const validator = new VegaPlugin().validator;
+    expect(validator.parse({ spec })).toEqual({
+      spec,
+      chartSelection: true,
+      fieldSelection: true,
+      embedOptions: {},
+    });
+    expect(
+      validator.parse({ spec, embedOptions: { actions: false } }).embedOptions,
+    ).toEqual({ actions: false });
+  });
+
+  it("rejects scalar values where plugin objects are required", () => {
+    expect(
+      DataTablePlugin.validator.safeParse({ ...tableData, data: [42] }).success,
+    ).toBe(false);
+    expect(
+      new PlotlyPlugin().validator.safeParse({ figure: 42, config: {} })
+        .success,
+    ).toBe(false);
+    expect(
+      new VegaPlugin().validator.safeParse({ spec: {}, embedOptions: 42 })
+        .success,
+    ).toBe(false);
+  });
+});

package/src/plugins/core/RenderHTML.tsx

@@ -300,11 +300,9 @@ const CopyableCode = ({ children }: { children: ReactNode }) => {
           tooltip={false}
           className="p-1"
           value={() => {
-            const codeElement = ref.current?.firstChild;
-            if (codeElement) {
-              return codeElement.textContent || "";
-            }
-            return "";
+            const codeElement =
+              ref.current?.querySelector("pre code") ?? ref.current?.firstChild;
+            return codeElement?.textContent ?? "";
           }}
         />
       </div>

package/src/plugins/core/__test__/PluginSlot.test.tsx

@@ -0,0 +1,202 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+import { act, render } from "@testing-library/react";
+import { createRef, StrictMode } from "react";
+import { afterEach, describe, expect, test } from "vitest";
+import { z } from "zod";
+import {
+  MarimoValueInputEvent,
+  MarimoValueUpdateEvent,
+} from "@/core/dom/events";
+import type { IPlugin, IPluginProps, Setter } from "../../types";
+import { PluginSlot, type PluginSlotHandle } from "../registerReactComponent";
+
+type Value = number;
+
+interface Mounted {
+  host: HTMLElement;
+  handle: React.RefObject<PluginSlotHandle | null>;
+  events: unknown[];
+  setValue: () => Setter<Value>;
+  value: () => Value;
+  unmount: () => void;
+}
+
+function mount(options: { strict?: boolean; initial?: Value } = {}): Mounted {
+  const host = document.createElement("div");
+  document.body.append(host);
+  const events: unknown[] = [];
+  host.addEventListener(MarimoValueInputEvent.TYPE, (e) => {
+    events.push((e as CustomEvent<{ value: unknown }>).detail.value);
+  });
+
+  let latestProps: IPluginProps<Value, unknown> | undefined;
+  const plugin: IPlugin<Value, unknown> = {
+    tagName: "marimo-test-slot",
+    validator: z.any(),
+    render: (props) => {
+      latestProps = props;
+      return <span data-testid="value">{props.value}</span>;
+    },
+  };
+
+  const handle = createRef<PluginSlotHandle>();
+  const tree = (
+    <PluginSlot
+      hostElement={host}
+      plugin={plugin}
+      ref={handle}
+      getInitialValue={() => options.initial ?? 0}
+    />
+  );
+  const result = render(
+    options.strict ? <StrictMode>{tree}</StrictMode> : tree,
+  );
+
+  return {
+    host,
+    handle,
+    events,
+    setValue: () => {
+      if (!latestProps) {
+        throw new Error("plugin did not render");
+      }
+      return latestProps.setValue;
+    },
+    value: () => {
+      if (!latestProps) {
+        throw new Error("plugin did not render");
+      }
+      return latestProps.value;
+    },
+    unmount: result.unmount,
+  };
+}
+
+function pushHostValue(host: HTMLElement, value: Value) {
+  host.dispatchEvent(
+    MarimoValueUpdateEvent.create({ detail: { value, element: host } }),
+  );
+}
+
+describe.each([
+  { strict: false, label: "default" },
+  { strict: true, label: "StrictMode" },
+])("PluginSlot value dispatch ($label)", ({ strict }) => {
+  afterEach(() => {
+    document.body.innerHTML = "";
+  });
+
+  test("a plain value sends exactly one input event", () => {
+    const m = mount({ strict });
+    act(() => m.setValue()(5));
+    expect(m.events).toEqual([5]);
+    expect(m.value()).toBe(5);
+  });
+
+  test("an equal value sends no input event", () => {
+    const m = mount({ strict, initial: 5 });
+    act(() => m.setValue()(5));
+    expect(m.events).toEqual([]);
+  });
+
+  test("functional updates chain within one tick and send one event each", () => {
+    const m = mount({ strict });
+    act(() => {
+      m.setValue()((v) => v + 1);
+      m.setValue()((v) => v + 1);
+    });
+    expect(m.events).toEqual([1, 2]);
+    expect(m.value()).toBe(2);
+  });
+
+  test("a host update changes the value without an input event", () => {
+    const m = mount({ strict });
+    act(() => pushHostValue(m.host, 10));
+    expect(m.value()).toBe(10);
+    expect(m.events).toEqual([]);
+
+    act(() => m.setValue()((v) => v + 1));
+    expect(m.events).toEqual([11]);
+  });
+
+  test("reset restores the initial value without an input event", () => {
+    const m = mount({ strict, initial: 3 });
+    act(() => m.setValue()(7));
+    expect(m.events).toEqual([7]);
+
+    act(() => m.handle.current?.reset());
+    expect(m.value()).toBe(3);
+    expect(m.events).toEqual([7]);
+
+    act(() => m.setValue()((v) => v + 1));
+    expect(m.events).toEqual([7, 4]);
+  });
+
+  test("equal host and plugin values send nothing, then a real change sends one event", () => {
+    const m = mount({ strict, initial: 2 });
+    act(() => pushHostValue(m.host, 2));
+    act(() => m.setValue()(2));
+    expect(m.events).toEqual([]);
+
+    act(() => m.setValue()(3));
+    expect(m.events).toEqual([3]);
+  });
+});
+
+describe("PluginSlot value dispatch (shared behavior)", () => {
+  afterEach(() => {
+    document.body.innerHTML = "";
+  });
+
+  test("the input event bubbles to document and names the host element", () => {
+    const m = mount();
+    const seen: Array<{ value: unknown; element: HTMLElement }> = [];
+    const listener = (e: Event) => {
+      seen.push(
+        (e as CustomEvent<{ value: unknown; element: HTMLElement }>).detail,
+      );
+    };
+    document.addEventListener(MarimoValueInputEvent.TYPE, listener);
+    act(() => m.setValue()(9));
+    document.removeEventListener(MarimoValueInputEvent.TYPE, listener);
+    expect(seen).toEqual([{ value: 9, element: m.host }]);
+  });
+
+  test("a shallow-equal object skips the event but adopts the new reference", () => {
+    const host = document.createElement("div");
+    document.body.append(host);
+    const events: unknown[] = [];
+    host.addEventListener(MarimoValueInputEvent.TYPE, (e) => {
+      events.push((e as CustomEvent<{ value: unknown }>).detail.value);
+    });
+    let setValue: Setter<{ a: number; b: string }> | undefined;
+    let committed: { a: number; b: string } | undefined;
+    const plugin: IPlugin<{ a: number; b: string }, unknown> = {
+      tagName: "marimo-test-object",
+      validator: z.any(),
+      render: (props) => {
+        setValue = props.setValue;
+        committed = props.value;
+        return <span />;
+      },
+    };
+    render(
+      <PluginSlot
+        hostElement={host}
+        plugin={plugin}
+        getInitialValue={() => ({ a: 1, b: "x" })}
+      />,
+    );
+
+    // A form resubmits an equal value with a fresh reference. The slot must
+    // adopt that reference so the form's identity check reads as saved.
+    const resubmitted = { a: 1, b: "x" };
+    act(() => setValue?.(resubmitted));
+    expect(events).toEqual([]);
+    expect(committed).toBe(resubmitted);
+
+    act(() => setValue?.((prev) => ({ ...prev, a: 2 })));
+    expect(events).toEqual([{ a: 2, b: "x" }]);
+  });
+});