feat(experimental): Add command system for calling Python from JS

[manzt]

This PR introduces an experimental API designed to make sending custom messages easier in the front end. This is achieved through a new RPC mechanism, which allows the frontend to invoke a command registered in Python and await the response.

A minimal example:

import anywidget
import traitlets

class Widget(anywidget.AnyWidget):
    _esm = """
    async function render({ model, el, experimental }) {
      let [msg, buffers] = await experimental.invoke("_echo", "hello");
      console.log(msg) // "hello"
    }
    export default { render };
    """
    
    @anywidget.experimental.command
    def _echo(self, msg, buffers):
        return msg, buffers

this can make implementing things with dynamic data much easier, like the ZarrWidget:

import zarr
import zarr.storage

store = zarr.storage.MemoryStore()
grp = zarr.open(store)
grp.create_dataset("foo", shape=[10, 10])

ZarrWidget(store=store)

ZarrWidget implementation

import anywidget
import traitlets

class ZarrWidget(anywidget.AnyWidget):
    _esm = """
    import * as zarr from "https://esm.sh/zarrita@next";
    
    function render({ model, el, experimental }) {
      let root = zarr.root({
        async get(key) {
          const [data, buffers] = await experimental.invoke("_zarr_get", key);
          if (!data.success) return undefined;
          return buffers[0].buffer;
        },
      });
    
      let pre = document.createElement("pre");

      let btn = document.createElement("button");
      btn.innerHTML = `click to load`;
      btn.addEventListener("click", async () => {
        const node = await zarr.open(root.resolve("/foo"), { kind: "array" });
        const { data, shape, stride } = await zarr.get(node);
        pre.innerText = JSON.stringify({ data: data.constructor.name, shape, stride });
      });
      el.append(btn, pre);
    }
    export default { render };
    """
    def __init__(self, store, **kwargs):
        super().__init__(**kwargs)
        self._store = store

    @anywidget.experimental.command
    def _zarr_get(self, key, buffers):
        try:
             buffers = [store[key.lstrip("/")]]
        except KeyError:
             buffers = []
        return { "success": len(buffers) == 1 }, buffers

cc: @keller-mark

TODO:

• Add tests
• Support sending binary data from invoke (via another argument)
• Documentation

Another option would be to move this kind of dispatching mechanism to a separate library (e.g., @anywidget/std), but, since it is coupled to the Python, I prefer we implement something within anywidget that makes this pattern frictionless for applications that need it.

EDIT: This PR has changed from a messaging callback to registering methods like RPC.

[netlify]

✅ Deploy Preview for anywidget ready!

Name	Link
🔨 Latest commit	9617127
🔍 Latest deploy log	https://app.netlify.com/sites/anywidget/deploys/65e9ccc6a35b8f000882198a
😎 Deploy Preview	https://deploy-preview-453--anywidget.netlify.app
📱 Preview on mobile	Toggle QR Code... Use your smartphone camera to open QR code link.

---

To edit notification comments on pull requests, go to your Netlify site configuration.

[keller-mark]

This is awesome! Very nice how the on_msg only gets called if the user has defined the _anywidget_experimental_reducer

[manzt]

Very nice how the on_msg only gets called if the user has defined the _anywidget_experimental_reducer

Thanks! Yeah, I think having something opinionated but opt-in would probably be preferable. I renamed to _experimental_anywidget_reducer in anticipation that we would remove the leading _experimental when stabilized.

[kylebarron]

I'm still curious whether the Python callback can be made async or whether that messes with the Jupyter event loop

[manzt]

I've been reviewing this, and I'm realizing this doesn't really follow the redux pattern. In Redux, dispatching emits and action, which updates the state, and then view updates according to state changes. Therefore, model.send is much closer conceptually to Redux-y patterns (i.e., emit a message to Python, Python updates the model, JS view responds to changes on the model).

This PR is narrowing in on something different but very important: RPC. So I wonder about taking this one step further, and baking in a mechanism like Tauri's invoke. Essentially, I wonder about something like:

class Widget(anywidget.AnyWidget):
    _esm = """
    async function render({ model, el, experimental }) {
      let [msg, buffers] = await experimental.invoke("_echo", "hello, world");
    }
    """
   @anywidget.command
   def _echo(msg, buffers):
       return msg, buffers

I think this is the direction I'd like to take this PR, what do you think @kylebarron @keller-mark ?

The previous API would still be possible, just register one "command" that handles many different types of messages. This would make the most simple case (a single callback) very easy. It would be nice if we could have something other than msg/buffers, I suppose we could require "msg" to be an dict and **msg when invoking the callback... but then not sure what to do about buffers.

A more opinionated version would be to use something like msgpack for serializing args and responses so callers wouldn't need to worry about handling msg/buffers themselves:

class Widget(anywidget.AnyWidget):
    _esm = """
    async function render({ model, el, experimental }) {
      let response = await experimental.invoke("_echo", {
        text: "hello, world",
        btext: new TextEncoder().encode("hello, world"),
      });
    }
    """
   @anywidget.command
   def _echo(text: str, btext: bytes):
       return { "text": text, "btext": btext }

Probably overkill, and something we could optionally implement ontop of the core mechanism in the future.

[keller-mark]

I think framing it more like RPC conceptually makes sense. I do not have a super strong preference towards dispatch or invoke as long as it allows the ZarrWidget to be implemented.

Is the @anywidget.command decorator necessary? Maybe it helps with security or future-proof things to know which functions are potentially invoked?

It would be nice if we could have something other than msg/buffers, I suppose we could require "msg" to be an dict and **msg when invoking the callback... but then not sure what to do about buffers.

