@marimo-team/islands 0.23.14-dev6 → 0.23.14-dev61

package/src/plugins/impl/anywidget/AnyWidgetPlugin.tsx

@@ -1,32 +1,21 @@
 /* Copyright 2026 Marimo. All rights reserved. */
-/* oxlint-disable typescript/no-explicit-any */
 
-import type { AnyWidget } from "@anywidget/types";
-import { useEffect, useRef } from "react";
+import { useEffect, useRef, useState } from "react";
 import { z } from "zod";
-import { useAsyncData } from "@/hooks/useAsyncData";
-import type { HTMLElementNotDerivedFromRef } from "@/hooks/useEventListener";
 import { createPlugin } from "@/plugins/core/builder";
 import type { IPluginProps } from "@/plugins/types";
-import { prettyError } from "@/utils/errors";
-import { Logger } from "@/utils/Logger";
-import { hasFunctionProperty, isRecord } from "@/utils/records";
 import { ErrorBanner } from "../common/error-banner";
-import { getMarimoInternal, MODEL_MANAGER, type Model } from "./model";
-import type { ModelState, WidgetModelId } from "./types";
-import { BINDING_MANAGER, WIDGET_DEF_REGISTRY } from "./widget-binding";
+import { WIDGET_REGISTRY } from "./registry";
+import { isWidgetModelId, type WidgetModelId } from "./types";
 
 /**
- * AnyWidget asset data
+ * The component's only data attribute; everything else arrives through
+ * the model's comm messages, keyed by this id.
  */
 interface Data {
-  jsUrl: string;
-  jsHash: string;
   modelId: WidgetModelId;
 }
 
-type AnyWidgetState = ModelState;
-
 /**
  * Value payload sent by the frontend on state updates.
  *
@@ -37,264 +26,54 @@ interface ModelIdRef {
   model_id?: WidgetModelId;
 }
 
-export function useAnyWidgetModule(opts: { jsUrl: string; jsHash: string }) {
-  const { jsUrl, jsHash } = opts;
-
-  // JS is an ESM file with a render function on it
-  // export function render({ model, el }) {
-  //   ...
-  const {
-    data: jsModule,
-    error,
-    refetch,
-  } = useAsyncData(async () => {
-    return await WIDGET_DEF_REGISTRY.getModule(jsUrl, jsHash);
-    // Re-render on jsHash change (which is a hash of the contents of the file)
-    // instead of a jsUrl change because URLs may change without the contents
-    // actually changing (and we don't want to re-render on every change).
-    // If there is an error loading the URL (e.g. maybe an invalid or old URL),
-    // we also want to re-render.
-  }, [jsHash]);
-
-  // If there is an error and the jsUrl has changed, we want to re-render
-  // because the URL may have changed to a valid URL.
-  const hasError = Boolean(error);
-  useEffect(() => {
-    if (hasError && jsUrl) {
-      WIDGET_DEF_REGISTRY.invalidate(jsHash);
-      refetch();
-    }
-  }, [hasError, jsUrl]);
-
-  return {
-    jsModule,
-    error,
-  };
-}
-
-export function useMountCss(css: string | null | undefined, host: HTMLElement) {
-  // Mount the CSS
-  useEffect(() => {
-    const shadowRoot = host.shadowRoot;
-    if (!css || !shadowRoot) {
-      return;
-    }
-
-    // Try constructed stylesheets first
-    if (
-      "adoptedStyleSheets" in Document.prototype &&
-      "replace" in CSSStyleSheet.prototype
-    ) {
-      const sheet = new CSSStyleSheet();
-      try {
-        sheet.replaceSync(css);
-        if (shadowRoot) {
-          shadowRoot.adoptedStyleSheets = [
-            ...shadowRoot.adoptedStyleSheets,
-            sheet,
-          ];
-        }
-        return () => {
-          if (shadowRoot) {
-            shadowRoot.adoptedStyleSheets =
-              shadowRoot.adoptedStyleSheets.filter((s) => s !== sheet);
-          }
-        };
-      } catch {
-        // Fall through to inline styles if constructed sheets fail
-      }
-    }
-
-    // Fallback to inline styles
-    const style = document.createElement("style");
-    style.innerHTML = css;
-    shadowRoot.append(style);
-    return () => {
-      style.remove();
-    };
-  }, [css, host]);
-}
-
 export const AnyWidgetPlugin = createPlugin<ModelIdRef>("marimo-anywidget")
   .withData(
     z.object({
-      jsUrl: z.string(),
-      jsHash: z.string(),
-      modelId: z.string().transform((v) => v as WidgetModelId),
+      modelId: z.custom<WidgetModelId>(isWidgetModelId, {
+        message: "Expected a non-empty widget model id",
+      }),
     }),
   )
   .withFunctions({})
   .renderer((props) => <AnyWidgetSlot {...props} />);
 
-const AnyWidgetSlot = (props: IPluginProps<ModelIdRef, Data>) => {
-  const { jsUrl, jsHash, modelId } = props.data;
-  const host = props.host as HTMLElementNotDerivedFromRef;
-
-  const { jsModule, error } = useAnyWidgetModule({ jsUrl, jsHash });
-
-  if (error) {
-    return <ErrorBanner error={error} />;
-  }
-
-  if (!jsModule) {
-    return null;
-  }
-
-  if (!isAnyWidgetModule(jsModule)) {
-    return (
-      <ErrorBanner error={getInvalidAnyWidgetModuleError(jsModule, jsUrl)} />
-    );
-  }
-
-  return (
-    <LoadedSlot
-      // Force remount when the widget module or model changes (cell re-run).
-      key={`${jsHash}:${modelId}`}
-      widget={jsModule.default}
-      modelId={modelId}
-      host={host}
-    />
-  );
-};
-
 /**
- * Run the anywidget module
- *
- * Per AFM spec (anywidget.dev/en/afm):
- * - initialize() is called once per model lifetime
- * - render() is called once per view (can be multiple per model)
+ * React adapter for a runtime-owned view. React supplies only the render
+ * target and its lifetime; the runtime owns models, generations, CSS,
+ * composition, and hot reload.
  */
