@marimo-team/frontend 0.19.8-dev9 → 0.19.9-dev0

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

@@ -1,99 +1,58 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
-import type { AnyWidget, Experimental } from "@anywidget/types";
-import { isEqual } from "lodash-es";
-import { useEffect, useMemo, useRef } from "react";
-import useEvent from "react-use-event-hook";
+import type { AnyWidget } from "@anywidget/types";
+import { useEffect, useRef } from "react";
 import { z } from "zod";
-import { MarimoIncomingMessageEvent } from "@/core/dom/events";
 import { asRemoteURL } from "@/core/runtime/config";
+import { resolveVirtualFileURL } from "@/core/static/files";
+import { isStaticNotebook } from "@/core/static/static-state";
 import { useAsyncData } from "@/hooks/useAsyncData";
-import { useDeepCompareMemoize } from "@/hooks/useDeepCompareMemoize";
-import {
-  type HTMLElementNotDerivedFromRef,
-  useEventListener,
-} from "@/hooks/useEventListener";
+import type { HTMLElementNotDerivedFromRef } from "@/hooks/useEventListener";
 import { createPlugin } from "@/plugins/core/builder";
-import { rpc } from "@/plugins/core/rpc";
 import type { IPluginProps } from "@/plugins/types";
-import {
-  decodeFromWire,
-  isWireFormat,
-  serializeBuffersToBase64,
-  type WireFormat,
-} from "@/utils/data-views";
 import { prettyError } from "@/utils/errors";
-import type { Base64String } from "@/utils/json/base64";
 import { Logger } from "@/utils/Logger";
 import { ErrorBanner } from "../common/error-banner";
-import { MODEL_MANAGER, Model } from "./model";
+import { getMarimoInternal, MODEL_MANAGER, type Model } from "./model";
+import type { ModelState, WidgetModelId } from "./types";
 
+/**
+ * AnyWidget asset data
+ */
 interface Data {
   jsUrl: string;
   jsHash: string;
   css?: string | null;
 }
 
-type T = Record<string, unknown>;
-
-// eslint-disable-next-line @typescript-eslint/consistent-type-definitions
-type PluginFunctions = {
-  send_to_widget: <T>(req: {
-    content: unknown;
-    buffers: Base64String[];
-  }) => Promise<null | undefined>;
-};
+type AnyWidgetState = ModelState;
 
-export const AnyWidgetPlugin = createPlugin<WireFormat<T>>("marimo-anywidget")
-  .withData(
-    z.object({
-      jsUrl: z.string(),
-      jsHash: z.string(),
-      css: z.string().nullish(),
-    }),
-  )
-  .withFunctions<PluginFunctions>({
-    send_to_widget: rpc
-      .input(
-        z.object({
-          content: z.unknown(),
-          buffers: z.array(z.string().transform((v) => v as Base64String)),
-        }),
-      )
-      .output(z.null().optional()),
-  })
-  .renderer((props) => <AnyWidgetSlot {...props} />);
-
-const AnyWidgetSlot = (
-  props: IPluginProps<WireFormat<T>, Data, PluginFunctions>,
-) => {
-  const { css, jsUrl, jsHash } = props.data;
+/**
+ * Initial value is a model_id reference.
+ * The backend sends just { model_id: string } and the frontend
+ * retrieves the actual state from the 'open' message.
+ */
+interface ModelIdRef {
+  model_id: WidgetModelId;
+}
 
-  // Decode wire format { state, bufferPaths, buffers } to state with DataViews
-  const valueWithBuffers = useMemo(() => {
-    if (isWireFormat(props.value)) {
-      const decoded = decodeFromWire(props.value);
-      Logger.debug("AnyWidget decoded wire format:", {
-        bufferPaths: props.value.bufferPaths,
-        buffersCount: props.value.buffers?.length,
-        decodedKeys: Object.keys(decoded),
-      });
-      return decoded;
-    }
-    Logger.warn("AnyWidget value is not wire format:", props.value);
-    return props.value;
-  }, [props.value]);
+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: module,
+    data: jsModule,
     error,
     refetch,
   } = useAsyncData(async () => {
-    const url = asRemoteURL(jsUrl).toString();
+    let url = asRemoteURL(jsUrl).toString();
+    // In static notebooks, resolve virtual files to blob URLs for import()
+    if (isStaticNotebook()) {
+      url = resolveVirtualFileURL(url);
+    }
     return await import(/* @vite-ignore */ url);
     // 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
@@ -111,9 +70,16 @@ const AnyWidgetSlot = (
     }
   }, [hasError, jsUrl]);
 
