Introduce a dedicated channel for JS widget communication (#843)

* Introduce a dedicated channel for JS widget communication

* Handle payload serialization errors

* Tie channel lifetime to the session

* Catch serialization errors instead of encoding twice

* Merge JS static and dynamic outputs

* Authenticate socket connection from session

* Update JS output format

* Remove unused helper

* Apply review comments

assets/js/js_output/index.js

@@ -3,21 +3,44 @@ import { randomToken } from "../lib/utils";
 
 import iframeHtml from "./iframe.html";
 
+const global = {
+  channel: null,
+};
+
+// Returns channel responsible for JS communication in the current session
+function getChannel(socket, { create = true } = {}) {
+  if (!global.channel && create) {
+    global.channel = socket.channel("js_output", {});
+    global.channel.join();
+  }
+
+  return global.channel;
+}
+
+export function leaveChannel() {
+  if (global.channel) {
+    global.channel.leave();
+    global.channel = null;
+  }
+}
+
 /**
  * A hook used to render JS-enabled cell output.
  *
  * The JavaScript is defined by the user, so we sandbox the script
  * execution inside an iframe.
  *
- * The hook expects `js_output:<id>:init` event with `{ data }` payload,
- * the data is then used in the initial call to the custom JS module.
+ * The hook connects to a dedicated channel, sending the token and
+ * output ref in an initial message. It expects `init:<ref>` message
+ * with `{ data }` payload, the data is then used in the initial call
+ * to the custom JS module.
  *
- * Then, a number of `js_output:<id>:event` with `{ event }` payload can
- * be sent. The `event` is forwarded to the initialized component.
+ * Then, a number of `event:<ref>` with `{ event, payload }` payload
+ * can be sent. The `event` is forwarded to the initialized component.
  *
  * Configuration:
  *
- *   * `data-id` - a unique identifier used as messages scope
+ *   * `data-ref` - a unique identifier used as messages scope
  *
  *   * `data-assets-base-url` - the URL to resolve all relative paths
  *     against in the iframe
@@ -25,17 +48,24 @@ import iframeHtml from "./iframe.html";
  *   * `data-js-path` - a relative path for the initial output-specific
  *     JS module
  *
+ *   * `data-session-token` - token is sent in the "connect" message to
+ *     the channel
+ *
  */
 const JSOutput = {
   mounted() {
     this.props = getProps(this);
     this.state = {
-      token: randomToken(),
+      childToken: randomToken(),
       childReadyPromise: null,
       childReady: false,
       iframe: null,
+      channelUnsubscribe: null,
+      errorContainer: null,
     };
 
+    const channel = getChannel(this.__liveSocket.getSocket());
+
     const iframePlaceholder = document.createElement("div");
     const iframe = document.createElement("iframe");
     this.state.iframe = iframe;
@@ -57,7 +87,7 @@ const JSOutput = {
         if (message.type === "ready" && !this.state.childReady) {
           postMessage({
             type: "readyReply",
-            token: this.state.token,
+            token: this.state.childToken,
             baseUrl: this.props.assetsBaseUrl,
             jsPath: this.props.jsPath,
           });
@@ -71,7 +101,7 @@ const JSOutput = {
           // any of those messages, so we can treat this as a possible
           // surface for attacks. In this case the most "critical" actions
           // are shortcuts, neither of which is particularly dangerous.
-          if (message.token !== this.state.token) {
+          if (message.token !== this.state.childToken) {
             throw new Error("Token mismatch");
           }
 
@@ -85,7 +115,7 @@ const JSOutput = {
             this.el.dispatchEvent(event);
           } else if (message.type === "event") {
             const { event, payload } = message;
-            this.pushEvent("event", { event, payload });
+            channel.push("event", { event, payload, ref: this.props.ref });
           }
         }
       };
@@ -163,20 +193,41 @@ const JSOutput = {
 
     // Event handlers
 
-    this.handleEvent(`js_output:${this.props.id}:init`, ({ data }) => {
+    channel.push("connect", {
+      session_token: this.props.sessionToken,
+      ref: this.props.ref,
+    });
+
+    const initRef = channel.on(`init:${this.props.ref}`, ({ data }) => {
       this.state.childReadyPromise.then(() => {
         postMessage({ type: "init", data });
       });
     });
 
-    this.handleEvent(
-      `js_output:${this.props.id}:event`,
+    const eventRef = channel.on(
+      `event:${this.props.ref}`,
       ({ event, payload }) => {
         this.state.childReadyPromise.then(() => {
           postMessage({ type: "event", event, payload });
         });
       }
     );
+
+    const errorRef = channel.on(`error:${this.props.ref}`, ({ message }) => {
+      if (!this.state.errorContainer) {
+        this.state.errorContainer = document.createElement("div");
+        this.state.errorContainer.classList.add("error-box", "mb-4");
+        this.el.prepend(this.state.errorContainer);
+      }
+
+      this.state.errorContainer.textContent = message;
+    });
+
+    this.state.channelUnsubscribe = () => {
+      channel.off(`init:${this.props.ref}`, initRef);
+      channel.off(`event:${this.props.ref}`, eventRef);
+      channel.off(`error:${this.props.ref}`, errorRef);
+    };
   },
 
   updated() {
@@ -188,14 +239,24 @@ const JSOutput = {
     this.resizeObserver.disconnect();
     this.intersectionObserver.disconnect();
     this.state.iframe.remove();
+
+    const channel = getChannel(this.__liveSocket.getSocket(), {
+      create: false,
+    });
+
+    if (channel) {
+      this.state.channelUnsubscribe();
+      channel.push("disconnect", { ref: this.props.ref });
+    }
   },
 };
 
 function getProps(hook) {
   return {
-    id: getAttributeOrThrow(hook.el, "data-id"),
+    ref: getAttributeOrThrow(hook.el, "data-ref"),
     assetsBaseUrl: getAttributeOrThrow(hook.el, "data-assets-base-url"),
     jsPath: getAttributeOrThrow(hook.el, "data-js-path"),
+    sessionToken: getAttributeOrThrow(hook.el, "data-session-token"),
   };
 }
 

assets/js/session/index.js

@@ -11,6 +11,7 @@ import { getAttributeOrDefault } from "../lib/attribute";
 import KeyBuffer from "./key_buffer";
 import { globalPubSub } from "../lib/pub_sub";
 import monaco from "../cell/live_editor/monaco";
+import { leaveChannel } from "../js_output";
 
 /**
  * A hook managing the whole session.
@@ -244,6 +245,8 @@ const Session = {
     document.removeEventListener("dblclick", this.handleDocumentDoubleClick);
 
     setFavicon("favicon");
+
+    leaveChannel();
   },
 };
 

lib/livebook/live_markdown/export.ex

@@ -14,15 +14,45 @@ defmodule Livebook.LiveMarkdown.Export do
   """
   @spec notebook_to_markdown(Notebook.t(), keyword()) :: String.t()
   def notebook_to_markdown(notebook, opts \\ []) do
-    ctx = %{
-      include_outputs?: Keyword.get(opts, :include_outputs, notebook.persist_outputs)
-    }
+    include_outputs? = Keyword.get(opts, :include_outputs, notebook.persist_outputs)
+
+    js_ref_with_data = if include_outputs?, do: collect_js_output_data(notebook), else: %{}
+
+    ctx = %{include_outputs?: include_outputs?, js_ref_with_data: js_ref_with_data}
 
     iodata = render_notebook(notebook, ctx)
     # Add trailing newline
     IO.iodata_to_binary([iodata, "\n"])
   end
 
+  defp collect_js_output_data(notebook) do
+    for section <- notebook.sections,
+        %Cell.Elixir{} = cell <- section.cells,
+        {:js, %{export: %{}, ref: ref, pid: pid}} <- cell.outputs do
+      Task.async(fn ->
+        {ref, get_js_output_data(pid, ref)}
+      end)
+    end
+    |> Task.await_many(:infinity)
+    |> Map.new()
+  end
+
+  defp get_js_output_data(pid, ref) do
+    send(pid, {:connect, self(), %{origin: self(), ref: ref}})
+
+    monitor_ref = Process.monitor(pid)
+
+    data =
+      receive do
+        {:connect_reply, data, %{ref: ^ref}} -> data
+        {:DOWN, ^monitor_ref, :process, _pid, _reason} -> nil
+      end
+
+    Process.demonitor(monitor_ref, [:flush])
+
+    data
+  end
+
   defp render_notebook(notebook, ctx) do
     comments =
       Enum.map(notebook.leading_comments, fn
@@ -94,7 +124,7 @@ defmodule Livebook.LiveMarkdown.Export do
   defp render_cell(%Cell.Elixir{} = cell, ctx) do
     delimiter = MarkdownHelpers.code_block_delimiter(cell.source)
     code = get_elixir_cell_code(cell)
-    outputs = if ctx.include_outputs?, do: render_outputs(cell), else: []
+    outputs = if ctx.include_outputs?, do: render_outputs(cell, ctx), else: []
 
     metadata = cell_metadata(cell)
 
@@ -119,33 +149,37 @@ defmodule Livebook.LiveMarkdown.Export do
 
   defp cell_metadata(_cell), do: %{}
 
-  defp render_outputs(cell) do
+  defp render_outputs(cell, ctx) do
     cell.outputs
     |> Enum.reverse()
-    |> Enum.map(&render_output/1)
+    |> Enum.map(&render_output(&1, ctx))
     |> Enum.reject(&(&1 == :ignored))
     |> Enum.intersperse("\n\n")
   end
 
-  defp render_output(text) when is_binary(text) do
+  defp render_output(text, _ctx) when is_binary(text) do
     text = String.replace_suffix(text, "\n", "")
     delimiter = MarkdownHelpers.code_block_delimiter(text)
     text = strip_ansi(text)
     [delimiter, "output\n", text, "\n", delimiter]
   end
 
-  defp render_output({:text, text}) do
+  defp render_output({:text, text}, _ctx) do
     delimiter = MarkdownHelpers.code_block_delimiter(text)
     text = strip_ansi(text)
     [delimiter, "output\n", text, "\n", delimiter]
   end
 
-  defp render_output({:vega_lite_static, spec}) do
+  defp render_output({:vega_lite_static, spec}, _ctx) do
     ["```", "vega-lite\n", Jason.encode!(spec), "\n", "```"]
   end
 
-  defp render_output({:js_static, %{export: %{info_string: info_string, key: key}}, data})
+  defp render_output(
+         {:js, %{export: %{info_string: info_string, key: key}, ref: ref}},
+         ctx
+       )
        when is_binary(info_string) do
+    data = ctx.js_ref_with_data[ref]
     payload = if key && is_map(data), do: data[key], else: data
 
     case encode_js_data(payload) do
@@ -157,7 +191,7 @@ defmodule Livebook.LiveMarkdown.Export do
     end
   end
 
-  defp render_output(_output), do: :ignored
+  defp render_output(_output, _ctx), do: :ignored
 
   defp encode_js_data(data) when is_binary(data), do: {:ok, data}
   defp encode_js_data(data), do: Jason.encode(data)

lib/livebook/notebook.ex

@@ -520,8 +520,7 @@ defmodule Livebook.Notebook do
       Enum.find_value(section.cells, fn cell ->
         is_struct(cell, Cell.Elixir) &&
           Enum.find_value(cell.outputs, fn
-            {:js_static, %{assets: %{hash: ^hash} = assets_info}, _data} -> assets_info
-            {:js_dynamic, %{assets: %{hash: ^hash} = assets_info}, _pid} -> assets_info
+            {:js, %{assets: %{hash: ^hash} = assets_info}} -> assets_info
             _ -> nil
           end)
       end)

lib/livebook/notebook/cell/elixir.ex

@@ -36,10 +36,8 @@ defmodule Livebook.Notebook.Cell.Elixir do
           | {:vega_lite_static, spec :: map()}
           # Vega-Lite graphic with dynamic data
           | {:vega_lite_dynamic, widget_process :: pid()}
-          # JavaScript powered output with static data
-          | {:js_static, info :: map(), data :: term()}
-          # JavaScript powered output with server process
-          | {:js_dynamic, info :: map(), widget_process :: pid()}
+          # JavaScript powered output
+          | {:js, info :: map()}
           # Interactive data table
           | {:table_dynamic, widget_process :: pid()}
           # Dynamic wrapper for static output

lib/livebook/session.ex

@@ -1129,10 +1129,11 @@ defmodule Livebook.Session do
 
     if file && should_save_notebook?(state) do
       pid = self()
-      content = LiveMarkdown.Export.notebook_to_markdown(state.data.notebook)
+      notebook = state.data.notebook
 
       {:ok, pid} =
         Task.start(fn ->
+          content = LiveMarkdown.Export.notebook_to_markdown(notebook)
           result = FileSystem.File.write(file, content)
           send(pid, {:save_finished, self(), result, file, default?})
         end)

lib/livebook_web/channels/js_output_channel.ex

@@ -0,0 +1,74 @@
+defmodule LivebookWeb.JSOutputChannel do
+  use Phoenix.Channel
+
+  @impl true
+  def join("js_output", %{}, socket) do
+    {:ok, assign(socket, ref_with_pid: %{}, ref_with_count: %{})}
+  end
+
+  @impl true
+  def handle_in("connect", %{"session_token" => session_token, "ref" => ref}, socket) do
+    {:ok, data} = Phoenix.Token.verify(LivebookWeb.Endpoint, "js output", session_token)
+    %{pid: pid} = data
+
+    send(pid, {:connect, self(), %{origin: self(), ref: ref}})
+
+    ref_with_pid = Map.put(socket.assigns.ref_with_pid, ref, pid)
+    ref_with_count = Map.update(socket.assigns.ref_with_count, ref, 1, &(&1 + 1))
+    socket = assign(socket, ref_with_pid: ref_with_pid, ref_with_count: ref_with_count)
+    {:noreply, socket}
+  end
+
+  def handle_in("event", %{"event" => event, "payload" => payload, "ref" => ref}, socket) do
+    pid = socket.assigns.ref_with_pid[ref]
+    send(pid, {:event, event, payload, %{origin: self(), ref: ref}})
+    {:noreply, socket}
+  end
+
+  def handle_in("disconnect", %{"ref" => ref}, socket) do
+    socket =
+      if socket.assigns.ref_with_count[ref] == 1 do
+        {_, ref_with_count} = Map.pop!(socket.assigns.ref_with_count, ref)
+        {_, ref_with_pid} = Map.pop!(socket.assigns.ref_with_pid, ref)
+        assign(socket, ref_with_count: ref_with_count, ref_with_pid: ref_with_pid)
+      else
+        ref_with_count = Map.update!(socket.assigns.ref_with_count, ref, &(&1 - 1))
+        assign(socket, ref_with_count: ref_with_count)
+      end
+
+    {:noreply, socket}
+  end
+
+  @impl true
+  def handle_info({:connect_reply, data, %{ref: ref}}, socket) do
+    with {:error, error} <- try_push(socket, "init:#{ref}", %{"data" => data}) do
+      message = "Failed to serialize initial widget data, " <> error
+      push(socket, "error:#{ref}", %{"message" => message})
+    end
+
+    {:noreply, socket}
+  end
+
+  def handle_info({:event, event, payload, %{ref: ref}}, socket) do
+    with {:error, error} <-
+           try_push(socket, "event:#{ref}", %{"event" => event, "payload" => payload}) do
+      message = "Failed to serialize event payload, " <> error
+      push(socket, "error:#{ref}", %{"message" => message})
+    end
+
+    {:noreply, socket}
+  end
+
+  # In case the payload fails to encode we catch the error
+  defp try_push(socket, event, payload) do
+    try do
+      push(socket, event, payload)
+    catch
+      :error, %Protocol.UndefinedError{protocol: Jason.Encoder, value: value} ->
+        {:error, "value #{inspect(value)} is not JSON-serializable, use another data type"}
+
+      :error, error ->
+        {:error, Exception.message(error)}
+    end
+  end
+end

lib/livebook_web/channels/socket.ex

@@ -0,0 +1,26 @@
+defmodule LivebookWeb.Socket do
+  use Phoenix.Socket
+
+  # App channels
+  channel "js_output", LivebookWeb.JSOutputChannel
+
+  # LiveView channels
+  channel "lvu:*", Phoenix.LiveView.UploadChannel
+  channel "lv:*", Phoenix.LiveView.Channel
+
+  @impl true
+  def connect(params, socket, info) do
+    auth_mode = Livebook.Config.auth_mode()
+
+    if LivebookWeb.AuthPlug.authenticated?(info.session || %{}, info.uri.port, auth_mode) do
+      Phoenix.LiveView.Socket.connect(params, socket, info)
+    else
+      :error
+    end
+  end
+
+  @impl true
+  def id(socket) do
+    Phoenix.LiveView.Socket.id(socket)
+  end
+end

lib/livebook_web/endpoint.ex

@@ -11,11 +11,11 @@ defmodule LivebookWeb.Endpoint do
     same_site: "Lax"
   ]
 
-  socket "/live", Phoenix.LiveView.Socket,
+  socket "/live", LivebookWeb.Socket,
     # Don't check the origin as we don't know how the web app is gonna be accessed.
     # It runs locally, but may be exposed via IP or domain name.
     # The WebSocket connection is already protected from CSWSH by using CSRF token.
-    websocket: [check_origin: false, connect_info: [:user_agent, session: @session_options]]
+    websocket: [check_origin: false, connect_info: [:user_agent, :uri, session: @session_options]]
 
   # We use Escript for distributing Livebook, so we don't have access to the static
   # files at runtime in the prod environment. To overcome this we load contents of

lib/livebook_web/live/output.ex

@@ -94,20 +94,8 @@ defmodule LivebookWeb.Output do
     )
   end
 
-  defp render_output({:js_static, info, data}, %{id: id, session_id: session_id}) do
-    live_component(LivebookWeb.Output.JSStaticComponent,
-      id: id,
-      info: info,
-      data: data,
-      session_id: session_id
-    )
-  end
-
-  defp render_output({:js_dynamic, info, pid}, %{id: id, socket: socket, session_id: session_id}) do
-    live_render(socket, LivebookWeb.Output.JSDynamicLive,
-      id: id,
-      session: %{"id" => id, "info" => info, "pid" => pid, "session_id" => session_id}
-    )
+  defp render_output({:js, info}, %{id: id, session_id: session_id}) do
+    live_component(LivebookWeb.Output.JSComponent, id: id, info: info, session_id: session_id)
   end
 
   defp render_output({:table_dynamic, pid}, %{id: id, socket: socket}) do

lib/livebook_web/live/output/js_component.ex

@@ -0,0 +1,21 @@
+defmodule LivebookWeb.Output.JSComponent do
+  use LivebookWeb, :live_component
+
+  @impl true
+  def render(assigns) do
+    ~H"""
+    <div id={"js-output-#{@id}"}
+      phx-hook="JSOutput"
+      phx-update="ignore"
+      data-ref={@info.ref}
+      data-assets-base-url={Routes.session_url(@socket, :show_asset, @session_id, @info.assets.hash, [])}
+      data-js-path={@info.assets.js_path}
+      data-session-token={session_token(@info.pid)}>
+    </div>
+    """
+  end
+
+  defp session_token(pid) do
+    Phoenix.Token.sign(LivebookWeb.Endpoint, "js output", %{pid: pid})
+  end
+end

lib/livebook_web/live/output/js_dynamic_live.ex

@@ -1,56 +0,0 @@
-defmodule LivebookWeb.Output.JSDynamicLive do
-  use LivebookWeb, :live_view
-
-  @impl true
-  def mount(
-        _params,
-        %{"pid" => pid, "id" => id, "info" => info, "session_id" => session_id},
-        socket
-      ) do
-    if connected?(socket) do
-      send(pid, {:connect, self(), %{origin: self()}})
-    end
-
-    assets_base_url = Routes.session_url(socket, :show_asset, session_id, info.assets.hash, [])
-
-    {:ok,
-     assign(socket,
-       widget_pid: pid,
-       id: id,
-       assets_base_url: assets_base_url,
-       js_path: info.assets.js_path
-     )}
-  end
-
-  @impl true
-  def render(assigns) do
-    ~H"""
-    <div id={"js-output-#{@id}"}
-      phx-hook="JSOutput"
-      phx-update="ignore"
-      data-id={@id}
-      data-assets-base-url={@assets_base_url}
-      data-js-path={@js_path}>
-    </div>
-    """
-  end
-
-  @impl true
-  def handle_event("event", %{"event" => event, "payload" => payload}, socket) do
-    send(socket.assigns.widget_pid, {:event, event, payload, %{origin: self()}})
-    {:noreply, socket}
-  end
-
-  @impl true
-  def handle_info({:connect_reply, data}, socket) do
-    {:noreply, push_event(socket, "js_output:#{socket.assigns.id}:init", %{"data" => data})}
-  end
-
-  def handle_info({:event, event, payload}, socket) do
-    {:noreply,
-     push_event(socket, "js_output:#{socket.assigns.id}:event", %{
-       "event" => event,
-       "payload" => payload
-     })}
-  end
-end

lib/livebook_web/live/output/js_static_component.ex

@@ -1,45 +0,0 @@
-defmodule LivebookWeb.Output.JSStaticComponent do
-  use LivebookWeb, :live_component
-
-  @impl true
-  def mount(socket) do
-    {:ok, assign(socket, initialized: false)}
-  end
-
-  @impl true
-  def update(assigns, socket) do
-    assets_base_url =
-      Routes.session_url(socket, :show_asset, assigns.session_id, assigns.info.assets.hash, [])
-
-    socket =
-      assign(socket,
-        id: assigns.id,
-        assets_base_url: assets_base_url,
-        js_path: assigns.info.assets.js_path
-      )
-
-    socket =
-      if connected?(socket) and not socket.assigns.initialized do
-        socket
-        |> assign(initialized: true)
-        |> push_event("js_output:#{socket.assigns.id}:init", %{"data" => assigns.data})
-      else
-        socket
-      end
-
-    {:ok, socket}
-  end
-
-  @impl true
-  def render(assigns) do
-    ~H"""
-    <div id={"js-output-#{@id}"}
-      phx-hook="JSOutput"
-      phx-update="ignore"
-      data-id={@id}
-      data-assets-base-url={@assets_base_url}
-      data-js-path={@js_path}>
-    </div>
-    """
-  end
-end

lib/livebook_web/plugs/auth_plug.ex

@@ -28,21 +28,29 @@ defmodule LivebookWeb.AuthPlug do
   Stores in the session the secret for the given mode.
   """
   def store(conn, mode, value) do
-    put_session(conn, key(conn, mode), hash(value))
+    put_session(conn, key(conn.port, mode), hash(value))
   end
 
   @doc """
   Checks if given connection is already authenticated.
   """
   @spec authenticated?(Plug.Conn.t(), Livebook.Config.auth_mode()) :: boolean()
-  def authenticated?(conn, mode)
+  def authenticated?(conn, mode) do
+    authenticated?(get_session(conn), conn.port, mode)
+  end
 
-  def authenticated?(conn, mode) when mode in [:token, :password] do
-    secret = get_session(conn, key(conn, mode))
+  @doc """
+  Checks if the given session is authenticated.
+  """
+  @spec authenticated?(map(), non_neg_integer(), Livebook.Config.auth_mode()) :: boolean()
+  def authenticated?(session, port, mode)
+
+  def authenticated?(session, port, mode) when mode in [:token, :password] do
+    secret = session[key(port, mode)]
     is_binary(secret) and Plug.Crypto.secure_compare(secret, expected(mode))
   end
 
-  def authenticated?(_conn, _mode) do
+  def authenticated?(_session, _port, _mode) do
     true
   end
 
@@ -66,7 +74,7 @@ defmodule LivebookWeb.AuthPlug do
     end
   end
 
-  defp key(conn, mode), do: "#{conn.port}:#{mode}"
+  defp key(port, mode), do: "#{port}:#{mode}"
   defp expected(mode), do: hash(Application.fetch_env!(:livebook, mode))
   defp hash(value), do: :crypto.hash(:sha256, value)
 end

test/livebook/live_markdown/export_test.exs

@@ -678,7 +678,7 @@ defmodule Livebook.LiveMarkdown.ExportTest do
       assert expected_document == document
     end
 
-    test "does not include js_static output with no export info" do
+    test "does not include js output with no export info" do
       notebook = %{
         Notebook.new()
         | name: "My Notebook",
@@ -691,11 +691,13 @@ defmodule Livebook.LiveMarkdown.ExportTest do
                     Notebook.Cell.new(:elixir)
                     | source: ":ok",
                       outputs: [
-                        {:js_static,
+                        {:js,
                          %{
+                           ref: "1",
+                           pid: spawn_widget_with_data("1", "data"),
                            assets: %{archive_path: "", hash: "abcd", js_path: "main.js"},
                            export: nil
-                         }, "data"}
+                         }}
                       ]
                   }
                 ]
@@ -718,7 +720,7 @@ defmodule Livebook.LiveMarkdown.ExportTest do
       assert expected_document == document
     end
 
-    test "includes js_static output if export info is set" do
+    test "includes js output if export info is set" do
       notebook = %{
         Notebook.new()
         | name: "My Notebook",
@@ -731,11 +733,13 @@ defmodule Livebook.LiveMarkdown.ExportTest do
                     Notebook.Cell.new(:elixir)
                     | source: ":ok",
                       outputs: [
-                        {:js_static,
+                        {:js,
                          %{
+                           ref: "1",
+                           pid: spawn_widget_with_data("1", "graph TD;\nA-->B;"),
                            assets: %{archive_path: "", hash: "abcd", js_path: "main.js"},
                            export: %{info_string: "mermaid", key: nil}
-                         }, "graph TD;\nA-->B;"}
+                         }}
                       ]
                   }
                 ]
@@ -763,7 +767,7 @@ defmodule Livebook.LiveMarkdown.ExportTest do
       assert expected_document == document
     end
 
-    test "serializes js_static data to JSON if not binary" do
+    test "serializes js output data to JSON if not binary" do
       notebook = %{
         Notebook.new()
         | name: "My Notebook",
@@ -776,11 +780,13 @@ defmodule Livebook.LiveMarkdown.ExportTest do
                     Notebook.Cell.new(:elixir)
                     | source: ":ok",
                       outputs: [
-                        {:js_static,
+                        {:js,
                          %{
+                           ref: "1",
+                           pid: spawn_widget_with_data("1", %{height: 50, width: 50}),
                            assets: %{archive_path: "", hash: "abcd", js_path: "main.js"},
                            export: %{info_string: "box", key: nil}
-                         }, %{height: 50, width: 50}}
+                         }}
                       ]
                   }
                 ]
@@ -807,7 +813,7 @@ defmodule Livebook.LiveMarkdown.ExportTest do
       assert expected_document == document
     end
 
-    test "exports partial js_static data when export_key is set" do
+    test "exports partial js output data when export_key is set" do
       notebook = %{
         Notebook.new()
         | name: "My Notebook",
@@ -820,11 +826,17 @@ defmodule Livebook.LiveMarkdown.ExportTest do
                     Notebook.Cell.new(:elixir)
                     | source: ":ok",
                       outputs: [
-                        {:js_static,
+                        {:js,
                          %{
+                           ref: "1",
+                           pid:
+                             spawn_widget_with_data("1", %{
+                               spec: %{"height" => 50, "width" => 50},
+                               datasets: []
+                             }),
                            assets: %{archive_path: "", hash: "abcd", js_path: "main.js"},
                            export: %{info_string: "vega-lite", key: :spec}
-                         }, %{spec: %{"height" => 50, "width" => 50}, datasets: []}}
+                         }}
                       ]
                   }
                 ]
@@ -951,4 +963,12 @@ defmodule Livebook.LiveMarkdown.ExportTest do
 
     assert expected_document == document
   end
+
+  defp spawn_widget_with_data(ref, data) do
+    spawn(fn ->
+      receive do
+        {:connect, pid, %{ref: ^ref}} -> send(pid, {:connect_reply, data, %{ref: ref}})
+      end
+    end)
+  end
 end

test/livebook/notebook_test.exs

@@ -259,7 +259,7 @@ defmodule Livebook.NotebookTest do
     test "returns asset info matching the given type if found" do
       assets_info = %{archive: "/path/to/archive.tar.gz", hash: "abcd", js_path: "main.js"}
       js_info = %{assets: assets_info}
-      output = {:js_static, js_info, %{}}
+      output = {:js, js_info}
 
       notebook = %{
         Notebook.new()

test/livebook_web/channels/js_output_channel_test.exs

@@ -0,0 +1,42 @@
+defmodule LivebookWeb.JSOutputChannelTest do
+  use LivebookWeb.ChannelCase
+
+  setup do
+    {:ok, _, socket} =
+      LivebookWeb.Socket
+      |> socket()
+      |> subscribe_and_join(LivebookWeb.JSOutputChannel, "js_output")
+
+    %{socket: socket}
+  end
+
+  test "loads initial data from the widget server and pushes to the client", %{socket: socket} do
+    push(socket, "connect", %{"session_token" => session_token(), "ref" => "1"})
+
+    assert_receive {:connect, from, %{}}
+    send(from, {:connect_reply, [1, 2, 3], %{ref: "1"}})
+
+    assert_push "init:1", %{"data" => [1, 2, 3]}
+  end
+
+  test "sends events received from widget server to the client", %{socket: socket} do
+    send(socket.channel_pid, {:event, "ping", [1, 2, 3], %{ref: "1"}})
+
+    assert_push "event:1", %{"event" => "ping", "payload" => [1, 2, 3]}
+  end
+
+  test "sends client events to the corresponding widget server", %{socket: socket} do
+    push(socket, "connect", %{"session_token" => session_token(), "ref" => "1"})
+
+    assert_receive {:connect, from, %{}}
+    send(from, {:connect_reply, [1, 2, 3], %{ref: "1"}})
+
+    push(socket, "event", %{"event" => "ping", "payload" => [1, 2, 3], "ref" => "1"})
+
+    assert_receive {:event, "ping", [1, 2, 3], %{origin: _origin}}
+  end
+
+  defp session_token() do
+    Phoenix.Token.sign(LivebookWeb.Endpoint, "js output", %{pid: self()})
+  end
+end

test/livebook_web/controllers/session_controller_test.exs

@@ -204,7 +204,7 @@ defmodule LivebookWeb.SessionControllerTest do
     archive_path = Path.expand("../../support/assets.tar.gz", __DIR__)
     hash = "test-" <> Livebook.Utils.random_id()
     assets_info = %{archive_path: archive_path, hash: hash, js_path: "main.js"}
-    output = {:js_static, %{assets: assets_info}, %{}}
+    output = {:js, %{assets: assets_info}}
 
     notebook = %{
       Notebook.new()

test/livebook_web/live/session_live_test.exs

@@ -290,46 +290,6 @@ defmodule LivebookWeb.SessionLiveTest do
 
       assert render(view) =~ "Dynamic output in frame"
     end
-
-    test "static js output sends the embedded data to the client", %{conn: conn, session: session} do
-      js_info = %{assets: %{archive_path: "", hash: "abcd", js_path: "main.js"}}
-      js_static_output = {:js_static, js_info, [1, 2, 3]}
-
-      section_id = insert_section(session.pid)
-      cell_id = insert_text_cell(session.pid, section_id, :elixir)
-      # Evaluate the cell
-      Session.queue_cell_evaluation(session.pid, cell_id)
-      # Send an additional output
-      send(session.pid, {:evaluation_output, cell_id, js_static_output})
-
-      {:ok, view, _} = live(conn, "/sessions/#{session.id}")
-
-      assert_push_event(view, "js_output:" <> _, %{"data" => [1, 2, 3]})
-    end
-
-    test "dynamic js output loads initial data from the widget server",
-         %{conn: conn, session: session} do
-      widget_pid =
-        spawn(fn ->
-          receive do
-            {:connect, pid, %{}} -> send(pid, {:connect_reply, [1, 2, 3]})
-          end
-        end)
-
-      js_info = %{assets: %{archive_path: "", hash: "abcd", js_path: "main.js"}}
-      js_dynamic_output = {:js_dynamic, js_info, widget_pid}
-
-      section_id = insert_section(session.pid)
-      cell_id = insert_text_cell(session.pid, section_id, :elixir)
-      # Evaluate the cell
-      Session.queue_cell_evaluation(session.pid, cell_id)
-      # Send an additional output
-      send(session.pid, {:evaluation_output, cell_id, js_dynamic_output})
-
-      {:ok, view, _} = live(conn, "/sessions/#{session.id}")
-
-      assert_push_event(view, "js_output:" <> _, %{"data" => [1, 2, 3]})
-    end
   end
 
   describe "runtime settings" do

test/support/channel_case.ex

@@ -0,0 +1,13 @@
+defmodule LivebookWeb.ChannelCase do
+  use ExUnit.CaseTemplate
+
+  using do
+    quote do
+      # Import conveniences for testing with channels
+      import Phoenix.ChannelTest
+
+      # The default endpoint for testing
+      @endpoint LivebookWeb.Endpoint
+    end
+  end
+end