Smart cell editor - intellisense node (#2232)

Co-authored-by: Jonatan Kłosko <jonatanklosko@gmail.com>

assets/js/hooks/js_view.js

@@ -381,12 +381,18 @@ const JSView = {
       } else if (message.type === "syncReply") {
         this.pongCallbackQueue.push(this.syncCallbackQueue.shift());
         this.channel.push("ping", { ref: this.props.ref });
-      } else if (message.type == "selectSecret") {
+      } else if (message.type === "selectSecret") {
         this.pushEvent("select_secret", {
           js_view_ref: this.props.ref,
           preselect_name: message.preselectName,
           options: message.options,
         });
+      } else if (message.type === "setSmartCellEditorIntellisenseNode") {
+        this.pushEvent("set_smart_cell_editor_intellisense_node", {
+          js_view_ref: this.props.ref,
+          node: message.node,
+          cookie: message.cookie,
+        });
       }
     }
   },

lib/livebook/notebook/cell/smart.ex

@@ -10,7 +10,17 @@ defmodule Livebook.Notebook.Cell.Smart do
   # The available smart cells come from the runtime, therefore they
   # are one Livebook's extension points.
 
-  defstruct [:id, :source, :chunks, :outputs, :kind, :attrs, :js_view, :editor]
+  defstruct [
+    :id,
+    :source,
+    :chunks,
+    :outputs,
+    :kind,
+    :attrs,
+    :js_view,
+    :editor,
+    :editor_intellisense_node
+  ]
 
   alias Livebook.Utils
   alias Livebook.Notebook.Cell
@@ -23,7 +33,8 @@ defmodule Livebook.Notebook.Cell.Smart do
           kind: String.t() | nil,
           attrs: attrs() | :__pruned__,
           js_view: Livebook.Runtime.js_view() | nil,
-          editor: Livebook.Runtime.editor() | nil
+          editor: Livebook.Runtime.editor() | nil,
+          editor_intellisense_node: {String.t(), String.t()} | nil
         }
 
   @type attrs :: map()
@@ -41,7 +52,8 @@ defmodule Livebook.Notebook.Cell.Smart do
       kind: nil,
       attrs: %{},
       js_view: nil,
-      editor: nil
+      editor: nil,
+      editor_intellisense_node: nil
     }
   end
 end

lib/livebook/runtime.ex

@@ -871,8 +871,14 @@ defprotocol Livebook.Runtime do
   The given `parent_locators` identifies a sequence of evaluations
   that may be used as the context when resolving the request (if relevant).
   """
-  @spec handle_intellisense(t(), pid(), intellisense_request(), parent_locators()) :: reference()
-  def handle_intellisense(runtime, send_to, request, parent_locators)
+  @spec handle_intellisense(
+          t(),
+          pid(),
+          intellisense_request(),
+          parent_locators(),
+          {String.t(), String.t()} | nil
+        ) :: reference()
+  def handle_intellisense(runtime, send_to, request, parent_locators, node)
 
   @doc """
   Reads file at the given absolute path within the runtime file system.

lib/livebook/runtime/attached.ex

@@ -131,8 +131,8 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.Attached do
     RuntimeServer.drop_container(runtime.server_pid, container_ref)
   end
 
-  def handle_intellisense(runtime, send_to, request, parent_locators) do
-    RuntimeServer.handle_intellisense(runtime.server_pid, send_to, request, parent_locators)
+  def handle_intellisense(runtime, send_to, request, parent_locators, node) do
+    RuntimeServer.handle_intellisense(runtime.server_pid, send_to, request, parent_locators, node)
   end
 
   def read_file(runtime, path) do

lib/livebook/runtime/elixir_standalone.ex

@@ -129,8 +129,8 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.ElixirStandalone do
     RuntimeServer.drop_container(runtime.server_pid, container_ref)
   end
 
