Accept module plugs in Kino.Proxy.listen/1

lib/kino/bridge.ex

@@ -259,12 +259,12 @@ defmodule Kino.Bridge do
   end
 
   @doc """
-  Requests the child spec for proxy handler with the given function.
+  Requests the child spec for proxy handler with the given plug.
   """
-  @spec get_proxy_handler_child_spec((Plug.Conn.t() -> Plug.Conn.t())) ::
+  @spec get_proxy_handler_child_spec(Kino.Proxy.plug()) ::
           {:ok, {module(), term()}} | request_error()
-  def get_proxy_handler_child_spec(fun) do
-    io_request({:livebook_get_proxy_handler_child_spec, fun})
+  def get_proxy_handler_child_spec(plug) do
+    io_request({:livebook_get_proxy_handler_child_spec, plug})
   end
 
   defp io_request(request) do

lib/kino/proxy.ex

@@ -72,15 +72,45 @@ defmodule Kino.Proxy do
   > ```
   """
 
+  @type plug() ::
+          (Plug.Conn.t() -> Plug.Conn.t())
+          | module()
+          | {module(), term()}
+
   @doc """
   Registers a request listener.
 
-  Expects the listener to be a function that handles a request
-  `Plug.Conn`.
+  Expects the listener to be a plug, that is, one of:
+
+    * a function plug: a `fun(conn)` function that takes a `Plug.Conn` and returns a `Plug.Conn`.
+
+    * a module plug: a `module` atom or a `{module, options}` tuple.
   """
-  @spec listen((Plug.Conn.t() -> Plug.Conn.t())) :: DynamicSupervisor.on_start_child()
-  def listen(fun) when is_function(fun, 1) do
-    case Kino.Bridge.get_proxy_handler_child_spec(fun) do
+  @spec listen(plug()) :: DynamicSupervisor.on_start_child()
+  def listen(plug) do
+    case plug do
+      fun when is_function(fun, 1) ->
+        :ok
+
+      mod when is_atom(mod) ->
+        :ok
+
+      {mod, _opts} when is_atom(mod) ->
+        :ok
+
+      other ->
+        raise """
+        expected plug to be one of:
+
+          * fun(conn)
+          * module
+          * {module, options}
+
+        got: #{inspect(other)}
+        """
+    end
+
+    case Kino.Bridge.get_proxy_handler_child_spec(plug) do
       {:ok, child_spec} ->
         Kino.start_child(child_spec)