Add send_event/4 for Kino.JS.Live

lib/kino/bridge.ex

@@ -110,6 +110,16 @@ defmodule Kino.Bridge do
     :ok
   end
 
+  @doc """
+  Starts monitoring the given Livebook process.
+
+  Provides the same semantics as `Process.monitor/1`.
+  """
+  @spec monitor(pid()) :: reference()
+  def monitor(pid) do
+    Process.monitor(pid)
+  end
+
   defp io_request(request) do
     gl = Process.group_leader()
     ref = Process.monitor(gl)

lib/kino/js/live.ex

@@ -253,7 +253,7 @@ defmodule Kino.JS.Live do
     quote location: :keep do
       @behaviour Kino.JS.Live
 
-      import Kino.JS.Live.Context, only: [assign: 2, update: 3, broadcast_event: 3]
+      import Kino.JS.Live.Context, only: [assign: 2, update: 3, broadcast_event: 3, send_event: 4]
 
       @before_compile Kino.JS.Live
     end

lib/kino/js/live/context.ex

@@ -53,7 +53,7 @@ defmodule Kino.JS.Live.Context do
   end
 
   @doc """
-  Sends an event to the client.
+  Sends an event to all clients.
 
   The event is dispatched to the registered JavaScript callback
   on all connected clients.
@@ -66,4 +66,19 @@ defmodule Kino.JS.Live.Context do
   def broadcast_event(%__MODULE__{} = ctx, event, payload \\ nil) when is_binary(event) do
     Kino.JS.Live.Server.broadcast_event(ctx, event, payload)
   end
+
+  @doc """
+  Sends an event to a specific client.
+
+  The event is dispatched to the registered JavaScript callback
+  on the specific connected client.
+
+  ## Examples
+
+      send_event(ctx, origin, "new_point", %{x: 10, y: 10})
+  """
+  @spec send_event(t(), term(), String.t(), term()) :: :ok
+  def send_event(%__MODULE__{} = ctx, origin, event, payload \\ nil) when is_binary(event) do
+    Kino.JS.Live.Server.send_event(ctx, origin, event, payload)
+  end
 end

lib/kino/js/live/server.ex

@@ -23,6 +23,19 @@ defmodule Kino.JS.Live.Server do
     :ok
   end
 
+  def send_event(ctx, origin, event, payload) do
+    ref = ctx.__private__.ref
+
+    pid =
+      case ctx.__private__.clients[origin] do
+        {pid, _} -> pid
+        _ -> raise "could not find a connected client with origin #{inspect(origin)}"
+      end
+
+    Kino.Bridge.send(pid, {:event, event, payload, %{ref: ref}})
+    :ok
+  end
+
   @impl true
   def init({module, init_arg, ref}) do
     {:ok, ctx, _opts} = call_init(module, init_arg, ref)
@@ -59,6 +72,7 @@ defmodule Kino.JS.Live.Server do
   def call_init(module, init_arg, ref) do
     ctx = Context.new()
     ctx = put_in(ctx.__private__[:ref], ref)
+    ctx = put_in(ctx.__private__[:clients], %{})
 
     if has_function?(module, :init, 2) do
       case module.init(init_arg, ctx) do
@@ -73,6 +87,8 @@ defmodule Kino.JS.Live.Server do
   def call_handle_info(msg, module, ctx)
 
   def call_handle_info({:connect, pid, %{origin: origin}}, module, ctx) do
+    ctx = add_client(ctx, pid, origin)
+
     ctx = %{ctx | origin: origin}
     {:ok, data, ctx} = module.handle_connect(ctx)
     ctx = %{ctx | origin: nil}
@@ -96,7 +112,17 @@ defmodule Kino.JS.Live.Server do
     {:ok, ctx}
   end
 
+  def call_handle_info({:DOWN, ref, :process, _pid, _reason} = msg, module, ctx) do
+    with :error <- remove_client(ctx, ref) do
+      call_handle_info_fallback(msg, module, ctx)
+    end
+  end
+
   def call_handle_info(msg, module, ctx) do
+    call_handle_info_fallback(msg, module, ctx)
+  end
+
+  def call_handle_info_fallback(msg, module, ctx) do
     if has_function?(module, :handle_info, 2) do
       {:noreply, ctx} = module.handle_info(msg, ctx)
       {:ok, ctx}
@@ -116,4 +142,28 @@ defmodule Kino.JS.Live.Server do
 
     :ok
   end
+
+  defp add_client(ctx, pid, origin) do
+    if Map.has_key?(ctx.__private__.clients, origin) do
+      ctx
+    else
+      monitor_ref = Kino.Bridge.monitor(pid)
+      put_in(ctx.__private__.clients[origin], {pid, monitor_ref})
+    end
+  end
+
+  defp remove_client(ctx, ref) do
+    origin =
+      Enum.find_value(ctx.__private__.clients, fn
+        {origin, {_pid, ^ref}} -> origin
+        _ -> nil
+      end)
+
+    if origin do
+      {_, ctx} = pop_in(ctx.__private__.clients[origin])
+      {:ok, ctx}
+    else
+      :error
+    end
+  end
 end

lib/kino/test.ex

@@ -55,6 +55,23 @@ defmodule Kino.Test do
     end
   end
 
+  @doc """
+  Asserts a `Kino.JS.Live` kino will send an event within `timeout`
+  to the caller.
+
+  ## Examples
+
+      assert_send_event(kino, "pong", %{})
+
+  """
+  defmacro assert_send_event(kino, event, payload, timeout \\ 100) do
+    quote do
+      %{ref: ref} = unquote(kino)
+
+      assert_receive {:event, unquote(event), unquote(payload), %{ref: ^ref}}, unquote(timeout)
+    end
+  end
+
   @doc """
   Sends a client event to a `Kino.JS.Live` kino.
 

test/kino/js/live_test.exs

@@ -38,6 +38,14 @@ defmodule Kino.JS.LiveTest do
       count = LiveCounter.read(kino)
       assert count == 2
     end
+
+    test "handle_event/3 with send event" do
+      kino = LiveCounter.new(0)
+      # Simulate a client event
+      _ = connect(kino)
+      push_event(kino, "ping", %{})
+      assert_send_event(kino, "pong", %{})
+    end
   end
 
   test "server ping" do

test/support/test_modules/live_counter.ex

@@ -29,6 +29,11 @@ defmodule Kino.TestModules.LiveCounter do
     {:noreply, bump_count(ctx, by)}
   end
 
+  def handle_event("ping", %{}, ctx) do
+    send_event(ctx, ctx.origin, "pong", %{})
+    {:noreply, ctx}
+  end
+
   @impl true
   def handle_cast({:bump, by}, ctx) do
     {:noreply, bump_count(ctx, by)}