-  def handle_intellisense(runtime, send_to, request, parent_locators) do
-    RuntimeServer.handle_intellisense(runtime.server_pid, send_to, request, parent_locators)
+  def handle_intellisense(runtime, send_to, request, parent_locators, node) do
+    RuntimeServer.handle_intellisense(runtime.server_pid, send_to, request, parent_locators, node)
   end
 
   def read_file(runtime, path) do

lib/livebook/runtime/embedded.ex

@@ -93,8 +93,8 @@ defimpl Livebook.Runtime, for: Livebook.Runtime.Embedded do
     RuntimeServer.drop_container(runtime.server_pid, container_ref)
   end
 
-  def handle_intellisense(runtime, send_to, request, parent_locators) do
-    RuntimeServer.handle_intellisense(runtime.server_pid, send_to, request, parent_locators)
+  def handle_intellisense(runtime, send_to, request, parent_locators, node) do
+    RuntimeServer.handle_intellisense(runtime.server_pid, send_to, request, parent_locators, node)
   end
 
   def read_file(runtime, path) do

lib/livebook/runtime/erl_dist/runtime_server.ex

@@ -129,11 +129,12 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
           pid(),
           pid(),
           Runtime.intellisense_request(),
-          Runtime.Runtime.parent_locators()
+          Runtime.Runtime.parent_locators(),
+          {String.t(), String.t()} | nil
         ) :: reference()
-  def handle_intellisense(pid, send_to, request, parent_locators) do
+  def handle_intellisense(pid, send_to, request, parent_locators, node) do
     ref = make_ref()
-    GenServer.cast(pid, {:handle_intellisense, send_to, ref, request, parent_locators})
+    GenServer.cast(pid, {:handle_intellisense, send_to, ref, request, parent_locators, node})
     ref
   end
 
@@ -497,7 +498,10 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     {:noreply, state}
   end
 
-  def handle_cast({:handle_intellisense, send_to, ref, request, parent_locators}, state) do
+  def handle_cast(
+        {:handle_intellisense, send_to, ref, request, parent_locators, node},
+        state
+      ) do
     {container_ref, parent_evaluation_refs} =
       case parent_locators do
         [] ->
@@ -525,7 +529,8 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
       end
 
     Task.Supervisor.start_child(state.task_supervisor, fn ->
-      response = Livebook.Intellisense.handle_request(request, intellisense_context, node())
+      node = intellisense_node(node)
+      response = Livebook.Intellisense.handle_request(request, intellisense_context, node)
       send(send_to, {:runtime_intellisense_response, ref, request, response})
     end)
 
@@ -907,4 +912,12 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     name = elem(dependency.dep, 0)
     Application.spec(name) != nil
   end
+
+  defp intellisense_node({node, cookie}) do
+    {node, cookie} = {String.to_atom(node), String.to_atom(cookie)}
+    Node.set_cookie(node, cookie)
+    if Node.connect(node), do: node, else: node()
+  end
+
+  defp intellisense_node(_node), do: node()
 end

lib/livebook_web/live/session_live.ex

@@ -1475,7 +1475,9 @@ defmodule LivebookWeb.SessionLive do
     with {:ok, cell, _section} <- Notebook.fetch_cell_and_section(data.notebook, cell_id) do
       if Runtime.connected?(data.runtime) do
         parent_locators = Session.parent_locators_for_cell(data, cell)
-        ref = Runtime.handle_intellisense(data.runtime, self(), request, parent_locators)
+        node = intellisense_node(cell)
+
+        ref = Runtime.handle_intellisense(data.runtime, self(), request, parent_locators, node)
 
         {:reply, %{"ref" => inspect(ref)}, socket}
       else
@@ -1701,6 +1703,23 @@ defmodule LivebookWeb.SessionLive do
      push_patch(socket, to: ~p"/sessions/#{socket.assigns.session.id}/settings/custom-view")}
   end
 
