Unify handling of widget state/buffers/bufferPaths (#7441)

This keeps the `state/buffer/bufferPaths` in sync when sending to/from
the backend for the widget. This changes some typings and adds a few
more tests. This is an improvement, but not the final state as we can
hopefully simplify further. This does fix a few errors when sending back
binary data.

Prior work done by @manzt in
9426b13

---------

Co-authored-by: Trevor Manz <trevor.j.manz@gmail.com>

Author: mscolnick

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

@@ -2,8 +2,9 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
 import type { AnyWidget, Experimental } from "@anywidget/types";
-import { get, isEqual, set } from "lodash-es";
+import { isEqual } from "lodash-es";
 import { useEffect, useMemo, useRef } from "react";
+import useEvent from "react-use-event-hook";
 import { z } from "zod";
 import { MarimoIncomingMessageEvent } from "@/core/dom/events";
 import { asRemoteURL } from "@/core/runtime/config";
@@ -17,10 +18,13 @@ import { createPlugin } from "@/plugins/core/builder";
 import { rpc } from "@/plugins/core/rpc";
 import type { IPluginProps } from "@/plugins/types";
 import {
-  type Base64String,
-  byteStringToBinary,
-  typedAtob,
-} from "@/utils/json/base64";
+  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";
@@ -29,44 +33,56 @@ interface Data {
   jsUrl: string;
   jsHash: string;
   css?: string | null;
-  bufferPaths?: (string | number)[][] | null;
-  initialValue: T;
 }
 
-type T = Record<string, any>;
+type T = Record<string, unknown>;
 
-// eslint-disable-next-line @typescript-eslint/consistent-type-definitions
 type PluginFunctions = {
-  send_to_widget: <T>(req: { content?: any }) => Promise<null | undefined>;
+  send_to_widget: <T>(req: {
+    content: unknown;
+    buffers: Base64String[];
+  }) => Promise<null | undefined>;
 };
 
-export const AnyWidgetPlugin = createPlugin<T>("marimo-anywidget")
+export const AnyWidgetPlugin = createPlugin<WireFormat<T>>("marimo-anywidget")
   .withData(
     z.object({
       jsUrl: z.string(),
       jsHash: z.string(),
       css: z.string().nullish(),
-      bufferPaths: z
-        .array(z.array(z.union([z.string(), z.number()])))
-        .nullish(),
-      initialValue: z.object({}).passthrough(),
     }),
   )
   .withFunctions<PluginFunctions>({
     send_to_widget: rpc
-      .input(z.object({ content: z.any() }))
+      .input(
+        z.object({
+          content: z.unknown(),
+          buffers: z.array(z.string().transform((v) => v as Base64String)),
+        }),
+      )
       .output(z.null().optional()),
   })
   .renderer((props) => <AnyWidgetSlot {...props} />);
 
-type Props = IPluginProps<T, Data, PluginFunctions>;
-
-const AnyWidgetSlot = (props: Props) => {
-  const { css, jsUrl, jsHash, bufferPaths } = props.data;
+const AnyWidgetSlot = (
+  props: IPluginProps<WireFormat<T>, Data, PluginFunctions>,
+) => {
+  const { css, jsUrl, jsHash } = props.data;
 
+  // Decode wire format { state, bufferPaths, buffers } to state with DataViews
   const valueWithBuffers = useMemo(() => {
-    return resolveInitialValue(props.value, bufferPaths ?? []);
-  }, [props.value, bufferPaths]);
+    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]);
 
   // JS is an ESM file with a render function on it
   // export function render({ model, el }) {
@@ -135,6 +151,12 @@ const AnyWidgetSlot = (props: Props) => {
     };
   }, [css, props.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)),
+  );
+
   if (error) {
     return <ErrorBanner error={error} />;
   }
@@ -162,6 +184,7 @@ const AnyWidgetSlot = (props: Props) => {
       key={key}
       {...props}
       widget={module.default}
+      setValue={wrappedSetValue}
       value={valueWithBuffers}
     />
   );
@@ -191,10 +214,19 @@ async function runAnyWidgetModule(
   const widget =
     typeof widgetDef === "function" ? await widgetDef() : widgetDef;
   await widget.initialize?.({ model, experimental });
-  const unsub = await widget.render?.({ model, el, experimental });
-  return () => {
-    unsub?.();
-  };
+  try {
+    const unsub = await widget.render?.({ model, el, experimental });
+    return () => {
+      unsub?.();
+    };
+  } 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 } {
@@ -218,6 +250,13 @@ function hasModelId(message: unknown): message is { model_id: string } {
   );
 }
 