I don't really follow this point but i am not familiar with msgpack. Is the issue that you want the API to be

-let [msg] = await experimental.invoke("_echo", "hello, world");
+let msg = await experimental.invoke("_echo", "hello, world");

when the user does not care about buffers?

[manzt]

do not have a super strong preference towards dispatch or invoke as long as it allows the ZarrWidget to be implemented.

Right, makes sense. I guess I should just stop bike-shedding and merge the thing :)

Is the @anywidget.command decorator necessary?

This decorator "marks" the methods on the class to "expose" to the front end. I guess we could try to inspect the signature of all the methods on the class and try to guess which ones implement the invoke signature:

(msg: Any, buffers: list[bytes]) -> tuple[Any, list[bytes]]

but this is more explicit and we can raise type errors if the wrapped method has the wrong signature. It also creates a future space where we could extend the behavior (see below).

I don't really follow this point but i am not familiar with msgpack. Is the issue that you want the API to be.

If you wanted to return a dict of buffers currently (for example), using getitems from zarr-python:

def getitems(keys: list[str]) -> dict[str, bytes] ... 

You would need to extract the buffers manually, both on the JS and Python side:

@anywidget.command
def _getitems(self, msg: Any, buffers: list[bytes]) -> tuple[list[str], list[bytes]]:
    keys = msg["keys"]
    items = self.store.getitems(keys)
    return list(items.keys()), list(items.values())

If this was a traitlet, Jupyter Widgets would automatically take care of "extracting" the buffers and assembling the deconstructed object on the JS side. But for custom messages, you need to do it yourself. I guess I was wondering if this is a place we could further make developing widgets easier, so developers didn't need to worry about packing/unpacking buffers.

One option would be to reuse the packing mechanism from Jupyter Widgets, or a more well defined binary protocol like msgpack. My current thinking is that we could extend command to allow this as an option (not for this PR):

@anywidget.command(serialize=None)
def _getitems(self, msg: Any, buffers: list[bytes]) -> tuple[list[str], list[bytes]]:
    keys = msg["keys"]
    items = self.store.getitems(keys)
    return list(items.keys()), list(items.values())

@anywidget.command(serialize="msgpack") # serialize="jupyter-widgets"
def _getitems(self, keys: list[str]) -> dict[str, bytes]:
    keys = msg["keys"]
    return self.store.getitems(keys)

Since serialization has some overhead, I think it makes sense to have the [msg, buffers] version always available... but I'm fairly confident the later would reduce some gotchas. Again, probably something to evaluate later after seeing this used in the wild.

[changeset-bot]

🦋 Changeset detected

Latest commit: 568e243

The changes in this PR will be included in the next version bump.

This PR includes changesets to release 4 packages

Name	Type
anywidget	Patch
@anywidget/types	Patch
@anywidget/react	Patch
@anywidget/svelte	Patch

Not sure what this means? Click here to learn what changesets are.

Click here if you're a maintainer who wants to add another changeset to this PR

[kylebarron]

I agree that an RPC-based API makes sense for this, and also +1 on exploring msgpack (at least as an option). I haven't used it myself but I'm guessing the client libraries for that are pretty stable and minimal

[manzt]

Awesome. I'm going to push forward and hopefully cut a release this week!

[manzt]

I haven't used it myself but I'm guessing the client libraries for that are pretty stable and minimal

Yes, I'm familiar with msgpack for neovim's extension mechanism. On the Python side there is msgspec, which is minimal and speedy, and on the JS side @msgpack/msgpack. Maybe @jcrist would have some thoughts on how appropriate it is in this context.

[keller-mark]

Thanks for the clarification, the ability to specify in the decorator whether the msgpack serialization should be used looks like a great way forward!

[manzt]

Cool, let's merge the thing! I will work on some docs in a separate PR, but they will be a draft likely until we stabilize the feature.

cc: @mscolnick, just FYI a new API we are iterating on. It would be great to learn if this is something that you think Marimo could support. Relevant code:

anywidget/packages/anywidget/src/widget.js

Lines 241 to 283 in 568e243

	/**
	* @template T
	* @param {import("@anywidget/types").AnyModel} model
	* @param {string} name
	* @param {any} [msg]
	* @param {DataView[]} [buffers]
	* @param {{ timeout?: number }} [options]
	* @return {Promise<[T, DataView[]]>}
	*/
	export function invoke(
	model,
	name,
	msg,
	buffers = [],
	{ timeout = 3000 } = {},
	) {
	// crypto.randomUUID() is not available in non-secure contexts (i.e., http://)
	// so we use simple (non-secure) polyfill.
	let id = uuid.v4();
	return new Promise((resolve, reject) => {
	let timer = setTimeout(() => {
	reject(new Error(`Promise timed out after ${timeout} ms`));
	model.off("msg:custom", handler);
	}, timeout);
	/**
	* @param {{ id: string, kind: "anywidget-command-response", response: T }} msg
	* @param {DataView[]} buffers
	*/
	function handler(msg, buffers) {
	if (!(msg.id === id)) return;
	clearTimeout(timer);
	resolve([msg.response, buffers]);
	model.off("msg:custom", handler);
	}
	model.on("msg:custom", handler);
	model.send(
	{ id, kind: "anywidget-command", name, msg },
	undefined,
	buffers,
	);
	});
	}

[mscolnick]

@manzt very cool addition! we can support this (we have a concept of RPC/Functions too that I can adapt to work with this as well.

Is there a timeline for when it will be out of experimental? Would like to avoid 2 changes if possible, but not against it.