@marimo-team/frontend 0.20.5-dev9 → 0.20.5-dev91

package/src/plugins/impl/mpl-interactive/MplInteractivePlugin.tsx

@@ -0,0 +1,309 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+/* eslint-disable @typescript-eslint/no-explicit-any */
+
+import { useCallback, useEffect, useRef } from "react";
+import { z } from "zod";
+import { useEventListener } from "@/hooks/useEventListener";
+import { createPlugin } from "@/plugins/core/builder";
+import { MODEL_MANAGER, type Model } from "@/plugins/impl/anywidget/model";
+import type { ModelState, WidgetModelId } from "@/plugins/impl/anywidget/types";
+import type { IPluginProps } from "@/plugins/types";
+import { downloadBlob } from "@/utils/download";
+import { Logger } from "@/utils/Logger";
+import { MplCommWebSocket } from "./mpl-websocket-shim";
+
+const MPL_SCOPE_CLASS = "mpl-interactive-figure";
+
+interface Data {
+  mplJsUrl: string;
+  cssUrl: string;
+  toolbarImages: Record<string, string>;
+  width: number;
+  height: number;
+}
+
+interface ModelIdRef {
+  model_id: WidgetModelId;
+}
+
+declare global {
+  interface Window {
+    mpl: {
+      figure: new (
+        id: string,
+        ws: MplCommWebSocket,
+        ondownload: (figure: MplFigure, format: string) => void,
+        element: HTMLElement,
+      ) => MplFigure;
+      toolbar_items: [
+        string | null,
+        string | null,
+        string | null,
+        string | null,
+      ][];
+    };
+  }
+}
+
+interface MplFigure {
+  id: string;
+  ws: MplCommWebSocket;
+  root: HTMLElement;
+  send_message: (type: string, properties: Record<string, unknown>) => void;
+}
+
+export const MplInteractivePlugin = createPlugin<ModelIdRef>(
+  "marimo-mpl-interactive",
+)
+  .withData(
+    z.object({
+      mplJsUrl: z.string(),
+      cssUrl: z.string(),
+      toolbarImages: z.record(z.string(), z.string()),
+      width: z.number(),
+      height: z.number(),
+    }),
+  )
+  .withFunctions({})
+  .renderer((props) => <MplInteractiveSlot {...props} />);
+
+let mplJsLoading: Promise<void> | null = null;
+
+async function ensureMplJs(jsUrl: string): Promise<void> {
+  if (window.mpl) {
+    return;
+  }
+  if (mplJsLoading) {
+    return mplJsLoading;
+  }
+  mplJsLoading = new Promise<void>((resolve, reject) => {
+    const script = document.createElement("script");
+    script.src = jsUrl;
+    script.onload = () => resolve();
+    script.onerror = () => {
+      mplJsLoading = null;
+      reject(new Error("Failed to load mpl.js"));
+    };
+    document.head.append(script);
+  });
+  return mplJsLoading;
+}
+
+/**
+ * Patch mpl.js toolbar image references to use inline data URIs.
+ *
+ * mpl.js sets `icon_img.src = '_images/' + image + '.png'` and
+ * `icon_img.srcset = '_images/' + image + '_large.png 2x'`.
+ *
+ * We observe the container for new <img> elements and rewrite their
+ * src/srcset to the inlined base64 data URIs.
+ */
+function patchToolbarImages(
+  container: HTMLElement,
+  toolbarImages: Record<string, string>,
+): () => void {
+  const patchImg = (img: HTMLImageElement) => {
+    const src = img.getAttribute("src") || "";
+    const match = src.match(/_images\/(.+)\.png$/);
+    if (match) {
+      const name = match[1];
+      const dataUri = toolbarImages[name];
+      if (dataUri) {
+        img.src = dataUri;
+      }
+    }
+    const srcset = img.getAttribute("srcset") || "";
+    const srcsetMatch = srcset.match(/_images\/(.+)\.png\s+2x$/);
+    if (srcsetMatch) {
+      const name = srcsetMatch[1];
+      const dataUri = toolbarImages[name];
+      if (dataUri) {
+        img.srcset = `${dataUri} 2x`;
+      }
+    }
+  };
+
+  // Patch any existing images
+  for (const img of container.querySelectorAll("img")) {
+    patchImg(img);
+  }
+
+  // Observe for new images added by mpl.js
+  const observer = new MutationObserver((mutations) => {
+    for (const mutation of mutations) {
+      for (const node of mutation.addedNodes) {
+        if (node instanceof HTMLImageElement) {
+          patchImg(node);
+        } else if (node instanceof HTMLElement) {
+          for (const img of node.querySelectorAll("img")) {
+            patchImg(img);
+          }
+        }
+      }
+    }
+  });
+
+  observer.observe(container, { childList: true, subtree: true });
+  return () => observer.disconnect();
+}
+
+function injectCss(container: HTMLElement, cssUrl: string): () => void {
+  const link = document.createElement("link");
+  link.rel = "stylesheet";
+  link.href = cssUrl;
+  container.append(link);
+  return () => link.remove();
+}
+
+const MplInteractiveSlot = (props: IPluginProps<ModelIdRef, Data>) => {
+  const { mplJsUrl, cssUrl, toolbarImages, width, height } = props.data;
+  const { model_id: modelId } = props.value;
+  const containerRef = useRef<HTMLDivElement>(null);
+  const figureRef = useRef<MplFigure | null>(null);
+  const wsRef = useRef<MplCommWebSocket | null>(null);
+
+  const setupFigure = useCallback(
+    async (container: HTMLElement) => {
+      // Load mpl.js globally (only once, via <script src>)
+      await ensureMplJs(mplJsUrl);
+
+      if (!window.mpl) {
+        Logger.error("mpl.js failed to load");
+        return;
+      }
+
+      // Get the model from MODEL_MANAGER
+      let model: Model<ModelState>;
+      try {
+        model = await MODEL_MANAGER.get(modelId);
+      } catch {
+        Logger.error("Failed to get model for mpl interactive", modelId);
+        return;
+      }
+
+      // Create the fake WebSocket
+      const fakeWs = new MplCommWebSocket((msg: unknown) => {
+        // Send from frontend → backend via model custom message
+        model.send(msg);
+      });
+      wsRef.current = fakeWs;
+
+      // Listen for backend → frontend messages via model custom events
+      const handleCustomMessage = (
+        content: { type: string; data?: unknown; format?: string },
+        buffers?: readonly DataView[],
+      ) => {
+        if (!content) {
+          return;
+        }
+
+        if (content.type === "json") {
+          fakeWs.receiveJson(content.data);
+        } else if (content.type === "binary" && buffers && buffers.length > 0) {
+          fakeWs.receiveBinary(buffers[0]);
+        } else if (
+          content.type === "download" &&
+          buffers &&
+          buffers.length > 0
+        ) {
+          const fmt = content.format || "png";
+          const dv = buffers[0];
+          const ab = dv.buffer.slice(
+            dv.byteOffset,
+            dv.byteOffset + dv.byteLength,
+          ) as ArrayBuffer;
+          downloadBlob(
+            new Blob([ab], { type: `image/${fmt}` }),
+            `figure.${fmt}`,
+          );
+        }
+      };
+
+      model.on("msg:custom", handleCustomMessage as any);
+
+      // Create the mpl figure
+      const figId = modelId;
+      const ondownload = (_figure: MplFigure, format: string) => {
+        // Send download request to backend
+        model.send({ type: "download", format });
+      };
+
+      const fig = new window.mpl.figure(figId, fakeWs, ondownload, container);
+      figureRef.current = fig;
+
+      // Set the canvas_div to the backend's figure size so the
+      // ResizeObserver doesn't trigger an immediate resize cycle.
+      // mpl.js creates: fig.root > [titlebar, canvas_div, toolbar]
+      const canvasDiv = fig.root.querySelector<HTMLElement>("div[tabindex]");
+      if (canvasDiv) {
+        canvasDiv.style.width = `${width}px`;
+        canvasDiv.style.height = `${height}px`;
+      }
+
+      // Trigger the onopen callback to start communication
+      // mpl.js sends initial messages in onopen
+      setTimeout(() => {
+        fakeWs.onopen?.();
+      }, 0);
+
+      return () => {
+        model.off("msg:custom", handleCustomMessage as any);
+        fakeWs.close();
+      };
+    },
+    [modelId, mplJsUrl, width, height],
+  );
+
+  useEffect(() => {
+    const container = containerRef.current;
+    if (!container) {
+      return;
+    }
+
+    // Clear any previous content (handles re-render / cell re-run)
+    container.innerHTML = "";
+
+    // Inject CSS
+    const removeCss = injectCss(container, cssUrl);
+
+    // Patch toolbar images
+    const removeImageObserver = patchToolbarImages(container, toolbarImages);
+
+    let cleanup: (() => void) | undefined;
+    let cancelled = false;
+
+    setupFigure(container)
+      .then((cleanupFn) => {
+        if (cancelled) {
+          cleanupFn?.();
+          return;
+        }
+        cleanup = cleanupFn;
+      })
+      .catch((error) => {
+        if (!cancelled) {
+          Logger.error("Failed to set up MPL interactive figure", error);
+        }
+      });
+
+    return () => {
+      cancelled = true;
+      removeCss();
+      removeImageObserver();
+      cleanup?.();
+      // Clear DOM on unmount so stale content doesn't linger
+      container.innerHTML = "";
+    };
+  }, [modelId, cssUrl, toolbarImages, setupFigure]);
+
+  // Re-request figure when tab becomes visible
+  useEventListener(document, "visibilitychange", () => {
+    const fig = figureRef.current;
+    if (!document.hidden && fig?.ws?.readyState === WebSocket.OPEN) {
+      fig.send_message("refresh", {});
+    }
+  });
+
+  // Must match _MPL_SCOPE in from_mpl_interactive.py
+  return <div ref={containerRef} className={MPL_SCOPE_CLASS} />;
+};