+interface Props
+  extends Omit<IPluginProps<T, Data, PluginFunctions>, "setValue"> {
+  widget: AnyWidget;
+  value: T;
+  setValue: (value: Partial<T>) => void;
+}
+
 const LoadedSlot = ({
   value,
   setValue,
@@ -228,15 +267,9 @@ const LoadedSlot = ({
 }: Props & { widget: AnyWidget }) => {
   const htmlRef = useRef<HTMLDivElement>(null);
 
+  // value is already decoded from wire format
   const model = useRef<Model<T>>(
-    new Model(
-      // Merge the initial value with the current value
-      // since we only send partial updates to the backend
-      { ...data.initialValue, ...value },
-      setValue,
-      functions.send_to_widget,
-      getDirtyFields(value, data.initialValue),
-    ),
+    new Model(value, setValue, functions.send_to_widget, new Set()),
   );
 
   // Listen to incoming messages
@@ -289,16 +322,3 @@ export const visibleForTesting = {
   isAnyWidgetModule,
   getDirtyFields,
 };
-
-export function resolveInitialValue(
-  raw: Record<string, any>,
-  bufferPaths: readonly (readonly (string | number)[])[],
-) {
-  const out = structuredClone(raw);
-  for (const bufferPath of bufferPaths) {
-    const base64String: Base64String = get(raw, bufferPath);
-    const bytes = byteStringToBinary(typedAtob(base64String));
-    set(out, bufferPath, new DataView(bytes.buffer));
-  }
-  return out;
-}

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

@@ -5,11 +5,7 @@ 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,
-  resolveInitialValue,
-  visibleForTesting,
-} from "../AnyWidgetPlugin";
+import { getDirtyFields, visibleForTesting } from "../AnyWidgetPlugin";
 import { Model } from "../model";
 
 const { LoadedSlot } = visibleForTesting;
@@ -132,6 +128,7 @@ describe("LoadedSlot", () => {
         message: {
           method: "update",
           state: { count: 10 },
+          buffer_paths: [],
         },
         buffers: [],
       },
@@ -183,55 +180,3 @@ describe("LoadedSlot", () => {
     });
   });
 });
-
-describe("resolveInitialValue", () => {
-  it("should convert base64 strings to DataView at specified paths", () => {
-    const result = resolveInitialValue(
-      {
-        a: 10,
-        b: "aGVsbG8=", // "hello" in base64
-        c: [1, "d29ybGQ="], // "world" in base64
-        d: {
-          foo: "bWFyaW1vCg==", // "marimo" in base64
-          baz: 20,
-        },
-      },
-      [["b"], ["c", 1], ["d", "foo"]],
-    );
-
-    expect(result).toMatchInlineSnapshot(`
-      {
-        "a": 10,
-        "b": DataView [
-          104,
-          101,
-          108,
-          108,
-          111,
-        ],
-        "c": [
-          1,
-          DataView [
-            119,
-            111,
-            114,
-            108,
-            100,
-          ],
-        ],
-        "d": {
-          "baz": 20,
-          "foo": DataView [
-            109,
-            97,
-            114,
-            105,
-            109,
-            111,
-            10,
-          ],
-        },
-      }
-    `);
-  });
-});

frontend/src/plugins/impl/anywidget/__tests__/model.test.ts

@@ -9,6 +9,7 @@ import {
   vi,
 } from "vitest";
 import { TestUtils } from "@/__tests__/test-helpers";
+import type { Base64String } from "@/utils/json/base64";
 import {
   type AnyWidgetMessage,
   handleWidgetMessage,
@@ -23,7 +24,10 @@ describe("Model", () => {
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   let onChange: (value: any) => void;
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
-  let sendToWidget: (req: { content?: any }) => Promise<null | undefined>;
+  let sendToWidget: (req: {
+    content: unknown;
+    buffers: Base64String[];
+  }) => Promise<null | undefined>;
 
   beforeEach(() => {
     onChange = vi.fn();
@@ -72,7 +76,7 @@ describe("Model", () => {
       });
     });
 
-    it("should send all dirty fields", () => {
+    it("should clear dirty fields after save", () => {
       model.set("foo", "new value");
       model.save_changes();
 
@@ -83,14 +87,13 @@ describe("Model", () => {
       model.set("bar", 456);
       model.save_changes();
 
+      // After clearing, only the newly changed field is sent
       expect(onChange).toHaveBeenCalledWith({
-        foo: "new value",
         bar: 456,
       });
     });
 
-    // Skip because we don't clear the dirty fields after save
-    it.skip("should clear dirty fields after save", () => {
+    it("should not call onChange when no dirty fields", () => {
       model.set("foo", "new value");
       model.save_changes();
       model.save_changes(); // Second save should not call onChange
@@ -144,21 +147,16 @@ describe("Model", () => {
       const callback = vi.fn();
       model.send({ test: true }, callback);
 
-      expect(sendToWidget).toHaveBeenCalledWith({ content: { test: true } });
+      expect(sendToWidget).toHaveBeenCalledWith({
+        content: {
+          state: { test: true },
+          bufferPaths: [],
+        },
+        buffers: [],
+      });
       await TestUtils.nextTick(); // flush
       expect(callback).toHaveBeenCalledWith(null);
     });
-
-    it("should warn when buffers are provided", () => {
-      const consoleSpy = vi.spyOn(console, "warn").mockImplementation(() => {
-        // noop
-      });
-      model.send({ test: true }, null, [new ArrayBuffer(8)]);
-
-      expect(consoleSpy).toHaveBeenCalledWith(
-        "buffers not supported in marimo anywidget.send",
-      );
-    });
   });
 
   describe("widget_manager", () => {
@@ -228,7 +226,11 @@ describe("Model", () => {
     it("should handle update messages", () => {
       model.receiveCustomMessage({
         method: "update",
-        state: { foo: "updated", bar: 789 },
+        state: {
+          foo: "updated",
+          bar: 789,
+        },
+        buffer_paths: [],
       });
 
       expect(model.get("foo")).toBe("updated");
@@ -333,7 +335,9 @@ describe("ModelManager", () => {
 
     const updateMessage: AnyWidgetMessage = {
       method: "update",
-      state: { count: 1 },
+      state: {
+        count: 1,
+      },
       buffer_paths: [],
     };