Update

lib/kino/bridge.ex

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

lib/kino/proxy.ex

@@ -88,29 +88,32 @@ defmodule Kino.Proxy do
   """
   @spec listen(plug()) :: DynamicSupervisor.on_start_child()
   def listen(plug) do
-    case plug do
-      fun when is_function(fun, 1) ->
-        :ok
+    fun =
+      case plug do
+        fun when is_function(fun, 1) ->
+          fun
 
-      mod when is_atom(mod) ->
-        :ok
+        mod when is_atom(mod) ->
+          opts = mod.init([])
+          &mod.call(&1, opts)
 
-      {mod, _opts} when is_atom(mod) ->
-        :ok
+        {mod, opts} when is_atom(mod) ->
+          opts = mod.init(opts)
+          &mod.call(&1, opts)
 
-      other ->
-        raise """
-        expected plug to be one of:
+        other ->
+          raise """
+          expected plug to be one of:
 
-          * fun(conn)
-          * module
-          * {module, options}
+            * fun(conn)
+            * module
+            * {module, options}
 
-        got: #{inspect(other)}
-        """
-    end
+          got: #{inspect(other)}
+          """
+      end
 
-    case Kino.Bridge.get_proxy_handler_child_spec(plug) do
+    case Kino.Bridge.get_proxy_handler_child_spec(fun) do
       {:ok, child_spec} ->
         Kino.start_child(child_spec)