package/src/plugins/impl/mpl-interactive/__tests__/mpl-websocket-shim.test.ts

@@ -0,0 +1,110 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+import { describe, expect, it, vi } from "vitest";
+import { MplCommWebSocket } from "../mpl-websocket-shim";
+
+describe("MplCommWebSocket", () => {
+  it("starts in OPEN state", () => {
+    const ws = new MplCommWebSocket(vi.fn());
+    expect(ws.readyState).toBe(WebSocket.OPEN);
+  });
+
+  it("send() parses JSON and calls sendFn with parsed object", () => {
+    const sendFn = vi.fn();
+    const ws = new MplCommWebSocket(sendFn);
+
+    ws.send(JSON.stringify({ type: "resize", width: 640, height: 480 }));
+
+    expect(sendFn).toHaveBeenCalledOnce();
+    expect(sendFn).toHaveBeenCalledWith({
+      type: "resize",
+      width: 640,
+      height: 480,
+    });
+  });
+
+  it("receiveJson() dispatches MessageEvent with JSON string data", () => {
+    const ws = new MplCommWebSocket(vi.fn());
+    const handler = vi.fn();
+    ws.onmessage = handler;
+
+    ws.receiveJson({ type: "figure_size", size: [640, 480] });
+
+    expect(handler).toHaveBeenCalledOnce();
+    const event: MessageEvent = handler.mock.calls[0][0];
+    expect(event).toBeInstanceOf(MessageEvent);
+    expect(event.type).toBe("message");
+    expect(JSON.parse(event.data as string)).toEqual({
+      type: "figure_size",
+      size: [640, 480],
+    });
+  });
+
+  it("receiveJson() is a no-op if onmessage is not set", () => {
+    const ws = new MplCommWebSocket(vi.fn());
+    ws.receiveJson({ type: "test" });
+    expect(ws.onmessage).toBeNull();
+  });
+
+  it("receiveBinary() dispatches MessageEvent with Blob data", () => {
+    const ws = new MplCommWebSocket(vi.fn());
+    const handler = vi.fn();
+    ws.onmessage = handler;
+
+    // Simulate a PNG-like binary buffer
+    const bytes = new Uint8Array([0x89, 0x50, 0x4e, 0x47, 0, 1, 2, 3]);
+    const dv = new DataView(bytes.buffer);
+
+    ws.receiveBinary(dv);
+
+    expect(handler).toHaveBeenCalledOnce();
+    const event: MessageEvent = handler.mock.calls[0][0];
+    expect(event).toBeInstanceOf(MessageEvent);
+    expect(event.data).toBeInstanceOf(Blob);
+    expect((event.data as Blob).size).toBe(8);
+  });
+
+  it("receiveBinary() handles DataView with offset", () => {
+    const ws = new MplCommWebSocket(vi.fn());
+    const handler = vi.fn();
+    ws.onmessage = handler;
+
+    // Create a DataView that's a slice of a larger buffer
+    const fullBuffer = new ArrayBuffer(16);
+    const fullView = new Uint8Array(fullBuffer);
+    fullView.set([0, 0, 0, 0, 0x89, 0x50, 0x4e, 0x47, 0, 1, 2, 3, 0, 0, 0, 0]);
+    const dv = new DataView(fullBuffer, 4, 8);
+
+    ws.receiveBinary(dv);
+
+    expect(handler).toHaveBeenCalledOnce();
+    const blob = handler.mock.calls[0][0].data as Blob;
+    expect(blob.size).toBe(8);
+  });
+
+  it("close() sets readyState to CLOSED and fires onclose", () => {
+    const ws = new MplCommWebSocket(vi.fn());
+    const closeHandler = vi.fn();
+    ws.onclose = closeHandler;
+
+    ws.close();
+
+    expect(ws.readyState).toBe(WebSocket.CLOSED);
+    expect(closeHandler).toHaveBeenCalledOnce();
+  });
+
+  it("close() does not throw if onclose is not set", () => {
+    const ws = new MplCommWebSocket(vi.fn());
+    ws.close();
+    expect(ws.readyState).toBe(WebSocket.CLOSED);
+  });
+
+  it("onopen can be triggered externally", () => {
+    const ws = new MplCommWebSocket(vi.fn());
+    const openHandler = vi.fn();
+    ws.onopen = openHandler;
+
+    ws.onopen?.();
+
+    expect(openHandler).toHaveBeenCalledOnce();
+  });
+});

