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

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

@@ -9,35 +9,73 @@ import {
   vi,
 } from "vitest";
 import { TestUtils } from "@/__tests__/test-helpers";
-import type { Base64String } from "@/utils/json/base64";
 import {
-  type AnyWidgetMessage,
+  getMarimoInternal,
   handleWidgetMessage,
   Model,
   visibleForTesting,
 } from "../model";
+import type { WidgetModelId } from "../types";
 
 const { ModelManager } = visibleForTesting;
 
+// Helper to create typed model IDs for tests
+const asModelId = (id: string): WidgetModelId => id as WidgetModelId;
+
+// Mock the request client
+const mockSendModelValue = vi.fn().mockResolvedValue(null);
+vi.mock("@/core/network/requests", () => ({
+  getRequestClient: () => ({
+    sendModelValue: mockSendModelValue,
+  }),
+}));
+
+// Helper to create a mock MarimoComm
+function createMockComm<T>() {
+  return {
+    sendUpdate: vi.fn().mockResolvedValue(undefined),
+    sendCustomMessage: vi.fn().mockResolvedValue(undefined),
+  };
+}
+
 describe("Model", () => {
   let model: Model<{ foo: string; bar: number }>;
-  // 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: unknown;
-    buffers: Base64String[];
-  }) => Promise<null | undefined>;
+  let mockComm: ReturnType<typeof createMockComm<{ foo: string; bar: number }>>;
 
   beforeEach(() => {
-    onChange = vi.fn();
-    sendToWidget = vi.fn().mockResolvedValue(null);
-    model = new Model(
-      { foo: "test", bar: 123 },
-      onChange,
-      sendToWidget,
-      new Set(),
-    );
+    mockComm = createMockComm();
+    mockSendModelValue.mockClear();
+    model = new Model({ foo: "test", bar: 123 }, mockComm);
+  });
+
+  describe("public API", () => {
+    it("should only expose AFM-compliant interface", () => {
+      // Get all enumerable own properties
+      const ownProperties = Object.keys(model).sort();
+      // Get prototype methods (excluding constructor)
+      const prototypeMethods = Object.getOwnPropertyNames(
+        Object.getPrototypeOf(model),
+      )
+        .filter((name) => name !== "constructor")
+        .sort();
+
+      // Snapshot the public API to catch accidental leaks of internal methods
+      expect({ ownProperties, prototypeMethods }).toMatchInlineSnapshot(`
+        {
+          "ownProperties": [
+            "widget_manager",
+          ],
+          "prototypeMethods": [
+            "get",
+            "off",
+            "on",
+            "save_changes",
+            "send",
+            "set",
+          ],
+        }
+      `);
+    });
   });
 
   describe("get/set", () => {
@@ -70,7 +108,7 @@ describe("Model", () => {
       model.set("bar", 456);
       model.save_changes();
 
-      expect(onChange).toHaveBeenCalledWith({
+      expect(mockComm.sendUpdate).toHaveBeenCalledWith({
         foo: "new value",
         bar: 456,
       });
@@ -80,7 +118,7 @@ describe("Model", () => {
       model.set("foo", "new value");
       model.save_changes();
 
-      expect(onChange).toHaveBeenCalledWith({
+      expect(mockComm.sendUpdate).toHaveBeenCalledWith({
         foo: "new value",
       });
 
@@ -88,17 +126,17 @@ describe("Model", () => {
       model.save_changes();
 
       // After clearing, only the newly changed field is sent
-      expect(onChange).toHaveBeenCalledWith({
+      expect(mockComm.sendUpdate).toHaveBeenCalledWith({
         bar: 456,
       });
     });
 
-    it("should not call onChange when no dirty fields", () => {
+    it("should not call sendUpdate when no dirty fields", () => {
       model.set("foo", "new value");
       model.save_changes();
-      model.save_changes(); // Second save should not call onChange
+      model.save_changes(); // Second save should not call sendUpdate
 
-      expect(onChange).toHaveBeenCalledTimes(1);
+      expect(mockComm.sendUpdate).toHaveBeenCalledTimes(1);
     });
   });
 
@@ -145,23 +183,28 @@ describe("Model", () => {
   describe("send", () => {
     it("should send message and handle callbacks", async () => {
       const callback = vi.fn();
-      model.send({ test: true }, callback);
+      await model.send({ test: true }, callback);
 
-      expect(sendToWidget).toHaveBeenCalledWith({
-        content: {
-          state: { test: true },
-          bufferPaths: [],
-        },
-        buffers: [],
-      });
-      await TestUtils.nextTick(); // flush
-      expect(callback).toHaveBeenCalledWith(null);
+      expect(mockComm.sendCustomMessage).toHaveBeenCalledWith(
+        { test: true },
+        [],
+      );
+      expect(callback).toHaveBeenCalled();
+    });
+
+    it("should convert buffers to DataViews", async () => {
+      const buffer = new ArrayBuffer(8);
+      await model.send({ test: true }, undefined, [buffer]);
+
+      expect(mockComm.sendCustomMessage).toHaveBeenCalledWith({ test: true }, [
+        expect.any(DataView),
+      ]);
     });
   });
 
   describe("widget_manager", () => {
-    const childModelId = "test-id";
-    const childModel = new Model({ foo: "test" }, vi.fn(), vi.fn(), new Set());
+    const childModelId = asModelId("test-id");
+    const childModel = new Model({ foo: "test" }, createMockComm());
     const manager = new ModelManager(10);
     let previousModelManager = Model._modelManager;
 
@@ -177,9 +220,9 @@ describe("Model", () => {
     });
 
     it("should throw error when accessing a model that is not registered", async () => {
-      await expect(model.widget_manager.get_model("random-id")).rejects.toThrow(
-        "Model not found for key: random-id",
-      );
+      await expect(
+        model.widget_manager.get_model(asModelId("random-id")),
+      ).rejects.toThrow("Model not found for key: random-id");
     });
 
     it("should return the registered model", async () => {
@@ -194,7 +237,7 @@ describe("Model", () => {
       const callback = vi.fn();
       model.on("change:foo", callback);
 
-      model.updateAndEmitDiffs({ foo: "test", bar: 456 });
+      getMarimoInternal(model).updateAndEmitDiffs({ foo: "test", bar: 456 });
       expect(callback).not.toHaveBeenCalled(); // foo didn't change
       expect(model.get("bar")).toBe(456);
     });
@@ -202,47 +245,33 @@ describe("Model", () => {
     it("should update and emit for deep changes", () => {
       const modelWithObject = new Model<{ foo: { nested: string } }>(
         { foo: { nested: "test" } },
-        onChange,
-        sendToWidget,
-        new Set(),
+        createMockComm(),
       );
       const callback = vi.fn();
       modelWithObject.on("change:foo", callback);
 
-      modelWithObject.updateAndEmitDiffs({ foo: { nested: "changed" } });
+      getMarimoInternal(modelWithObject).updateAndEmitDiffs({
+        foo: { nested: "changed" },
+      });
       expect(callback).toHaveBeenCalledTimes(1);
     });
 
     it("should emit change event for any changes", async () => {
       const callback = vi.fn();
       model.on("change", callback);
-      model.updateAndEmitDiffs({ foo: "changed", bar: 456 });
+      getMarimoInternal(model).updateAndEmitDiffs({ foo: "changed", bar: 456 });
       await TestUtils.nextTick(); // flush
       expect(callback).toHaveBeenCalledTimes(1);
     });
   });
 
-  describe("receiveCustomMessage", () => {
-    it("should handle update messages", () => {
-      model.receiveCustomMessage({
-        method: "update",
-        state: {
-          foo: "updated",
-          bar: 789,
-        },
-        buffer_paths: [],
-      });
-
-      expect(model.get("foo")).toBe("updated");
-      expect(model.get("bar")).toBe(789);
-    });
-
+  describe("emitCustomMessage", () => {
     it("should handle custom messages", () => {
       const callback = vi.fn();
       model.on("msg:custom", callback);
 
       const content = { type: "test" };
-      model.receiveCustomMessage({
+      getMarimoInternal(model).emitCustomMessage({
         method: "custom",
         content,
       });
@@ -256,7 +285,7 @@ describe("Model", () => {
 
       const content = { type: "test" };
       const buffer = new DataView(new ArrayBuffer(8));
-      model.receiveCustomMessage(
+      getMarimoInternal(model).emitCustomMessage(
         {
           method: "custom",
           content,
@@ -266,108 +295,85 @@ describe("Model", () => {
 
       expect(callback).toHaveBeenCalledWith(content, [buffer]);
     });
-
-    it("should log error for invalid messages", () => {
-      const consoleSpy = vi.spyOn(console, "error").mockImplementation(() => {
-        // noop
-      });
-      model.receiveCustomMessage({ invalid: "message" });
-
-      expect(consoleSpy).toHaveBeenCalledTimes(2);
-    });
   });
 });
 
 describe("ModelManager", () => {
   let modelManager = new ModelManager(50);
-  const handle = ({
-    modelId,
-    message,
-    buffers,
-  }: {
-    modelId: string;
-    message: AnyWidgetMessage;
-    buffers: readonly DataView[];
-  }) => {
-    return handleWidgetMessage({
-      modelId,
-      msg: message,
-      buffers,
-      modelManager,
-    });
-  };
+  const testId = asModelId("test-id");
 
   beforeEach(() => {
     // Clear the model manager before each test
     modelManager = new ModelManager(50);
+    mockSendModelValue.mockClear();
   });
 
   it("should set and get models", async () => {
-    const model = new Model({ count: 0 }, vi.fn(), vi.fn(), new Set());
-    modelManager.set("test-id", model);
-    const retrievedModel = await modelManager.get("test-id");
+    const model = new Model({ count: 0 }, createMockComm());
+    modelManager.set(testId, model);
+    const retrievedModel = await modelManager.get(testId);
     expect(retrievedModel).toBe(model);
   });
 
   it("should handle model not found", async () => {
-    await expect(modelManager.get("non-existent")).rejects.toThrow(
+    await expect(modelManager.get(asModelId("non-existent"))).rejects.toThrow(
       "Model not found for key: non-existent",
     );
   });
 
   it("should delete models", async () => {
-    const model = new Model({ count: 0 }, vi.fn(), vi.fn(), new Set());
-    modelManager.set("test-id", model);
-    modelManager.delete("test-id");
-    await expect(modelManager.get("test-id")).rejects.toThrow();
+    const model = new Model({ count: 0 }, createMockComm());
+    modelManager.set(testId, model);
+    modelManager.delete(testId);
+    await expect(modelManager.get(testId)).rejects.toThrow();
   });
 
   it("should handle widget messages", async () => {
-    const openMessage: AnyWidgetMessage = {
-      method: "open",
-      state: { count: 0 },
-      buffer_paths: [],
-    };
-
-    await handle({ modelId: "test-id", message: openMessage, buffers: [] });
-    const model = await modelManager.get("test-id");
+    await handleWidgetMessage(modelManager, {
+      model_id: testId,
+      message: {
+        method: "open",
+        state: { count: 0 },
+        buffer_paths: [],
+        buffers: [],
+      },
+    });
+    const model = await modelManager.get(testId);
     expect(model.get("count")).toBe(0);
 
-    const updateMessage: AnyWidgetMessage = {
-      method: "update",
-      state: {
-        count: 1,
+    await handleWidgetMessage(modelManager, {
+      model_id: testId,
+      message: {
+        method: "update",
+        state: { count: 1 },
+        buffer_paths: [],
+        buffers: [],
       },
-      buffer_paths: [],
-    };
-
-    await handle({ modelId: "test-id", message: updateMessage, buffers: [] });
+    });
     expect(model.get("count")).toBe(1);
   });
 
   it("should handle custom messages", async () => {
-    const model = new Model({ count: 0 }, vi.fn(), vi.fn(), new Set());
+    const model = new Model({ count: 0 }, createMockComm());
     const callback = vi.fn();
     model.on("msg:custom", callback);
-    modelManager.set("test-id", model);
+    modelManager.set(testId, model);
 
-    await handle({
-      modelId: "test-id",
-      message: { method: "custom", content: { count: 1 } },
-      buffers: [],
+    await handleWidgetMessage(modelManager, {
+      model_id: testId,
+      message: { method: "custom", content: { count: 1 }, buffers: [] },
     });
     expect(callback).toHaveBeenCalledWith({ count: 1 }, []);
   });
 
   it("should handle close messages", async () => {
-    const model = new Model({ count: 0 }, vi.fn(), vi.fn(), new Set());
-    modelManager.set("test-id", model);
+    const model = new Model({ count: 0 }, createMockComm());
+    modelManager.set(testId, model);
 
-    await handle({
-      modelId: "test-id",
+    await handleWidgetMessage(modelManager, {
+      model_id: testId,
       message: { method: "close" },
-      buffers: [],
     });
-    await expect(modelManager.get("test-id")).rejects.toThrow();
+    await expect(modelManager.get(testId)).rejects.toThrow();
   });
 });

package/src/{utils/__tests__/data-views.test.ts → plugins/impl/anywidget/__tests__/serialization.test.ts}

@@ -1,12 +1,8 @@
 /* Copyright 2026 Marimo. All rights reserved. */
 import { describe, expect, it } from "vitest";
-import {
-  decodeFromWire,
-  isWireFormat,
-  serializeBuffersToBase64,
-  type WireFormat,
-} from "../data-views";
-import type { Base64String } from "../json/base64";
+import type { WireFormat } from "@/plugins/impl/anywidget/types";
+import type { Base64String } from "../../../../utils/json/base64";
+import { decodeFromWire, serializeBuffersToBase64 } from "../serialization";
 
 describe("decodeFromWire with DataViews", () => {
   it("should return the original state if bufferPaths.length === 0", () => {
@@ -108,7 +104,7 @@ describe("Immutability Tests", () => {
     expect(input.array).toBe(input.array); // Array reference unchanged
   });
 
-  it("decodeFromWire (wire format) should not mutate input", () => {
+  it("decodeFromWire (wire format) should decode base64 to DataView", () => {
     const encoder = new TextEncoder();
     const data = encoder.encode("Hello");
     const base64 = btoa(String.fromCharCode(...data)) as Base64String;
@@ -119,19 +115,16 @@ describe("Immutability Tests", () => {
       buffers: [base64],
     };
 
-    const clone = structuredClone(input);
-
-    decodeFromWire(input);
-
-    // Check deep equality
-    expect(input).toEqual(clone);
+    const result = decodeFromWire(input);
 
-    // Check references are unchanged
-    expect(input.state).toBe(input.state);
-    expect(input.state.nested).toBe(input.state.nested);
+    // Result should have DataView instead of base64 string
+    expect(result.text).toBeInstanceOf(DataView);
+    // Other properties should be preserved
+    expect(result.number).toBe(42);
+    expect(result.nested).toEqual({ value: "test" });
   });
 
-  it("decodeFromWire (with DataViews) should not mutate input", () => {
+  it("decodeFromWire (with DataViews) should insert buffers at paths", () => {
     const encoder = new TextEncoder();
     const dataView = new DataView(encoder.encode("data").buffer);
 
@@ -143,21 +136,21 @@ describe("Immutability Tests", () => {
     const bufferPaths = [["data"]];
     const buffers = [dataView];
 
-    const input = { state, bufferPaths, buffers };
-    const clone = structuredClone(input);
-
-    decodeFromWire(input);
-
-    // Check deep equality
-    expect(input).toEqual(clone);
-
-    // Check references are unchanged
-    expect(input.state.nested).toBe(state.nested);
-    expect(input.state.array).toBe(state.array);
-    expect(input.buffers?.[0]).toBe(dataView); // Buffer reference unchanged
+    const result = decodeFromWire({
+      state,
+      bufferPaths,
+      buffers,
+    }) as typeof state & { data: DataView };
+
+    // Result should have the DataView at the path
+    expect(result.data).toBe(dataView);
+    // Other properties should be preserved
+    expect(result.placeholder).toBe("value");
+    expect(result.nested).toEqual({ key: "test" });
+    expect(result.array).toEqual([1, 2, 3]);
   });
 
-  it("decodeFromWire should return new object, not mutate", () => {
+  it("decodeFromWire should insert buffers and return state", () => {
     const encoder = new TextEncoder();
     const dataView = new DataView(encoder.encode("data").buffer);
 
@@ -165,16 +158,15 @@ describe("Immutability Tests", () => {
     const bufferPaths = [["c"]];
     const buffers = [dataView];
 
-    const result = decodeFromWire({ state, bufferPaths, buffers });
-
-    // Result should be different reference
-    expect(result).not.toBe(state);
-
-    // Input should not have new property
-    expect("c" in state).toBe(false);
+    const result = decodeFromWire({
+      state,
+      bufferPaths,
+      buffers,
+    }) as typeof state & { c: DataView };
 
     // Result should have new property
     expect("c" in result).toBe(true);
+    expect(result.c).toBe(dataView);
   });
 
   it("serializeBuffersToBase64 should return new state object", () => {
@@ -198,7 +190,7 @@ describe("Immutability Tests", () => {
     expect(typeof result.state.buffer).toBe("string");
   });
 
-  it("decodeFromWire with object input should not mutate state", () => {
+  it("decodeFromWire with object input should insert buffers", () => {
     const encoder = new TextEncoder();
     const dataView = new DataView(encoder.encode("data").buffer);
 
@@ -206,16 +198,17 @@ describe("Immutability Tests", () => {
     const bufferPaths = [["d"]];
     const buffers = [dataView];
 
-    const input = { state, bufferPaths, buffers };
-    const clone = structuredClone(input);
-
-    decodeFromWire(input);
-
-    // Check deep equality
-    expect(input).toEqual(clone);
-
-    // Check references are unchanged
-    expect(input.state.b).toBe(state.b);
+    const result = decodeFromWire({
+      state,
+      bufferPaths,
+      buffers,
+    }) as typeof state & { d: DataView };
+
+    // Result should have the buffer inserted
+    expect(result.d).toBe(dataView);
+    // Original properties should be preserved
+    expect(result.a).toBe(1);
+    expect(result.b).toEqual({ c: 2 });
   });
 
   it("nested objects should maintain independence after serialization", () => {
@@ -342,53 +335,6 @@ describe("serializeBuffersToBase64", () => {
   });
 });
 
-describe("isWireFormat", () => {
-  it("should return true for valid wire format", () => {
-    const wireFormat: WireFormat = {
-      state: { a: 1 },
-      bufferPaths: [],
-      buffers: [],
-    };
-    expect(isWireFormat(wireFormat)).toBe(true);
-  });
-
-  it("should return true for wire format with data", () => {
-    const wireFormat: WireFormat = {
-      state: { value: "SGVsbG8=" },
-      bufferPaths: [["value"]],
-      buffers: ["SGVsbG8=" as Base64String],
-    };
-    expect(isWireFormat(wireFormat)).toBe(true);
-  });
-
-  it("should return false for null", () => {
-    expect(isWireFormat(null)).toBe(false);
-  });
-
-  it("should return false for non-objects", () => {
-    expect(isWireFormat("string")).toBe(false);
-    expect(isWireFormat(123)).toBe(false);
-    expect(isWireFormat(true)).toBe(false);
-    expect(isWireFormat(undefined)).toBe(false);
-  });
-
-  it("should return false when missing state", () => {
-    expect(isWireFormat({ bufferPaths: [], buffers: [] })).toBe(false);
-  });
-
-  it("should return false when missing bufferPaths", () => {
-    expect(isWireFormat({ state: {}, buffers: [] })).toBe(false);
-  });
-
-  it("should return false when missing buffers", () => {
-    expect(isWireFormat({ state: {}, bufferPaths: [] })).toBe(false);
-  });
-
-  it("should return false for plain objects", () => {
-    expect(isWireFormat({ a: 1, b: 2 })).toBe(false);
-  });
-});
-
 describe("decodeFromWire from WireFormat", () => {
   it("should return state unchanged when no buffer paths", () => {
     const wire: WireFormat = {