Handle intellisense during evaluation (#941)

* Handle intellisense during evaluation

* Apply review comments

* Add TODOs

lib/livebook/evaluator.ex

@@ -48,6 +48,11 @@ defmodule Livebook.Evaluator do
   @type evaluation_response ::
           {:ok, any()} | {:error, Exception.kind(), any(), Exception.stacktrace()}
 
+  # We store evaluation envs in process dictionary, so that we can
+  # build intellisense context without asking the evaluator
+  @env_key :evaluation_env
+  @initial_env_key :initial_env
+
   ## API
 
   @doc """
@@ -154,21 +159,42 @@ defmodule Livebook.Evaluator do
   end
 
   @doc """
-  Asynchronously handles the given intellisense request.
+  Returns an empty intellisense context.
-
-  If `evaluation_ref` is given, its binding and environment are also
-  used as context for the intellisense. Response is sent to the `send_to`
-  process as `{:intellisense_response, ref, response}`.
   """
-  @spec handle_intellisense(
+  @spec intellisense_context() :: Livebook.Intellisense.intellisense_context()
-          t(),
+  def intellisense_context() do
-          pid(),
+    # TODO: Use Code.env_for_eval and eval_quoted_with_env on Elixir v1.14+
-          term(),
+    env = :elixir.env_for_eval([])
-          Livebook.Runtime.intellisense_request(),
+    map_binding = fn fun -> fun.([]) end
-          ref() | nil
+    %{env: env, map_binding: map_binding}
-        ) :: :ok
+  end
-  def handle_intellisense(evaluator, send_to, ref, request, evaluation_ref \\ nil) do
+
-    cast(evaluator, {:handle_intellisense, send_to, ref, request, evaluation_ref})
+  @doc """
+  Builds intellisense context from the given evaluation.
+  """
+  @spec intellisense_context(t(), ref()) :: Livebook.Intellisense.intellisense_context()
+  def intellisense_context(evaluator, ref) do
+    {:dictionary, dictionary} = Process.info(evaluator.pid, :dictionary)
+
+    env =
+      find_in_dictionary(dictionary, {@env_key, ref}) ||
+        find_in_dictionary(dictionary, @initial_env_key)
+
+    map_binding = fn fun -> map_binding(evaluator, ref, fun) end
+
+    %{env: env, map_binding: map_binding}
+  end
+
+  defp find_in_dictionary(dictionary, key) do
+    Enum.find_value(dictionary, fn
+      {^key, value} -> value
+      _pair -> nil
+    end)
+  end
+
+  # Applies the given function to evaluation binding
+  defp map_binding(evaluator, ref, fun) do
+    call(evaluator, {:map_binding, ref, fun})
   end
 
   defp cast(evaluator, message) do
@@ -221,13 +247,16 @@ defmodule Livebook.Evaluator do
   end
 
   defp initial_state(evaluator_ref, formatter, io_proxy, object_tracker) do
+    context = initial_context()
+    Process.put(@initial_env_key, context.env)
+
     %{
       evaluator_ref: evaluator_ref,
       formatter: formatter,
       io_proxy: io_proxy,
       object_tracker: object_tracker,
       contexts: %{},
-      initial_context: initial_context()
+      initial_context: context
     }
   end
 
@@ -245,6 +274,7 @@ defmodule Livebook.Evaluator do
   end
 
   defp initial_context() do
+    # TODO: Use Code.env_for_eval and eval_quoted_with_env on Elixir v1.14+
     env = :elixir.env_for_eval([])
     %{binding: [], env: env, id: random_id()}
   end
@@ -273,7 +303,7 @@ defmodule Livebook.Evaluator do
 
     evaluation_time_ms = get_execution_time_delta(start_time)
 
-    state = put_in(state.contexts[ref], result_context)
+    state = put_context(state, ref, result_context)
 
     Evaluator.IOProxy.flush(state.io_proxy)
     Evaluator.IOProxy.clear_input_cache(state.io_proxy)
@@ -287,30 +317,13 @@ defmodule Livebook.Evaluator do
   end
 
   defp handle_cast({:forget_evaluation, ref}, state) do
-    state = Map.update!(state, :contexts, &Map.delete(&1, ref))
+    state = delete_context(state, ref)
     Evaluator.ObjectTracker.remove_reference(state.object_tracker, {self(), ref})
 
     :erlang.garbage_collect(self())
     {:noreply, state}
   end
 
-  defp handle_cast({:handle_intellisense, send_to, ref, request, evaluation_ref}, state) do
-    context = get_context(state, evaluation_ref)
-
-    # Safely rescue from intellisense errors
-    response =
-      try do
-        Livebook.Intellisense.handle_request(request, context.binding, context.env)
-      rescue
-        error -> Logger.error(Exception.format(:error, error, __STACKTRACE__))
-      end
-
-    send(send_to, {:intellisense_response, ref, request, response})
-
-    :erlang.garbage_collect(self())
-    {:noreply, state}
-  end
-
   defp handle_call({:fetch_evaluation_context, ref, cached_id}, _from, state) do
     context = get_context(state, ref)
 
@@ -334,6 +347,8 @@ defmodule Livebook.Evaluator do
         {:ok, context} ->
           # If the context changed, mirror the process dictionary again
           copy_process_dictionary_from(source_evaluator)
+
+          Process.put(@initial_env_key, context.env)
           put_in(state.initial_context, context)
 
         {:error, :not_modified} ->
@@ -343,6 +358,23 @@ defmodule Livebook.Evaluator do
     {:reply, :ok, state}
   end
 
+  defp handle_call({:map_binding, ref, fun}, _from, state) do
+    context = get_context(state, ref)
+    result = fun.(context.binding)
+    {:reply, result, state}
+  end
+
+  defp put_context(state, ref, context) do
+    Process.put({@env_key, ref}, context.env)
+    put_in(state.contexts[ref], context)
+  end
+
+  defp delete_context(state, ref) do
+    Process.delete({@env_key, ref})
+    {_, state} = pop_in(state.contexts[ref])
+    state
+  end
+
   defp get_context(state, ref) do
     Map.get_lazy(state.contexts, ref, fn -> state.initial_context end)
   end
@@ -352,6 +384,8 @@ defmodule Livebook.Evaluator do
       quoted = Code.string_to_quoted!(code, file: env.file)
       # TODO: Use Code.eval_quoted_with_env/3 on Elixir v1.14
       {result, binding, env} = :elixir.eval_quoted(quoted, binding, env)
+      # TODO: Remove this line on Elixir v1.14 as binding propagates to env correctly
+      {_, binding, env} = :elixir.eval_forms(:ok, binding, env)
 
       {:ok, result, binding, env}
     catch
@@ -396,6 +430,8 @@ defmodule Livebook.Evaluator do
   end
 
   defp internal_dictionary_key?("$" <> _), do: true
+  defp internal_dictionary_key?({@env_key, _ref}), do: true
+  defp internal_dictionary_key?(@initial_env_key), do: true
   defp internal_dictionary_key?(_), do: false
 
   defp get_execution_time_delta(started_at) do

lib/livebook/intellisense.ex

@@ -8,38 +8,49 @@ defmodule Livebook.Intellisense do
   # language server that Livebook uses.
 
   alias Livebook.Intellisense.{IdentifierMatcher, SignatureMatcher, Docs}
+  alias Livebook.Runtime
 
   # Configures width used for inspect and specs formatting.
   @line_length 45
   @extended_line_length 80
 
+  @typedoc """
+  Evaluation state to consider for intellisense.
+
+  The `:map_binding` is only called when a value needs to
+  be extracted from binding.
+  """
+  @type context :: %{
+          env: Macro.Env.t(),
+          map_binding: (Code.binding() -> any())
+        }
+
   @doc """
-  Resolves an intellisense request as defined by `Livebook.Runtime`.
+  Resolves an intellisense request as defined by `Runtime`.
 
   In practice this function simply dispatches the request to one of
   the other public functions in this module.
   """
   @spec handle_request(
-          Livebook.Runtime.intellisense_request(),
+          Runtime.intellisense_request(),
-          Code.binding(),
+          context()
-          Macro.Env.t()
+        ) :: Runtime.intellisense_response()
-        ) :: Livebook.Runtime.intellisense_response()
+  def handle_request(request, context)
-  def handle_request(request, env, binding)
 
-  def handle_request({:completion, hint}, binding, env) do
+  def handle_request({:completion, hint}, context) do
-    items = get_completion_items(hint, binding, env)
+    items = get_completion_items(hint, context)
     %{items: items}
   end
 
-  def handle_request({:details, line, column}, binding, env) do
+  def handle_request({:details, line, column}, context) do
-    get_details(line, column, binding, env)
+    get_details(line, column, context)
   end
 
-  def handle_request({:signature, hint}, binding, env) do
+  def handle_request({:signature, hint}, context) do
-    get_signature_items(hint, binding, env)
+    get_signature_items(hint, context)
   end
 
-  def handle_request({:format, code}, _binding, _env) do
+  def handle_request({:format, code}, _context) do
     case format_code(code) do
       {:ok, code} -> %{code: code}
       :error -> nil
@@ -66,10 +77,9 @@ defmodule Livebook.Intellisense do
   @doc """
   Returns information about signatures matching the given `hint`.
   """
-  @spec get_signature_items(String.t(), Code.binding(), Macro.Env.t()) ::
+  @spec get_signature_items(String.t(), context()) :: Runtime.signature_response() | nil
-          Runtime.signature_response() | nil
+  def get_signature_items(hint, context) do
-  def get_signature_items(hint, binding, env) do
+    case SignatureMatcher.get_matching_signatures(hint, context) do
-    case SignatureMatcher.get_matching_signatures(hint, binding, env) do
       {:ok, [], _active_argument} ->
         nil
 
@@ -108,10 +118,9 @@ defmodule Livebook.Intellisense do
   @doc """
   Returns a list of completion suggestions for the given `hint`.
   """
-  @spec get_completion_items(String.t(), Code.binding(), Macro.Env.t()) ::
+  @spec get_completion_items(String.t(), context()) :: list(Runtime.completion_item())
-          list(Livebook.Runtime.completion_item())
+  def get_completion_items(hint, context) do
-  def get_completion_items(hint, binding, env) do
+    IdentifierMatcher.completion_identifiers(hint, context)
-    IdentifierMatcher.completion_identifiers(hint, binding, env)
     |> Enum.filter(&include_in_completion?/1)
     |> Enum.map(&format_completion_item/1)
     |> Enum.concat(extra_completion_items(hint))
@@ -128,7 +137,7 @@ defmodule Livebook.Intellisense do
 
   defp include_in_completion?(_), do: true
 
-  defp format_completion_item({:variable, name, _value}),
+  defp format_completion_item({:variable, name}),
     do: %{
       label: Atom.to_string(name),
       kind: :variable,
@@ -137,7 +146,7 @@ defmodule Livebook.Intellisense do
       insert_text: Atom.to_string(name)
     }
 
-  defp format_completion_item({:map_field, name, _value}),
+  defp format_completion_item({:map_field, name}),
     do: %{
       label: Atom.to_string(name),
       kind: :field,
@@ -340,10 +349,9 @@ defmodule Livebook.Intellisense do
   Returns detailed information about an identifier located
   in `column` in `line`.
   """
-  @spec get_details(String.t(), pos_integer(), Code.binding(), Macro.Env.t()) ::
+  @spec get_details(String.t(), pos_integer(), context()) :: Runtime.details_response() | nil
-          Livebook.Runtime.details_response() | nil
+  def get_details(line, column, context) do
-  def get_details(line, column, binding, env) do
+    case IdentifierMatcher.locate_identifier(line, column, context) do
-    case IdentifierMatcher.locate_identifier(line, column, binding, env) do
       %{matches: []} ->
         nil
 
@@ -357,9 +365,9 @@ defmodule Livebook.Intellisense do
     end
   end
 
-  defp format_details_item({:variable, name, _value}), do: code(name)
+  defp format_details_item({:variable, name}), do: code(name)
 
-  defp format_details_item({:map_field, name, _value}), do: code(name)
+  defp format_details_item({:map_field, name}), do: code(name)
 
   defp format_details_item({:in_struct_field, _struct, name, default}) do
     join_with_divider([

lib/livebook/intellisense/identifier_matcher.ex

@@ -13,14 +13,15 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
   # which is a very extensive implementation used in the
   # Elixir Language Server.
 
+  alias Livebook.Intellisense
   alias Livebook.Intellisense.Docs
 
   @typedoc """
   A single identifier together with relevant information.
   """
   @type identifier_item ::
-          {:variable, name(), value()}
+          {:variable, name()}
-          | {:map_field, name(), value()}
+          | {:map_field, name()}
           | {:in_struct_field, module(), name(), default :: value()}
           | {:module, module(), display_name(), Docs.documentation()}
           | {:function, module(), name(), arity(), function_type(), display_name(),
@@ -45,15 +46,13 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
 
   `hint` may be a single token or line fragment like `if Enum.m`.
   """
-  @spec completion_identifiers(String.t(), Code.binding(), Macro.Env.t()) ::
+  @spec completion_identifiers(String.t(), Intellisense.context()) :: list(identifier_item())
-          list(identifier_item())
+  def completion_identifiers(hint, intellisense_context) do
-  def completion_identifiers(hint, binding, env) do
     context = Code.Fragment.cursor_context(hint)
 
     ctx = %{
       fragment: hint,
-      binding: binding,
+      intellisense_context: intellisense_context,
-      env: env,
       matcher: @prefix_matcher,
       type: :completion
     }
@@ -68,20 +67,19 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
   The function returns range of columns where the identifier
   is located and a list of matching identifier items.
   """
-  @spec locate_identifier(String.t(), pos_integer(), Code.binding(), Macro.Env.t()) ::
+  @spec locate_identifier(String.t(), pos_integer(), Intellisense.context()) ::
           %{
             matches: list(identifier_item()),
             range: nil | %{from: pos_integer(), to: pos_integer()}
           }
-  def locate_identifier(line, column, binding, env) do
+  def locate_identifier(line, column, intellisense_context) do
     case Code.Fragment.surround_context(line, {1, column}) do
       %{context: context, begin: {_, from}, end: {_, to}} ->
         fragment = String.slice(line, 0, to - 1)
 
         ctx = %{
           fragment: fragment,
-          binding: binding,
+          intellisense_context: intellisense_context,
-          env: env,
           matcher: @exact_matcher,
           type: :locate
         }
@@ -172,10 +170,6 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
     end
   end
 
-  defp expand_dot_path({:var, var}, ctx) do
-    Keyword.fetch(ctx.binding, List.to_atom(var))
-  end
-
   defp expand_dot_path({:alias, alias}, ctx) do
     {:ok, expand_alias(List.to_string(alias), ctx)}
   end
@@ -188,10 +182,36 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
     :error
   end
 
-  defp expand_dot_path({:dot, parent, call}, ctx) do
+  defp expand_dot_path(path, ctx) do
-    case expand_dot_path(parent, ctx) do
+    with {:ok, path} <- recur_expand_dot_path(path, []) do
-      {:ok, %{} = map} -> Map.fetch(map, List.to_atom(call))
+      value_from_binding(path, ctx)
-      _ -> :error
+    end
+  end
+
+  defp recur_expand_dot_path({:var, var}, path) do
+    {:ok, [List.to_atom(var) | path]}
+  end
+
+  defp recur_expand_dot_path({:dot, parent, call}, path) do
+    recur_expand_dot_path(parent, [List.to_atom(call) | path])
+  end
+
+  defp recur_expand_dot_path(_, _path) do
+    :error
+  end
+
+  defp value_from_binding([var | map_path], ctx) do
+    if Macro.Env.has_var?(ctx.intellisense_context.env, {var, nil}) do
+      ctx.intellisense_context.map_binding.(fn binding ->
+        value = Keyword.fetch(binding, var)
+
+        Enum.reduce(map_path, value, fn
+          key, {:ok, map} when is_map(map) -> Map.fetch(map, key)
+          _key, _acc -> :error
+        end)
+      end)
+    else
+      :error
     end
   end
 
@@ -286,7 +306,7 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
 
   defp match_local(hint, ctx) do
     imports =
-      ctx.env
+      ctx.intellisense_context.env
       |> imports_from_env()
       |> Enum.flat_map(fn {mod, funs} ->
         match_module_function(mod, hint, ctx, funs)
@@ -298,20 +318,19 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
   end
 
   defp match_variable(hint, ctx) do
-    for {key, value} <- ctx.binding,
+    for {var, nil} <- Macro.Env.vars(ctx.intellisense_context.env),
-        is_atom(key),
+        name = Atom.to_string(var),
-        name = Atom.to_string(key),
         ctx.matcher.(name, hint),
-        do: {:variable, key, value}
+        do: {:variable, var}
   end
 
   defp match_map_field(map, hint, ctx) do
     # Note: we need Map.to_list/1 in case this is a struct
-    for {key, value} <- Map.to_list(map),
+    for {key, _value} <- Map.to_list(map),
         is_atom(key),
         name = Atom.to_string(key),
         ctx.matcher.(name, hint),
-        do: {:map_field, key, value}
+        do: {:map_field, key}
   end
 
   defp match_sigil(hint, ctx) do
@@ -329,14 +348,14 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
     for mod <- get_matching_modules(hint, ctx),
         usable_as_unquoted_module?(mod),
         name = ":" <> Atom.to_string(mod),
-        do: {:module, mod, name, Livebook.Intellisense.Docs.get_module_documentation(mod)}
+        do: {:module, mod, name, Intellisense.Docs.get_module_documentation(mod)}
   end
 
   # Converts alias string to module atom with regard to the given env
   defp expand_alias(alias, ctx) do
     [name | rest] = alias |> String.split(".") |> Enum.map(&String.to_atom/1)
 
-    case Keyword.fetch(ctx.env.aliases, Module.concat(Elixir, name)) do
+    case Macro.Env.fetch_alias(ctx.intellisense_context.env, name) do
       {:ok, name} when rest == [] -> name
       {:ok, name} -> Module.concat([name | rest])
       :error -> Module.concat([name | rest])
@@ -344,10 +363,10 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
   end
 
   defp match_env_alias(hint, ctx) do
-    for {alias, mod} <- ctx.env.aliases,
+    for {alias, mod} <- ctx.intellisense_context.env.aliases,
         [name] = Module.split(alias),
         ctx.matcher.(name, hint),
-        do: {:module, mod, name, Livebook.Intellisense.Docs.get_module_documentation(mod)}
+        do: {:module, mod, name, Intellisense.Docs.get_module_documentation(mod)}
   end
 
   defp match_module(base_mod, hint, nested?, ctx) do
@@ -396,7 +415,7 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
         valid_alias_piece?("." <> name),
         mod = Module.concat(parent_mod_parts ++ name_parts),
         uniq: true,
-        do: {:module, mod, name, Livebook.Intellisense.Docs.get_module_documentation(mod)}
+        do: {:module, mod, name, Intellisense.Docs.get_module_documentation(mod)}
   end
 
   defp valid_alias_piece?(<<?., char, rest::binary>>) when char in ?A..?Z,
@@ -460,7 +479,7 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
         end)
 
       doc_items =
-        Livebook.Intellisense.Docs.lookup_module_members(
+        Intellisense.Docs.lookup_module_members(
           mod,
           Enum.map(matching_funs, &Tuple.delete_at(&1, 2)),
           kinds: [:function, :macro]
@@ -508,8 +527,7 @@ defmodule Livebook.Intellisense.IdentifierMatcher do
         ctx.matcher.(name, hint)
       end)
 
-    doc_items =
+    doc_items = Intellisense.Docs.lookup_module_members(mod, matching_types, kinds: [:type])
-      Livebook.Intellisense.Docs.lookup_module_members(mod, matching_types, kinds: [:type])
 
     Enum.map(matching_types, fn {name, arity} ->
       doc_item =

lib/livebook/intellisense/signature_matcher.ex

@@ -15,10 +15,12 @@ defmodule Livebook.Intellisense.SignatureMatcher do
   Evaluation binding and environment is used to expand aliases,
   imports, access variable values, etc.
   """
-  @spec get_matching_signatures(String.t(), Code.binding(), Macro.Env.t()) ::
+  @spec get_matching_signatures(String.t(), Livebook.Intellisense.intellisense_context()) ::
           {:ok, list(signature_info()), active_argument :: non_neg_integer()} | :error
-  def get_matching_signatures(hint, binding, env) do
+  def get_matching_signatures(hint, intellisense_context) do
-    case matching_call(hint, binding, env) do
+    %{env: env} = intellisense_context
+
+    case matching_call(hint, intellisense_context) do
       {:ok, {:remote, mod, fun}, maybe_arity, active_argument} ->
         funs = [{fun, maybe_arity || :any}]
         signature_infos = signature_infos_for_members(mod, funs, active_argument)
@@ -107,22 +109,21 @@ defmodule Livebook.Intellisense.SignatureMatcher do
   end
 
   # Returns {:ok, call, exact_arity_or_nil, active_argument} or :error
-  defp matching_call(hint, binding, env) do
+  defp matching_call(hint, intellisense_context) do
     with {:ok, call_target, active_argument} <- call_target_and_argument(hint) do
       case call_target do
         local when is_atom(local) ->
           {:ok, {:local, local}, nil, active_argument}
 
         {:., _, [{:__aliases__, _, _} = alias, fun]} when is_atom(fun) ->
-          alias = Macro.expand(alias, env)
+          alias = Macro.expand(alias, intellisense_context.env)
           {:ok, {:remote, alias, fun}, nil, active_argument}
 
         {:., _, [mod, fun]} when is_atom(mod) and is_atom(fun) ->
           {:ok, {:remote, mod, fun}, nil, active_argument}
 
         {:., _, [{var, _, context}]} when is_atom(var) and is_atom(context) ->
-          case Keyword.fetch(binding, var) do
+          with {:ok, fun} <- function_from_binding(var, intellisense_context) do
-            {:ok, fun} when is_function(fun) ->
             info = :erlang.fun_info(fun)
 
             case info[:type] do
@@ -132,9 +133,6 @@ defmodule Livebook.Intellisense.SignatureMatcher do
               :local ->
                 {:ok, {:anonymous, var}, info[:arity], active_argument}
             end
-
-            :error ->
-              :error
           end
 
         _ ->
@@ -208,4 +206,17 @@ defmodule Livebook.Intellisense.SignatureMatcher do
       |> Macro.prewalker()
       |> Enum.any?(&match?({:__cursor__, _, []}, &1))
   end
+
+  defp function_from_binding(var, intellisense_context) do
+    if Macro.Env.has_var?(intellisense_context.env, {var, nil}) do
+      intellisense_context.map_binding.(fn binding ->
+        case Keyword.fetch(binding, var) do
+          {:ok, fun} when is_function(fun) -> {:ok, fun}
+          _ -> :error
+        end
+      end)
+    else
+      :error
+    end
+  end
 end

lib/livebook/runtime/erl_dist/runtime_server.ex

@@ -246,17 +246,17 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     {container_ref, evaluation_ref} = locator
     evaluator = state.evaluators[container_ref]
 
-    if evaluator != nil and elem(request, 0) not in [:format] do
+    intellisense_context =
-      Evaluator.handle_intellisense(evaluator, send_to, ref, request, evaluation_ref)
+      if evaluator == nil or elem(request, 0) in [:format] do
+        Evaluator.intellisense_context()
       else
-      # Handle the request in a temporary process using an empty evaluation context
+        Evaluator.intellisense_context(evaluator, evaluation_ref)
+      end
+
     Task.Supervisor.start_child(state.task_supervisor, fn ->
-        binding = []
+      response = Livebook.Intellisense.handle_request(request, intellisense_context)
-        env = :elixir.env_for_eval([])
-        response = Livebook.Intellisense.handle_request(request, binding, env)
       send(send_to, {:intellisense_response, ref, request, response})
     end)
-    end
 
     {:noreply, state}
   end

test/livebook/evaluator_test.exs

@@ -236,38 +236,6 @@ defmodule Livebook.EvaluatorTest do
     end
   end
 
-  describe "handle_intellisense/5 given completion request" do
-    test "sends completion response to the given process", %{evaluator: evaluator} do
-      request = {:completion, "System.ver"}
-      Evaluator.handle_intellisense(evaluator, self(), :ref, request)
-
-      assert_receive {:intellisense_response, :ref, ^request, %{items: [%{label: "version/0"}]}},
-                     2_000
-    end
-
-    test "given evaluation reference uses its bindings and env", %{evaluator: evaluator} do
-      code = """
-      alias IO.ANSI
-      number = 10
-      """
-
-      Evaluator.evaluate_code(evaluator, self(), code, :code_1)
-      assert_receive {:evaluation_response, :code_1, _, metadata()}
-
-      request = {:completion, "num"}
-      Evaluator.handle_intellisense(evaluator, self(), :ref, request, :code_1)
-
-      assert_receive {:intellisense_response, :ref, ^request, %{items: [%{label: "number"}]}},
-                     2_000
-
-      request = {:completion, "ANSI.brigh"}
-      Evaluator.handle_intellisense(evaluator, self(), :ref, request, :code_1)
-
-      assert_receive {:intellisense_response, :ref, ^request, %{items: [%{label: "bright/0"}]}},
-                     2_000
-    end
-  end
-
   describe "initialize_from/3" do
     setup %{object_tracker: object_tracker} do
       {:ok, _pid, parent_evaluator} =

test/livebook/runtime/erl_dist/runtime_server_test.exs

@@ -139,13 +139,23 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServerTest do
     end
 
     test "provides extended completion when previous evaluation reference is given", %{pid: pid} do
-      RuntimeServer.evaluate_code(pid, "number = 10", {:c1, :e1}, {:c1, nil})
+      code = """
+      alias IO.ANSI
+      number = 10
+      """
+
+      RuntimeServer.evaluate_code(pid, code, {:c1, :e1}, {:c1, nil})
       assert_receive {:evaluation_response, :e1, _, %{evaluation_time_ms: _time_ms}}
 
       request = {:completion, "num"}
       RuntimeServer.handle_intellisense(pid, self(), :ref, request, {:c1, :e1})
 
       assert_receive {:intellisense_response, :ref, ^request, %{items: [%{label: "number"}]}}
+
+      request = {:completion, "ANSI.brigh"}
+      RuntimeServer.handle_intellisense(pid, self(), :ref, request, {:c1, :e1})
+
+      assert_receive {:intellisense_response, :ref, ^request, %{items: [%{label: "bright/0"}]}}
     end
   end