-async function runAnyWidgetModule<T extends AnyWidgetState>(
-  widgetDef: AnyWidget<T>,
-  model: Model<T>,
-  modelId: WidgetModelId,
-  el: HTMLElement,
-  signal: AbortSignal,
-): Promise<void> {
-  // Clear the element, in case the widget is re-rendering
-  el.innerHTML = "";
-
-  try {
-    const binding = BINDING_MANAGER.getOrCreate(modelId);
-    const render = await binding.bind(widgetDef, model);
-    await render(el, signal);
-    // Replay current model values so render listeners observe hydrated state
-    // even if backend updates arrived before listeners were attached.
-    getMarimoInternal(model).reemitState();
-  } catch (error) {
-    Logger.error("Error rendering anywidget", error);
-    el.classList.add("text-error");
-    el.innerHTML = `Error rendering anywidget: ${prettyError(error)}`;
-  }
-}
-
-function isAnyWidgetModule(mod: any): mod is { default: AnyWidget } {
-  if (!mod.default) {
-    return false;
-  }
-
-  return (
-    typeof mod.default === "function" ||
-    typeof mod.default?.render === "function" ||
-    typeof mod.default?.initialize === "function"
-  );
-}
-
-function getInvalidAnyWidgetModuleError(mod: unknown, jsUrl: string): Error {
-  const afmDocs = "https://anywidget.dev/en/afm/";
-  const hasNamedRender = isRecord(mod) && hasFunctionProperty(mod, "render");
-  const hasNamedInitialize =
-    isRecord(mod) && hasFunctionProperty(mod, "initialize");
-
-  if (hasNamedRender || hasNamedInitialize) {
-    const namedExports = [
-      hasNamedRender ? "`render`" : null,
-      hasNamedInitialize ? "`initialize`" : null,
-    ]
-      .filter(Boolean)
-      .join(" and ");
-    const lifecycleHooks = [
-      hasNamedRender ? "render" : null,
-      hasNamedInitialize ? "initialize" : null,
-    ].filter((hook): hook is string => hook !== null);
-    const defaultExportExample = `export default { ${lifecycleHooks.join(", ")} }`;
-    const namedExportExample =
-      lifecycleHooks.length === 1
-        ? `export function ${lifecycleHooks[0]}`
-        : "named export function ...";
-    return new Error(
-      `Anywidget module at ${jsUrl} uses named exports (${namedExports}). ` +
-        "Per the AFM spec, use a default export instead: " +
-        `\`${defaultExportExample}\` (not \`${namedExportExample}\`). ` +
-        `See ${afmDocs}`,
-    );
-  }
-
-  if (!isRecord(mod) || mod.default === undefined) {
-    return new Error(
-      `Anywidget module at ${jsUrl} is missing a default export. ` +
-        "Per the AFM spec, use `export default { render }` or " +
-        "`export default async () => ({ render })`. " +
-        `See ${afmDocs}`,
-    );
-  }
-
-  return new Error(
-    `Anywidget module at ${jsUrl} has an invalid default export. ` +
-      "Expected a factory function or an object with `render` or `initialize`. " +
-      `See ${afmDocs}`,
-  );
-}
-
-interface Props<T extends AnyWidgetState> {
-  widget: AnyWidget<T>;
-  modelId: WidgetModelId;
-  host: HTMLElementNotDerivedFromRef;
-}
-
-const LoadedSlot = <T extends AnyWidgetState>({
-  widget,
-  modelId,
-  host,
-}: Props<T> & { widget: AnyWidget<T> }) => {
+const AnyWidgetSlot = (props: IPluginProps<ModelIdRef, Data>) => {
+  const { modelId } = props.data;
   const htmlRef = useRef<HTMLDivElement>(null);
-
-  // value is already decoded from wire format, may be null if waiting for open message
-  const model = MODEL_MANAGER.getSync(modelId);
-
-  if (!model) {
-    Logger.error("Model not found for modelId", modelId);
-  }
-
-  const css = model?.get("_css");
-  useMountCss(css, host);
+  const [error, setError] = useState<Error>();
 
   useEffect(() => {
-    if (!htmlRef.current || !model) {
+    const el = htmlRef.current;
+    if (!el) {
       return;
     }
     const controller = new AbortController();
-    runAnyWidgetModule(
-      widget,
-      model,
+    setError(undefined);
+    WIDGET_REGISTRY.createView({
       modelId,
-      htmlRef.current,
-      controller.signal,
-    );
+      el,
+      signal: controller.signal,
+    }).catch((error) => {
+      if (!controller.signal.aborted) {
+        setError(error instanceof Error ? error : new Error(String(error)));
+      }
+    });
     return () => controller.abort();
-    // We re-run the widget when the modelId changes, which means the cell
-    // that created the Widget has been re-run.
-    // We need to re-run the widget because it may contain initialization code
-    // that could be reset by the new widget.
-    // See example: https://github.com/marimo-team/marimo/issues/3962#issuecomment-2703184123
-  }, [widget, modelId, model]);
+  }, [modelId]);
 
-  return <div ref={htmlRef} />;
+  return (
+    <>
+      {error ? <ErrorBanner error={error} /> : null}
+      <div ref={htmlRef} />
+    </>
+  );
 };
 
 export const visibleForTesting = {
-  LoadedSlot,
-  runAnyWidgetModule,
-  isAnyWidgetModule,
-  getInvalidAnyWidgetModuleError,
+  AnyWidgetSlot,
 };

