Repository: livebook-dev/kino (elixir)

Issue:
Add support for messaging a single JS view client
Currently we have `broadcast_event` that pushes to everyone, but in certain cases we may want to talk to a specific client. See https://github.com/livebook-dev/livebook/discussions/1297#discussioncomment-3225199.

Pull request description:
Add send_event/4 for Kino.JS.Live
Closes #171.

[send_event.webm](https://user-images.githubusercontent.com/17034772/182946178-99d2b0ba-01b6-47db-ae92-eab9aa4fb2c7.webm)

Note that if the same JS view is rendered multiple times, all instances are still the same client, because client.

Patch:
diff --git a/lib/kino/bridge.ex b/lib/kino/bridge.ex
index 98e06db..f4e597d 100644
--- a/lib/kino/bridge.ex
+++ b/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)
diff --git a/lib/kino/js/live.ex b/lib/kino/js/live.ex
index b611202..720a07d 100644
--- a/lib/kino/js/live.ex
+++ b/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
diff --git a/lib/kino/js/live/context.ex b/lib/kino/js/live/context.ex
index c00d2ad..9446379 100644
--- a/lib/kino/js/live/context.ex
+++ b/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
diff --git a/lib/kino/js/live/server.ex b/lib/kino/js/live/server.ex
index 59e4d6c..b63a95c 100644
--- a/lib/kino/js/live/server.ex
+++ b/lib/kino/js/live/server.ex
@@ -23,6 +23,16 @@ defmodule Kino.JS.Live.Server do
     :ok
   end
 
+  def send_event(ctx, origin, event, payload) do
+    ref = ctx.__private__.ref
+
+    if pid = ctx.__private__.origins_with_pid[origin] do
+      Kino.Bridge.send(pid, {:event, event, payload, %{ref: ref}})
+    end
+
+    :ok
+  end
+
   @impl true
   def init({module, init_arg, ref}) do
     {:ok, ctx, _opts} = call_init(module, init_arg, ref)
@@ -59,6 +69,8 @@ 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__[:origins_with_pid], %{})
+    ctx = put_in(ctx.__private__[:monitors_with_origin], %{})
 
     if has_function?(module, :init, 2) do
       case module.init(init_arg, ctx) do
@@ -73,6 +85,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 +110,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 +140,25 @@ defmodule Kino.JS.Live.Server do
 
     :ok
   end
+
+  defp add_client(ctx, pid, origin) do
+    if Map.has_key?(ctx.__private__.origins_with_pid, origin) do
+      ctx
+    else
+      monitor_ref = Kino.Bridge.monitor(pid)
+      ctx = put_in(ctx.__private__.origins_with_pid[origin], pid)
+      put_in(ctx.__private__.monitors_with_origin[monitor_ref], origin)
+    end
+  end
+
+  defp remove_client(ctx, monitor_ref) do
+    {origin, ctx} = pop_in(ctx.__private__.monitors_with_origin[monitor_ref])
+
+    if origin do
+      {_, ctx} = pop_in(ctx.__private__.origins_with_pid[origin])
+      {:ok, ctx}
+    else
+      :error
+    end
+  end
 end
diff --git a/lib/kino/terminator.ex b/lib/kino/terminator.ex
index 6872388..3f85be0 100644
--- a/lib/kino/terminator.ex
+++ b/lib/kino/terminator.ex
@@ -25,7 +25,7 @@ defmodule Kino.Terminator do
 
   @impl true
   def handle_info({:terminate, pid}, state) do
-    :ok = DynamicSupervisor.terminate_child(Kino.DynamicSupervisor, pid)
+    DynamicSupervisor.terminate_child(Kino.DynamicSupervisor, pid)
     {:noreply, state}
   end
 end


Test patch:
diff --git a/lib/kino/test.ex b/lib/kino/test.ex
index fedc514..c75f978 100644
--- a/lib/kino/test.ex
+++ b/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.
 
diff --git a/test/kino/js/live_test.exs b/test/kino/js/live_test.exs
index dfe1c02..ff5bdc5 100644
--- a/test/kino/js/live_test.exs
+++ b/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
diff --git a/test/support/test_modules/live_counter.ex b/test/support/test_modules/live_counter.ex
index 394ee9f..2d40b18 100644
--- a/test/support/test_modules/live_counter.ex
+++ b/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)}