package/src/plugins/impl/mpl-interactive/mpl-websocket-shim.ts

@@ -0,0 +1,57 @@
+/* Copyright 2026 Marimo. All rights reserved. */
+
+/**
+ * Fake WebSocket that routes messages through MarimoComm / MODEL_MANAGER
+ * instead of a real network WebSocket.
+ *
+ * mpl.js expects a WebSocket-like object with:
+ *   - readyState
+ *   - send(data: string)
+ *   - onopen / onmessage / onclose callbacks
+ */
+export class MplCommWebSocket {
+  readyState: number = WebSocket.OPEN;
+  private sendFn: (msg: unknown) => void;
+
+  onopen: (() => void) | null = null;
+  onmessage: ((evt: MessageEvent) => void) | null = null;
+  onclose: (() => void) | null = null;
+
+  constructor(sendFn: (msg: unknown) => void) {
+    this.sendFn = sendFn;
+  }
+
+  /**
+   * Called by mpl.js to send a message to the backend.
+   * mpl.js always sends JSON strings.
+   */
+  send(data: string): void {
+    this.sendFn(JSON.parse(data));
+  }
+
+  /**
+   * Called when the backend pushes a JSON message via the model custom event.
+   */
+  receiveJson(data: unknown): void {
+    this.onmessage?.(
+      new MessageEvent("message", { data: JSON.stringify(data) }),
+    );
+  }
+
+  /**
+   * Called when the backend pushes binary data (PNG render) via model custom event.
+   */
+  receiveBinary(buffer: DataView): void {
+    const ab = buffer.buffer.slice(
+      buffer.byteOffset,
+      buffer.byteOffset + buffer.byteLength,
+    ) as ArrayBuffer;
+    const blob = new Blob([ab]);
+    this.onmessage?.(new MessageEvent("message", { data: blob }));
+  }
+
+  close(): void {
+    this.readyState = WebSocket.CLOSED;
+    this.onclose?.();
+  }
+}