package/src/plugins/impl/anywidget/__tests__/AnyWidgetPlugin.test.tsx

@@ -1,221 +1,163 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
 import { render, waitFor } from "@testing-library/react";
-import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
-import { TestUtils } from "@/__tests__/test-helpers";
-import type { HTMLElementNotDerivedFromRef } from "@/hooks/useEventListener";
+import {
+  afterEach,
+  beforeEach,
+  describe,
+  expect,
+  it,
+  type MockInstance,
+  vi,
+} from "vitest";
+import type { IPluginProps } from "@/plugins/types";
 import { visibleForTesting } from "../AnyWidgetPlugin";
-import { MODEL_MANAGER, Model } from "../model";
-import type { WidgetModelId } from "../types";
-import { BINDING_MANAGER } from "../widget-binding";
+import { Model } from "../model";
+import { WIDGET_REGISTRY } from "../registry";
+import type { EsmSpec, ModelState, WidgetModelId } from "../types";
+import { WIDGET_DEF_REGISTRY } from "../widget-binding";
 
-const { LoadedSlot, isAnyWidgetModule, getInvalidAnyWidgetModuleError } =
-  visibleForTesting;
+const { AnyWidgetSlot } = visibleForTesting;
 
 // Helper to create typed model IDs for tests
 const asModelId = (id: string): WidgetModelId => id as WidgetModelId;
 
