@marimo-team/frontend 0.19.8-dev41 → 0.19.8-dev49

package/src/components/editor/ai/ai-completion-editor.tsx

@@ -3,10 +3,8 @@
 import { useCompletion } from "@ai-sdk/react";
 import { EditorView } from "@codemirror/view";
 import {
-  AtSignIcon,
   CircleCheckIcon,
   Loader2Icon,
-  SendIcon,
   SparklesIcon,
   XIcon,
 } from "lucide-react";
@@ -20,6 +18,10 @@ import { storePrompt } from "@marimo-team/codemirror-ai";
 import type { ReactCodeMirrorRef } from "@uiw/react-codemirror";
 import { useAtom, useAtomValue } from "jotai";
 import { AIModelDropdown } from "@/components/ai/ai-model-dropdown";
+import {
+  AddContextButton,
+  SendButton,
+} from "@/components/chat/chat-components";
 import { Checkbox } from "@/components/ui/checkbox";
 import { Label } from "@/components/ui/label";
 import { Switch } from "@/components/ui/switch";
@@ -253,39 +255,6 @@ export const AiCompletionEditor: React.FC<Props> = ({
     }
   };
 
-  const loadingStopButton = (
-    <Button
-      data-testid="stop-completion-button"
-      variant="text"
-      size="xs"
-      className="mb-0"
-      onClick={stop}
-    >
-      <Loader2Icon className="animate-spin mr-1" size={14} />
-      Stop
-    </Button>
-  );
-
-  const submitButton = (
-    <Tooltip content="Submit">
-      <Button variant="text" size="icon" onClick={handleSubmit}>
-        <SendIcon className="h-3 w-3" />
-      </Button>
-    </Tooltip>
-  );
-
-  const contextButton = (
-    <Tooltip content="Add context">
-      <Button
-        variant="text"
-        size="icon"
-        onClick={() => addContextCompletion(inputRef)}
-      >
-        <AtSignIcon className="h-3 w-3" />
-      </Button>
-    </Tooltip>
-  );
-
   const completionButtons = (
     <>
       <AcceptCompletionButton
@@ -355,9 +324,17 @@ export const AiCompletionEditor: React.FC<Props> = ({
 
             <div className="-mr-1.5 py-1.5">
               <div className="flex flex-row items-center justify-end gap-0.5">
-                {isLoading && loadingStopButton}
-                {submitButton}
-                {contextButton}
+                <SendButton
+                  isLoading={isLoading}
+                  onStop={stop}
+                  onSendClick={handleSubmit}
+                  isEmpty={!input.trim()}
+                  showStopLabel={true}
+                />
+                <AddContextButton
+                  handleAddContext={() => addContextCompletion(inputRef)}
+                  isLoading={isLoading}
+                />
                 <AIModelDropdown
                   triggerClassName="h-7 text-xs"
                   iconSize="small"

package/src/components/editor/chrome/panels/packages-panel.tsx

@@ -349,8 +349,9 @@ const PackagesList: React.FC<{
 
 const UpgradeButton: React.FC<{
   packageName: string;
+  tags?: { kind: string; value: string }[];
   onSuccess: () => void;
-}> = ({ packageName, onSuccess }) => {
+}> = ({ packageName, tags, onSuccess }) => {
   const [loading, setLoading] = React.useState(false);
   const { addPackage } = useRequestClient();
 
@@ -362,9 +363,11 @@ const UpgradeButton: React.FC<{
   const handleUpgradePackage = async () => {
     try {
       setLoading(true);
+      const group = tags?.find((tag) => tag.kind === "group")?.value;
       const response = await addPackage({
         package: packageName,
         upgrade: true,
+        group,
       });
       if (response.success) {
         onSuccess();
@@ -395,12 +398,10 @@ const RemoveButton: React.FC<{
   const handleRemovePackage = async () => {
     try {
       setLoading(true);
-      const isDev = tags?.some(
-        (tag) => tag.kind === "group" && tag.value === "dev",
-      );
+      const group = tags?.find((tag) => tag.kind === "group")?.value;
       const response = await removePackage({
         package: packageName,
-        dev: isDev,
+        group,
       });
       if (response.success) {
         onSuccess();
@@ -604,12 +605,14 @@ const DependencyTreeNode: React.FC<{
         {/* Actions for top-level packages */}
         {isTopLevel && (
           <div className="flex gap-1 invisible group-hover:visible">
-            <UpgradeButton packageName={node.name} onSuccess={onSuccess} />
+            <UpgradeButton
+              packageName={node.name}
+              tags={node.tags}
+              onSuccess={onSuccess}
+            />
+
             <RemoveButton
               packageName={node.name}
-              // FIXME: Backend types are wrong/outdated.
-              // tags actually have the shape: Array<{ kind: string; value: string }>
-              // @ts-expect-error — backend tag types do not match frontend expectations yet
               tags={node.tags}
               onSuccess={onSuccess}
             />

package/src/core/islands/__tests__/bridge.test.ts

@@ -192,9 +192,11 @@ describe("IslandsPyodideBridge", () => {
       const request = {
         modelId: "widget-1",
         message: {
+          method: "update" as const,
           state: { value: 42 },
           bufferPaths: [],
         },
+        buffers: [],
       };
 
       await bridge.sendModelValue(request);
@@ -202,12 +204,14 @@ describe("IslandsPyodideBridge", () => {
       expect(mockBridge).toHaveBeenCalledWith({
         functionName: "put_control_request",
         payload: {
-          type: "update-widget-model",
+          type: "model",
           modelId: "widget-1",
           message: {
+            method: "update",
             state: { value: 42 },
             bufferPaths: [],
           },
+          buffers: [],
         },
       });
     });
@@ -226,7 +230,8 @@ describe("IslandsPyodideBridge", () => {
       await bridge.sendRun({ cellIds: [], codes: [] });
       await bridge.sendModelValue({
         modelId: "",
-        message: { state: {}, bufferPaths: [] },
+        message: { method: "update", state: {}, bufferPaths: [] },
+        buffers: [],
       });
 
       // All calls should have the type field

package/src/core/islands/bridge.ts

@@ -140,7 +140,7 @@ export class IslandsPyodideBridge implements RunRequests, EditRequests {
 
   sendModelValue: RunRequests["sendModelValue"] = async (request) => {
     await this.putControlRequest({
-      type: "update-widget-model",
+      type: "model",
       ...request,
     });
     return null;

package/src/core/islands/main.ts

@@ -13,6 +13,10 @@ import "iconify-icon";
 
 import { toast } from "@/components/ui/use-toast";
 import { renderHTML } from "@/plugins/core/RenderHTML";
+import {
+  handleWidgetMessage,
+  MODEL_MANAGER,
+} from "@/plugins/impl/anywidget/model";
 import { initializePlugins } from "@/plugins/plugins";
 import { logNever } from "@/utils/assertNever";
 import { Functions } from "@/utils/functions";
@@ -194,6 +198,9 @@ export async function initialize() {
         return;
       case "kernel-startup-error":
         return;
+      case "model-lifecycle":
+        handleWidgetMessage(MODEL_MANAGER, msg.data);
+        return;
       default:
         logNever(msg.data);
     }

package/src/core/network/types.ts

@@ -78,7 +78,7 @@ export interface SetCellConfigRequest {
   configs: Record<CellId, Partial<CellConfig>>;
 }
 export type UpdateUIElementRequest = schemas["UpdateUIElementRequest"];
-export type UpdateWidgetModelRequest = schemas["UpdateWidgetModelRequest"];
+export type ModelRequest = schemas["ModelRequest"];
 export type UpdateCellIdsRequest = schemas["UpdateCellIdsRequest"];
 export type UpdateUserConfigRequest = schemas["UpdateUserConfigRequest"];
 export type ShutdownSessionRequest = schemas["ShutdownSessionRequest"];
@@ -110,7 +110,7 @@ export type LspServerHealth = schemas["LspServerHealth"];
  */
 export interface RunRequests {
   sendComponentValues: (request: UpdateUIElementValuesRequest) => Promise<null>;
-  sendModelValue: (request: UpdateWidgetModelRequest) => Promise<null>;
+  sendModelValue: (request: ModelRequest) => Promise<null>;
   sendInstantiate: (request: InstantiateNotebookRequest) => Promise<null>;
   sendFunctionRequest: (request: InvokeFunctionRequest) => Promise<null>;
 }

package/src/core/wasm/bridge.ts

@@ -541,7 +541,7 @@ export class PyodideBridge implements RunRequests, EditRequests {
 
   sendModelValue: RunRequests["sendModelValue"] = async (request) => {
     await this.putControlRequest({
-      type: "update-widget-model",
+      type: "model",
       ...request,
     });
     return null;

package/src/core/websocket/useMarimoKernelConnection.tsx

@@ -11,7 +11,6 @@ import { useConnectionTransport } from "@/core/websocket/useWebSocket";
 import { renderHTML } from "@/plugins/core/RenderHTML";
 import {
   handleWidgetMessage,
-  isMessageWidgetState,
   MODEL_MANAGER,
 } from "@/plugins/impl/anywidget/model";
 import { logNever } from "@/utils/assertNever";
@@ -141,31 +140,22 @@ export function useMarimoKernelConnection(opts: {
         return;
 
       case "send-ui-element-message": {
-        const modelId = msg.data.model_id;
         const uiElement = msg.data.ui_element;
-        const message = msg.data.message;
-        const buffers = safeExtractSetUIElementMessageBuffers(msg.data);
-
-        if (modelId && isMessageWidgetState(message)) {
-          handleWidgetMessage({
-            modelId,
-            msg: message,
-            buffers,
-            modelManager: MODEL_MANAGER,
-          });
-        }
-
         if (uiElement) {
+          const buffers = safeExtractSetUIElementMessageBuffers(msg.data);
           UI_ELEMENT_REGISTRY.broadcastMessage(
             uiElement as UIElementId,
             msg.data.message,
             buffers,
           );
         }
-
         return;
       }
 
+      case "model-lifecycle":
+        handleWidgetMessage(MODEL_MANAGER, msg.data);
+        return;
+
       case "remove-ui-elements":
         handleRemoveUIElements(msg.data);
         return;

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

@@ -1,97 +1,50 @@
 /* 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 () => {
@@ -117,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;
     }
@@ -156,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`,
     );
@@ -189,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}
     />
   );
 };
@@ -200,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} />;
 };
@@ -327,5 +247,4 @@ export const visibleForTesting = {
   LoadedSlot,
   runAnyWidgetModule,
   isAnyWidgetModule,
-  getDirtyFields,
 };