+  return {
+    jsModule,
+    error,
+  };
+}
+
+export function useMountCss(css: string | null | undefined, host: HTMLElement) {
   // Mount the CSS
   useEffect(() => {
-    const shadowRoot = props.host.shadowRoot;
+    const shadowRoot = host.shadowRoot;
     if (!css || !shadowRoot) {
       return;
     }
@@ -150,23 +116,38 @@ const AnyWidgetSlot = (
     return () => {
       style.remove();
     };
-  }, [css, props.host]);
+  }, [css, host]);
+}
 
-  // Wrap setValue to serialize DataViews back to base64 before sending
-  // Structure matches ipywidgets protocol: { state, bufferPaths, buffers }
-  const wrappedSetValue = useEvent((partialValue: Partial<T>) =>
-    props.setValue(serializeBuffersToBase64(partialValue)),
-  );
+export const AnyWidgetPlugin = createPlugin<ModelIdRef>("marimo-anywidget")
+  .withData(
+    z.object({
+      jsUrl: z.string(),
+      jsHash: z.string(),
+      css: z.string().nullish(),
+    }),
+  )
+  .withFunctions({})
+  .renderer((props) => <AnyWidgetSlot {...props} />);
+
+const AnyWidgetSlot = (props: IPluginProps<ModelIdRef, Data>) => {
+  const { css, jsUrl, jsHash } = props.data;
+  const { model_id: modelId } = props.value;
+  const host = props.host as HTMLElementNotDerivedFromRef;
+
+  const { jsModule, error } = useAnyWidgetModule({ jsUrl, jsHash });
+
+  useMountCss(css, host);
 
   if (error) {
     return <ErrorBanner error={error} />;
   }
 
-  if (!module) {
+  if (!jsModule) {
     return null;
   }
 
-  if (!isAnyWidgetModule(module)) {
+  if (!isAnyWidgetModule(jsModule)) {
     const error = new Error(
       `Module at ${jsUrl} does not appear to be a valid anywidget`,
     );
@@ -183,10 +164,8 @@ const AnyWidgetSlot = (
       // Plugins may be stateful and we cannot make assumptions that we won't be
       // so it is safer to just re-render.
       key={key}
-      {...props}
-      widget={module.default}
-      setValue={wrappedSetValue}
-      value={valueWithBuffers}
+      widget={jsModule.default}
+      modelId={modelId}
     />
   );
 };
@@ -194,125 +173,72 @@ const AnyWidgetSlot = (
 /**
  * Run the anywidget module
  *
- * @param widgetDef - The anywidget definition
- * @param model - The model to pass to the widget
+ * 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)
  */
-async function runAnyWidgetModule(
-  widgetDef: AnyWidget,
+async function runAnyWidgetModule<T extends AnyWidgetState>(
+  widgetDef: AnyWidget<T>,
   model: Model<T>,
   el: HTMLElement,
-): Promise<() => void> {
-  const experimental: Experimental = {
-    invoke: async (_name, _msg, _options) => {
-      const message =
-        "anywidget.invoke not supported in marimo. Please file an issue at https://github.com/marimo-team/marimo/issues";
-      Logger.warn(message);
-      throw new Error(message);
-    },
-  };
+  signal: AbortSignal,
+): Promise<void> {
   // Clear the element, in case the widget is re-rendering
   el.innerHTML = "";
-  const widget =
-    typeof widgetDef === "function" ? await widgetDef() : widgetDef;
-  await widget.initialize?.({ model, experimental });
+
   try {
-    const unsub = await widget.render?.({ model, el, experimental });
-    return () => {
-      unsub?.();
-    };
+    const render = await getMarimoInternal(model).resolveWidget(widgetDef);
+    await render(el, signal);
   } catch (error) {
     Logger.error("Error rendering anywidget", error);
     el.classList.add("text-error");
     el.innerHTML = `Error rendering anywidget: ${prettyError(error)}`;
-    return () => {
-      // No-op
-    };
   }
 }
 
 function isAnyWidgetModule(mod: any): mod is { default: AnyWidget } {
-  return (
-    mod.default &&
-    (typeof mod.default === "function" ||
-      mod.default?.render ||
-      mod.default?.initialize)
-  );
-}
-
-export function getDirtyFields(value: T, initialValue: T): Set<keyof T> {
-  return new Set(
-    Object.keys(value).filter((key) => !isEqual(value[key], initialValue[key])),
-  );
-}
+  if (!mod.default) {
+    return false;
+  }
 
-function hasModelId(message: unknown): message is { model_id: string } {
   return (
-    typeof message === "object" && message !== null && "model_id" in message
+    typeof mod.default === "function" ||
+    typeof mod.default?.render === "function" ||
+    typeof mod.default?.initialize === "function"
   );
 }
 
-interface Props
-  extends Omit<IPluginProps<T, Data, PluginFunctions>, "setValue"> {
-  widget: AnyWidget;
-  value: T;
-  setValue: (value: Partial<T>) => void;
+interface Props<T extends AnyWidgetState> {
+  widget: AnyWidget<T>;
+  modelId: WidgetModelId;
 }
 
-const LoadedSlot = ({
-  value,
-  setValue,
+const LoadedSlot = <T extends AnyWidgetState>({
   widget,
-  functions,
-  data,
-  host,
-}: Props & { widget: AnyWidget }) => {
+  modelId,
+}: Props<T> & { widget: AnyWidget<T> }) => {
   const htmlRef = useRef<HTMLDivElement>(null);
 
-  // value is already decoded from wire format
-  const model = useRef<Model<T>>(
-    new Model(value, setValue, functions.send_to_widget, new Set()),
-  );
+  // value is already decoded from wire format, may be null if waiting for open message
+  const model = MODEL_MANAGER.getSync(modelId);
 
-  // Listen to incoming messages
-  useEventListener(
-    host as HTMLElementNotDerivedFromRef,
-    MarimoIncomingMessageEvent.TYPE,
-    (e) => {
-      const message = e.detail.message;
-      if (hasModelId(message)) {
-        MODEL_MANAGER.get(message.model_id).then((model) => {
-          model.receiveCustomMessage(message, e.detail.buffers);
-        });
-      } else {
-        model.current.receiveCustomMessage(message, e.detail.buffers);
-      }
-    },
-  );
+  if (!model) {
+    Logger.error("Model not found for modelId", modelId);
+  }
 
   useEffect(() => {
-    if (!htmlRef.current) {
+    if (!htmlRef.current || !model) {
       return;
     }
-    const unsubPromise = runAnyWidgetModule(
-      widget,
-      model.current,
-      htmlRef.current,
-    );
-    return () => {
-      unsubPromise.then((unsub) => unsub());
-    };
-    // We re-run the widget when the jsUrl changes, which means the cell
+    const controller = new AbortController();
+    runAnyWidgetModule(widget, model, htmlRef.current, controller.signal);
+    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, data.jsUrl]);
-
-  // When the value changes, update the model
-  const valueMemo = useDeepCompareMemoize(value);
-  useEffect(() => {
-    model.current.updateAndEmitDiffs(valueMemo);
-  }, [valueMemo]);
+  }, [widget, modelId, model]);
 
   return <div ref={htmlRef} />;
 };
@@ -321,5 +247,4 @@ export const visibleForTesting = {
   LoadedSlot,
   runAnyWidgetModule,
   isAnyWidgetModule,
-  getDirtyFields,
 };

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

@@ -1,65 +1,17 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 
-import { act, render, waitFor } from "@testing-library/react";
+import { render, waitFor } from "@testing-library/react";
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import { TestUtils } from "@/__tests__/test-helpers";
-import type { UIElementId } from "@/core/cells/ids";
-import { MarimoIncomingMessageEvent } from "@/core/dom/events";
-import { getDirtyFields, visibleForTesting } from "../AnyWidgetPlugin";
-import { Model } from "../model";
+import type { HTMLElementNotDerivedFromRef } from "@/hooks/useEventListener";
+import { visibleForTesting } from "../AnyWidgetPlugin";
+import { MODEL_MANAGER, Model } from "../model";
+import type { WidgetModelId } from "../types";
 
 const { LoadedSlot } = visibleForTesting;
 
-describe("getDirtyFields", () => {
-  it("should return empty set when values are equal", () => {
-    const value = { foo: "bar", baz: 123 };
-    const initialValue = { foo: "bar", baz: 123 };
-
-    const result = getDirtyFields(value, initialValue);
-
-    expect(result.size).toBe(0);
-  });
-
-  it("should return keys of changed values", () => {
-    const value = { foo: "changed", baz: 123 };
-    const initialValue = { foo: "bar", baz: 123 };
-
-    const result = getDirtyFields(value, initialValue);
-
-    expect(result.size).toBe(1);
-    expect(result.has("foo")).toBe(true);
-  });
-
-  it("should handle multiple changed values", () => {
-    const value = { foo: "changed", baz: 456, extra: "new" };
-    const initialValue = { foo: "bar", baz: 123, extra: "old" };
-
-    const result = getDirtyFields(value, initialValue);
-
-    expect(result.size).toBe(3);
-    expect(result.has("foo")).toBe(true);
-    expect(result.has("baz")).toBe(true);
-    expect(result.has("extra")).toBe(true);
-  });
-
-  it("should handle nested objects correctly", () => {
-    const value = { foo: "bar", nested: { a: 1, b: 2 } };
-    const initialValue = { foo: "bar", nested: { a: 1, b: 3 } };
-
-    const result = getDirtyFields(value, initialValue);
-
-    expect(result.size).toBe(1);
-    expect(result.has("nested")).toBe(true);
-  });
-
-  it("should handle subset of initial fields", () => {
-    const value = { foo: "bar", baz: 123 };
-    const initialValue = { foo: "bar", baz: 123, full: "value" };
-
-    const result = getDirtyFields(value, initialValue);
-    expect(result.size).toBe(0);
-  });
-});
+// Helper to create typed model IDs for tests
+const asModelId = (id: string): WidgetModelId => id as WidgetModelId;
 
 // Mock a minimal AnyWidget implementation
 const mockWidget = {
@@ -76,23 +28,32 @@ vi.mock("../AnyWidgetPlugin", async () => {
 });
 
 describe("LoadedSlot", () => {
+  const modelId = asModelId("test-model-id");
+  let mockModel: Model<{ count: number }>;
+
   const mockProps = {
-    value: { count: 0 },
-    setValue: vi.fn(),
     widget: mockWidget,
-    functions: {
-      send_to_widget: vi.fn().mockResolvedValue(null),
-    },
     data: {
       jsUrl: "http://example.com/widget.js",
       jsHash: "abc123",
-      initialValue: { count: 0 },
     },
-    host: document.createElement("div"),
+    host: document.createElement(
+      "div",
+    ) as unknown as HTMLElementNotDerivedFromRef,
+    modelId: modelId,
   };
 
   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);
   });
 
   it("should render a div with ref", () => {
@@ -100,83 +61,36 @@ describe("LoadedSlot", () => {
     expect(container.querySelector("div")).not.toBeNull();
   });
 
-  it("should initialize model with merged values", () => {
-    const modelSpy = vi.spyOn(Model.prototype, "updateAndEmitDiffs");
-    render(<LoadedSlot {...mockProps} />);
-
-    expect(modelSpy).toHaveBeenCalledExactlyOnceWith({ count: 0 });
-  });
-
-  it("should update model when value prop changes", async () => {
-    const { rerender } = render(<LoadedSlot {...mockProps} />);
-    const modelSpy = vi.spyOn(Model.prototype, "updateAndEmitDiffs");
-
-    // Update the value prop
-    rerender(<LoadedSlot {...mockProps} value={{ count: 5 }} />);
-
-    // Model should be updated with the new value
-    expect(modelSpy).toHaveBeenCalledWith({ count: 5 });
-  });
-
-  it("should listen for incoming messages", async () => {
+  it("should call runAnyWidgetModule on initialization", async () => {
     render(<LoadedSlot {...mockProps} />);
 
-    // Send a mock message
-    const mockMessageEvent = MarimoIncomingMessageEvent.create({
-      detail: {
-        objectId: "test-id" as UIElementId,
-        message: {
-          method: "update",
-          state: { count: 10 },
-          buffer_paths: [],
-        },
-        buffers: [],
-      },
-      bubbles: false,
-      composed: true,
-    });
-    const updateAndEmitDiffsSpy = vi.spyOn(Model.prototype, "set");
-
-    // Dispatch the event on the host element
-    act(() => {
-      mockProps.host.dispatchEvent(mockMessageEvent);
-    });
-
+    // Wait a render
     await waitFor(() => {
-      expect(updateAndEmitDiffsSpy).toHaveBeenCalledWith("count", 10);
+      expect(mockWidget.render).toHaveBeenCalled();
     });
   });
 
-  it("should call runAnyWidgetModule on initialization", async () => {
+  it("should re-run widget when widget prop changes", async () => {
     const { rerender } = render(<LoadedSlot {...mockProps} />);
 
-    // Wait a render
+    // Wait for initial render
     await waitFor(() => {
       expect(mockWidget.render).toHaveBeenCalled();
     });
 
-    // Render without any prop changes
-    rerender(<LoadedSlot {...mockProps} />);
-    await TestUtils.nextTick();
+    // Create a new widget mock
+    const newMockWidget = {
+      initialize: vi.fn(),
+      render: vi.fn(),
+    };
 
-    // Still only called once
-    expect(mockWidget.render).toHaveBeenCalledTimes(1);
-
-    // Change the jsUrl
-    rerender(
-      <LoadedSlot
-        {...mockProps}
-        data={{
-          ...mockProps.data,
-          jsUrl: "http://example.com/widget-updated.js",
-        }}
-      />,
-    );
+    // Change the widget
+    rerender(<LoadedSlot {...mockProps} widget={newMockWidget} />);
     await TestUtils.nextTick();
 
-    // Wait a render
+    // Wait for re-render with new widget
     await waitFor(() => {
-      expect(mockWidget.render).toHaveBeenCalledTimes(2);
+      expect(newMockWidget.render).toHaveBeenCalled();
     });
   });
 });