-// Mock a minimal AnyWidget implementation
-const mockWidget = {
-  initialize: vi.fn(),
-  render: vi.fn(),
-};
-
-describe("LoadedSlot", () => {
-  const modelId = asModelId("test-model-id");
-  let mockModel: Model<{ count: number }>;
-
-  const mockProps = {
-    widget: mockWidget,
-    data: {
-      jsUrl: "http://example.com/widget.js",
-      jsHash: "abc123",
-      modelId: modelId,
-    },
-    host: document.createElement(
-      "div",
-    ) as unknown as HTMLElementNotDerivedFromRef,
-    modelId: modelId,
+function createMockComm() {
+  return {
+    sendUpdate: vi.fn().mockResolvedValue(undefined),
+    sendCustomMessage: vi.fn().mockResolvedValue(undefined),
+  };
+}
+
+function createMockModel(state?: ModelState) {
+  return new Model<ModelState>(state ?? { count: 0 }, createMockComm());
+}
+
+const hostEl = () => document.createElement("div");
+
+describe("AnyWidgetSlot", () => {
+  const SPEC: EsmSpec = { url: "./@file/10-slot.js", hash: "slot-hash" };
+  let getModuleSpy: MockInstance<typeof WIDGET_DEF_REGISTRY.getModule>;
+  let nextId = 0;
+  let modelId: WidgetModelId;
+  let mockWidget: {
+    initialize: ReturnType<typeof vi.fn>;
+    render: ReturnType<typeof vi.fn>;
   };
 
   beforeEach(() => {
-    vi.clearAllMocks();
-    // Create and register a mock model before each test
-    mockModel = new Model(
-      { count: 0 },
-      {
-        sendUpdate: vi.fn().mockResolvedValue(undefined),
-        sendCustomMessage: vi.fn().mockResolvedValue(undefined),
-      },
-    );
-    MODEL_MANAGER.set(modelId, mockModel);
+    nextId += 1;
+    modelId = asModelId(`slot-test-${nextId}`);
+    mockWidget = { initialize: vi.fn(), render: vi.fn() };
+    getModuleSpy = vi
+      .spyOn(WIDGET_DEF_REGISTRY, "getModule")
+      .mockResolvedValue({ default: mockWidget });
+    WIDGET_REGISTRY.setModel(modelId, createMockModel());
+    WIDGET_REGISTRY.setSpec(modelId, {
+      url: SPEC.url,
+      // Unique hash per test so the module-cache spy is always hit.
+      hash: `${SPEC.hash}-${nextId}`,
+    });
   });
 
   afterEach(() => {
-    BINDING_MANAGER.destroy(modelId);
+    WIDGET_REGISTRY.delete(modelId);
+    getModuleSpy.mockRestore();
   });
 
-  it("should render a div with ref", () => {
-    const { container } = render(<LoadedSlot {...mockProps} />);
-    expect(container.querySelector("div")).not.toBeNull();
+  const props = (
+    id: WidgetModelId,
+    value: Record<string, unknown> = {},
+  ): IPluginProps<{ model_id?: WidgetModelId }, { modelId: WidgetModelId }> =>
+    ({
+      data: { modelId: id },
+      value,
+      host: hostEl(),
+      setValue: vi.fn(),
+      functions: {},
+    }) as unknown as IPluginProps<
+      { model_id?: WidgetModelId },
+      { modelId: WidgetModelId }
+    >;
+
+  it("resolves the widget through the registry and renders a view", async () => {
+    render(<AnyWidgetSlot {...props(modelId)} />);
+    await waitFor(() => {
+      expect(mockWidget.initialize).toHaveBeenCalledTimes(1);
+      expect(mockWidget.render).toHaveBeenCalledTimes(1);
+    });
   });
 
-  it("should call runAnyWidgetModule on initialization", async () => {
-    render(<LoadedSlot {...mockProps} />);
-
-    // Wait a render
+  it("aborts the runtime-owned view on unmount", async () => {
+    const { unmount } = render(<AnyWidgetSlot {...props(modelId)} />);
     await waitFor(() => {
-      expect(mockWidget.render).toHaveBeenCalled();
+      expect(mockWidget.render).toHaveBeenCalledTimes(1);
     });
+    const signal = mockWidget.render.mock.calls[0][0].signal as AbortSignal;
+    expect(signal.aborted).toBe(false);
+    unmount();
+    expect(signal.aborted).toBe(true);
   });
 
-  it("should not remount when value update drops model_id", async () => {
-    // Regression: when the frontend sends a state update (e.g. {zoom_level: 0}),
-    // it overwrites the UIElement value that originally held {model_id: "..."}.
-    // The key must stay stable because modelId comes from data, not value.
-    const { container, rerender } = render(<LoadedSlot {...mockProps} />);
-
+  it("does not remount when only the value changes", async () => {
+    // Regression: a state update rewrites the plugin value (dropping
+    // model_id), but the key comes from data attributes, so the view
+    // must survive.
+    const { container, rerender } = render(
+      <AnyWidgetSlot {...props(modelId)} />,
+    );
     await waitFor(() => {
       expect(mockWidget.render).toHaveBeenCalledTimes(1);
     });
-
     const divBefore = container.querySelector("div");
 
-    // Simulate a value update that does NOT include model_id
-    // (this is what happens when the widget sends trait state)
-    rerender(<LoadedSlot {...mockProps} />);
+    rerender(<AnyWidgetSlot {...props(modelId, { zoom_level: 0 })} />);
 
     await waitFor(() => {
-      // The div should be the same DOM node (no remount)
       expect(container.querySelector("div")).toBe(divBefore);
-      // render should not be called again (no remount)
       expect(mockWidget.render).toHaveBeenCalledTimes(1);
     });
   });
 
-  it("should re-run widget when widget prop changes", async () => {
-    const { rerender } = render(<LoadedSlot {...mockProps} />);
-
-    // Wait for initial render
+  it("remounts and re-initializes when modelId changes (cell re-run)", async () => {
+    // Regression for marimo-team/marimo#3962: a cell re-run reconstructs
+    // the widget, opening a new comm with a new model id. The plugin
+    // must remount so initialize runs against the fresh model.
+    const { rerender } = render(<AnyWidgetSlot {...props(modelId)} />);
     await waitFor(() => {
-      expect(mockWidget.render).toHaveBeenCalled();
+      expect(mockWidget.render).toHaveBeenCalledTimes(1);
     });
 
-    // Create a new widget mock
-    const newMockWidget = {
-      initialize: vi.fn(),
-      render: vi.fn(),
-    };
-
-    // Change the widget
-    rerender(<LoadedSlot {...mockProps} widget={newMockWidget} />);
-    await TestUtils.nextTick();
-
-    // Wait for re-render with new widget
-    await waitFor(() => {
-      expect(newMockWidget.render).toHaveBeenCalled();
+    const rerunId = asModelId(`slot-test-${nextId}-rerun`);
+    WIDGET_REGISTRY.setModel(rerunId, createMockModel());
+    WIDGET_REGISTRY.setSpec(rerunId, {
+      url: SPEC.url,
+      hash: `${SPEC.hash}-${nextId}-rerun`,
     });
+    try {
+      rerender(<AnyWidgetSlot {...props(rerunId)} />);
+      await waitFor(() => {
+        expect(mockWidget.initialize).toHaveBeenCalledTimes(2);
+        expect(mockWidget.render).toHaveBeenCalledTimes(2);
+      });
+    } finally {
+      WIDGET_REGISTRY.delete(rerunId);
+    }
   });
 
-  it("should hydrate view state even when listener attaches late", async () => {
-    mockModel = new Model(
-      { count: 8 },
-      {
-        sendUpdate: vi.fn().mockResolvedValue(undefined),
-        sendCustomMessage: vi.fn().mockResolvedValue(undefined),
-      },
-    );
-    MODEL_MANAGER.set(modelId, mockModel);
-
-    const lateListenerWidget = {
-      initialize: vi.fn(),
-      render: vi.fn(({ model, el }) => {
-        // Simulate a widget view that starts with a local default and
-        // relies on change events for hydration.
-        el.textContent = "count is 5";
-        const onCount = () => {
-          el.textContent = `count is ${model.get("count")}`;
-        };
-        model.on("change:count", onCount);
-        return () => model.off("change:count", onCount);
-      }),
-    };
-
-    const { container } = render(
-      <LoadedSlot {...mockProps} widget={lateListenerWidget} />,
-    );
-
+  it("shows an error banner when the widget cannot be resolved", async () => {
+    getModuleSpy.mockRejectedValue(new Error("import failed"));
+    const { container } = render(<AnyWidgetSlot {...props(modelId)} />);
     await waitFor(() => {
-      expect(lateListenerWidget.render).toHaveBeenCalled();
-      expect(container.textContent).toContain("count is 8");
+      expect(container.textContent).toContain("import failed");
     });
   });
-});
-
-describe("isAnyWidgetModule", () => {
-  it("should accept a default object with render", () => {
-    expect(isAnyWidgetModule({ default: { render: () => undefined } })).toBe(
-      true,
-    );
-  });
 
-  it("should accept a default factory function", () => {
-    expect(
-      isAnyWidgetModule({ default: async () => ({ render: () => {} }) }),
-    ).toBe(true);
-  });
-
-  it("should reject legacy named render exports", () => {
-    expect(isAnyWidgetModule({ render: () => undefined })).toBe(false);
-  });
-});
-
-describe("getInvalidAnyWidgetModuleError", () => {
-  const jsUrl = "./@file/widget.js";
-
-  it("should explain legacy named render exports", () => {
-    const error = getInvalidAnyWidgetModuleError(
-      { render: () => undefined },
-      jsUrl,
-    );
-    expect(error.message).toContain("named exports (`render`)");
-    expect(error.message).toContain("`export default { render }`");
-    expect(error.message).toContain("not `export function render`");
-  });
-
-  it("should explain legacy named initialize exports", () => {
-    const error = getInvalidAnyWidgetModuleError(
-      { initialize: () => undefined },
-      jsUrl,
-    );
-    expect(error.message).toContain("named exports (`initialize`)");
-    expect(error.message).toContain("`export default { initialize }`");
-    expect(error.message).toContain("not `export function initialize`");
-  });
-
-  it("should avoid nested backticks for multi-hook named exports", () => {
-    const error = getInvalidAnyWidgetModuleError(
-      { render: () => undefined, initialize: () => undefined },
-      jsUrl,
-    );
-    expect(error.message).toContain(
-      "`export default { render, initialize }` (not `named export function ...`).",
-    );
-    expect(error.message).not.toContain("`named `export");
-  });
-
-  it("should explain a missing default export", () => {
-    expect(getInvalidAnyWidgetModuleError({}, jsUrl).message).toContain(
-      "missing a default export",
-    );
-    expect(getInvalidAnyWidgetModuleError(null, jsUrl).message).toContain(
-      "missing a default export",
-    );
-  });
-
-  it("should explain an invalid default export", () => {
-    const error = getInvalidAnyWidgetModuleError({ default: {} }, jsUrl);
-    expect(error.message).toContain("invalid default export");
-    expect(error.message).toContain("https://anywidget.dev/en/afm/");
+  it("shows an error banner when render fails", async () => {
+    mockWidget.render.mockRejectedValue(new Error("widget exploded"));
+    const { container } = render(<AnyWidgetSlot {...props(modelId)} />);
+    await waitFor(() => {
+      expect(container.textContent).toContain("widget exploded");
+    });
   });
 });