+  def handle_event(
+        "set_smart_cell_editor_intellisense_node",
+        %{"js_view_ref" => cell_id, "node" => node, "cookie" => cookie},
+        socket
+      ) do
+    node =
+      if is_binary(node) and node =~ "@" and is_binary(cookie) and cookie != "" do
+        {node, cookie}
+      end
+
+    Session.set_cell_attributes(socket.assigns.session.pid, cell_id, %{
+      editor_intellisense_node: node
+    })
+
+    {:noreply, socket}
+  end
+
   @impl true
   def handle_info({:operation, operation}, socket) do
     {:noreply, handle_operation(socket, operation)}
@@ -2909,4 +2928,7 @@ defmodule LivebookWeb.SessionLive do
   defp app_status_color(%{execution: :executed}), do: "bg-green-bright-400"
   defp app_status_color(%{execution: :error}), do: "bg-red-400"
   defp app_status_color(%{execution: :interrupted}), do: "bg-gray-400"
+
+  defp intellisense_node(%Cell.Smart{editor_intellisense_node: node_cookie}), do: node_cookie
+  defp intellisense_node(_), do: nil
 end

test/livebook/runtime/erl_dist/runtime_server_test.exs

@@ -143,10 +143,10 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServerTest do
     end
   end
 
-  describe "handle_intellisense/5 given completion request" do
+  describe "handle_intellisense/6 given completion request" do
     test "provides basic completion when no evaluation reference is given", %{pid: pid} do
       request = {:completion, "System.ver"}
-      ref = RuntimeServer.handle_intellisense(pid, self(), request, [])
+      ref = RuntimeServer.handle_intellisense(pid, self(), request, [], nil)
 
       assert_receive {:runtime_intellisense_response, ^ref, ^request,
                       %{items: [%{label: "version/0"}]}}
@@ -162,33 +162,37 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServerTest do
       assert_receive {:runtime_evaluation_response, :e1, _, %{evaluation_time_ms: _time_ms}}
 
       request = {:completion, "num"}
-      ref = RuntimeServer.handle_intellisense(pid, self(), request, [{:c1, :e1}])
+
+      ref =
+        RuntimeServer.handle_intellisense(pid, self(), request, [{:c1, :e1}], nil)
 
       assert_receive {:runtime_intellisense_response, ^ref, ^request,
                       %{items: [%{label: "number"}]}}
 
       request = {:completion, "ANSI.brigh"}
-      ref = RuntimeServer.handle_intellisense(pid, self(), request, [{:c1, :e1}])
+
+      ref =
+        RuntimeServer.handle_intellisense(pid, self(), request, [{:c1, :e1}], nil)
 
       assert_receive {:runtime_intellisense_response, ^ref, ^request,
                       %{items: [%{label: "bright/0"}]}}
     end
   end
 
-  describe "handle_intellisense/5 given details request" do
+  describe "handle_intellisense/6 given details request" do
     test "responds with identifier details", %{pid: pid} do
       request = {:details, "System.version", 10}
-      ref = RuntimeServer.handle_intellisense(pid, self(), request, [])
+      ref = RuntimeServer.handle_intellisense(pid, self(), request, [], nil)
 
       assert_receive {:runtime_intellisense_response, ^ref, ^request,
                       %{range: %{from: 1, to: 15}, contents: [_]}}
     end
   end
 
-  describe "handle_intellisense/5 given format request" do
+  describe "handle_intellisense/6 given format request" do
     test "responds with a formatted code", %{pid: pid} do
       request = {:format, "System.version"}
-      ref = RuntimeServer.handle_intellisense(pid, self(), request, [])
+      ref = RuntimeServer.handle_intellisense(pid, self(), request, [], nil)
 
       assert_receive {:runtime_intellisense_response, ^ref, ^request, %{code: "System.version()"}}
     end

test/support/noop_runtime.ex

@@ -22,7 +22,7 @@ defmodule Livebook.Runtime.NoopRuntime do
     def evaluate_code(_, _, _, _, _, _ \\ []), do: :ok
     def forget_evaluation(_, _), do: :ok
     def drop_container(_, _), do: :ok
-    def handle_intellisense(_, _, _, _), do: make_ref()
+    def handle_intellisense(_, _, _, _, _), do: make_ref()
 
     def read_file(_, path) do
       case File.read(path) do