Apply review comments

lib/kino/bridge.ex

@@ -250,6 +250,15 @@ defmodule Kino.Bridge do
     match?({:ok, _}, io_request(:livebook_get_evaluation_file))
   end
 
+  @doc """
+  Requests the child spec for proxy handler with the given function.
+  """
+  @spec get_proxy_handler_child_spec((Plug.Conn.t() -> Plug.Conn.t())) ::
+          Supervisor.child_spec() | {module(), term()} | module()
+  def get_proxy_handler_child_spec(fun) do
+    with {:ok, reply} <- io_request({:livebook_get_proxy_handler_child_spec, fun}), do: reply
+  end
+
   defp io_request(request) do
     gl = Process.group_leader()
     ref = Process.monitor(gl)

lib/kino/proxy.ex

@@ -1,22 +1,16 @@
 defmodule Kino.Proxy do
   @moduledoc """
-  A kino for handling proxy requests from the host.
+  Functionality for handling proxy requests forwarded from Livebook.
 
-  ## Examples
-
-      Kino.Proxy.listen(fn conn ->
-        conn
-        |> Plug.Conn.put_resp_header("content-type", "application/text;charset=utf-8")
-        |> Plug.Conn.send_resp(200, "used " <> conn.method <> " method")
-      end)
+  TODO: Write an extensive docs here.
   """
 
   @doc """
   Persists the function to be listened by the proxy handler.
   """
-  @spec listen(atom(), module(), (Plug.Conn.t() -> Plug.Conn.t())) ::
-          DynamicSupervisor.on_start_child()
-  def listen(name \\ __MODULE__, mod \\ Livebook.Proxy.Handler, fun) when is_function(fun, 1) do
-    Kino.start_child({mod, name: name, listen: fun})
+  @spec listen((Plug.Conn.t() -> Plug.Conn.t())) :: DynamicSupervisor.on_start_child()
+  def listen(fun) when is_function(fun, 1) do
+    child_spec = Kino.Bridge.get_proxy_handler_child_spec(fun)
+    Kino.start_child(child_spec)
   end
 end