package/src/plugins/impl/plotly/PlotlyPlugin.tsx

@@ -250,6 +250,7 @@ const STANDARD_POINT_KEYS: string[] = [
   "lon",
   "curveNumber",
   "pointNumber",
+  "pointNumbers",
   "pointIndex",
 ];
 
@@ -263,6 +264,8 @@ function extractPoints(
   let parser: PlotlyTemplateParser | undefined;
 
   return points.map((point) => {
+    const standardPointFields = pick(point, STANDARD_POINT_KEYS);
+
     // Get the first hovertemplate
     const hovertemplate = Array.isArray(point.data.hovertemplate)
       ? point.data.hovertemplate[0]
@@ -271,13 +274,16 @@ function extractPoints(
     // For chart types with standard point keys (e.g. heatmaps),
     // or when there's no hovertemplate, pick keys directly from the point.
     if (!hovertemplate || point.data?.type === "heatmap") {
-      return pick(point, STANDARD_POINT_KEYS);
+      return standardPointFields;
     }
 
     // Update or create a parser
     parser = parser
       ? parser.update(hovertemplate)
       : createParser(hovertemplate);
-    return parser.parse(point);
+    return {
+      ...standardPointFields,
+      ...parser.parse(point),
+    };
   });
 }

package/src/plugins/impl/vega/vega-component.tsx

@@ -16,6 +16,7 @@ import { Tooltip } from "@/components/ui/tooltip";
 import { useAsyncData } from "@/hooks/useAsyncData";
 import { useDeepCompareMemoize } from "@/hooks/useDeepCompareMemoize";
 import { useTheme } from "@/theme/useTheme";
+import { cn } from "@/utils/cn";
 import { Events } from "@/utils/events";
 import { Logger } from "@/utils/Logger";
 import { Objects } from "@/utils/objects";
@@ -304,7 +305,12 @@ const LoadedVegaComponent = ({
         </Alert>
       )}
       <div
-        className="relative"
+        className={cn(
+          "relative",
+          "width" in selectableSpec &&
+            selectableSpec.width === "container" &&
+            "vega-container-width",
+        )}
         // Capture the pointer down event to prevent the parent from handling it
         onPointerDown={Events.stopPropagation()}
       >

package/src/plugins/impl/vega/vega.css

@@ -1,18 +1,11 @@
 @reference "../../../css/globals.css";
 
 .vega-embed {
-  width: 100%;
-  flex: 1;
-  display: inline-block;
-
-  @media (min-width: 500px) {
-    min-width: 300px;
-  }
+  max-width: 100%;
+}
 
-  /* For vega embeds in slots, reset the styles to let the user set the width */
-  @container style(--slot: true) {
-    min-width: unset;
-  }
+.vega-container-width .vega-embed {
+  width: 100%;
 }
 
 .vega-embed > .chart-wrapper {

package/src/plugins/plugins.ts

@@ -26,6 +26,7 @@ import { MatrixPlugin } from "./impl/MatrixPlugin";
 import { MicrophonePlugin } from "./impl/MicrophonePlugin";
 import { MultiselectPlugin } from "./impl/MultiselectPlugin";
 import { MatplotlibPlugin } from "./impl/matplotlib/MatplotlibPlugin";
+import { MplInteractivePlugin } from "./impl/mpl-interactive/MplInteractivePlugin";
 import { NumberPlugin } from "./impl/NumberPlugin";
 import { PanelPlugin } from "./impl/panel/PanelPlugin";
 import { PlotlyPlugin } from "./impl/plotly/PlotlyPlugin";
@@ -92,6 +93,7 @@ export const UI_PLUGINS: IPlugin<any, unknown>[] = [
   AnyWidgetPlugin,
   DataEditorPlugin,
   PanelPlugin,
+  MplInteractivePlugin,
 ];
 
 // List of output / layout plugins