Add support for scanning smart cell evaluation result (#1048)

* Add support for scanning smart cell evaluation result

* Don't allow moving an evaluating cell across sections

* Use consistent naming

lib/livebook/runtime.ex

@@ -286,6 +286,9 @@ defprotocol Livebook.Runtime do
     * `:file` - the file considered as the source during evaluation.
       This information is relevant for errors formatting and imparts
       the value of `__DIR__`
+
+    * `:smart_cell_ref` - a reference of the smart cell which code is
+      to be evaluated, if applicable
   """
   @spec evaluate_code(t(), String.t(), locator(), locator(), keyword()) :: :ok
   def evaluate_code(runtime, code, locator, base_locator, opts \\ [])

lib/livebook/runtime/erl_dist/runtime_server.ex

@@ -221,11 +221,11 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
      |> handle_down_scan_binding(message)}
   end
 
-  def handle_info({:evaluation_finished, pid, evaluation_ref}, state) do
+  def handle_info({:evaluation_finished, locator}, state) do
     {:noreply,
      state
      |> report_smart_cell_definitions()
-     |> scan_binding_after_evaluation(pid, evaluation_ref)}
+     |> scan_binding_after_evaluation(locator)}
   end
 
   def handle_info(:memory_usage, state) do
@@ -286,7 +286,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
   end
 
   def handle_cast(
-        {:evaluate_code, code, {container_ref, evaluation_ref}, base_locator, opts},
+        {:evaluate_code, code, {container_ref, evaluation_ref} = locator, base_locator, opts},
         state
       ) do
     state = ensure_evaluator(state, container_ref)
@@ -306,7 +306,23 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
           nil
       end
 
-    opts = Keyword.put(opts, :notify_to, self())
+    {smart_cell_ref, opts} = Keyword.pop(opts, :smart_cell_ref)
+    smart_cell_info = smart_cell_ref && state.smart_cells[smart_cell_ref]
+
+    myself = self()
+
+    opts =
+      Keyword.put(opts, :on_finish, fn result ->
+        with %{scan_eval_result: scan_eval_result} when scan_eval_result != nil <- smart_cell_info do
+          try do
+            smart_cell_info.scan_eval_result.(smart_cell_info.pid, result)
+          rescue
+            error -> Logger.error("scanning evaluation result raised an error: #{inspect(error)}")
+          end
+        end
+
+        send(myself, {:evaluation_finished, locator})
+      end)
 
     Evaluator.evaluate_code(
       state.evaluators[container_ref],
@@ -360,7 +376,12 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
              {definition.module, %{ref: ref, attrs: attrs, target_pid: state.owner}}
            ) do
         {:ok, pid, info} ->
-          %{js_view: js_view, source: source, scan_binding: scan_binding} = info
+          %{
+            js_view: js_view,
+            source: source,
+            scan_binding: scan_binding,
+            scan_eval_result: scan_eval_result
+          } = info
 
           send(
             state.owner,
@@ -372,7 +393,8 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
             scan_binding: scan_binding,
             base_locator: base_locator,
             scan_binding_pending: false,
-            scan_binding_monitor_ref: nil
+            scan_binding_monitor_ref: nil,
+            scan_eval_result: scan_eval_result
           }
 
           info = scan_binding_async(ref, info, state)
@@ -543,12 +565,7 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServer do
     end)
   end
 
-  defp scan_binding_after_evaluation(state, pid, evaluation_ref) do
-    {container_ref, _} =
-      Enum.find(state.evaluators, fn {_container_ref, evaluator} -> evaluator.pid == pid end)
-
-    locator = {container_ref, evaluation_ref}
-
+  defp scan_binding_after_evaluation(state, locator) do
     update_in(state.smart_cells, fn smart_cells ->
       Map.map(smart_cells, fn
         {ref, %{base_locator: ^locator} = info} -> scan_binding_async(ref, info, state)

lib/livebook/runtime/evaluator.ex

@@ -38,7 +38,7 @@ defmodule Livebook.Runtime.Evaluator do
   @typedoc """
   An evaluation context.
 
-  Each evaluation produces a new context, which may be optionally
+  Each evaluation produces a  new context, which may be optionally
   used by a later evaluation.
   """
   @type context :: %{binding: Code.binding(), env: Macro.Env.t(), id: binary()}
@@ -49,10 +49,10 @@ defmodule Livebook.Runtime.Evaluator do
   @type ref :: term()
 
   @typedoc """
-  An evaluation response, either the resulting value or an error
-  if raised.
+  An evaluation result, either the return value or an error if
+  raised.
   """
-  @type evaluation_response ::
+  @type evaluation_result ::
           {:ok, result :: any()}
           | {:error, Exception.kind(), error :: any(), Exception.stacktrace()}
 
@@ -75,7 +75,7 @@ defmodule Livebook.Runtime.Evaluator do
       events to. Defaults to the value of `:send_to`
 
     * `:formatter` - a module implementing the `Livebook.Runtime.Evaluator.Formatter`
-      behaviour, used for transforming evaluation response before sending
+      behaviour, used for transforming evaluation result before sending
       it to the client. Defaults to identity
   """
   @spec start_link(keyword()) :: {:ok, pid(), t()} | {:error, term()}
@@ -115,14 +115,14 @@ defmodule Livebook.Runtime.Evaluator do
   Asynchronously parses and evaluates the given code.
 
   Any exceptions are captured and transformed into an error
-  response.
+  result.
 
   The resulting contxt (binding and env) is stored under `ref`.
   Any subsequent calls may specify `base_ref` pointing to a
   previous evaluation, in which case the corresponding context
   is used as the entry point for evaluation.
 
-  The evaluation response is transformed with the configured
+  The evaluation result is transformed with the configured
   formatter send to the configured client (see `start_link/1`).
 
   See `Livebook.Runtime.evaluate_code/5` for the messages format
@@ -130,9 +130,9 @@ defmodule Livebook.Runtime.Evaluator do
 
   ## Options
 
-    * `:notify_to` - a process to be notified about finished
-      evaluation. The notification is sent as a message of the
-      form `{:evaluation_finished, pid, ref}`
+    * `:on_finish` - a function to run when the evaluation is
+      finished. The function receives `t:evaluation_result/0`
+      as an argument
   """
   @spec evaluate_code(t(), String.t(), ref(), ref() | nil, keyword()) :: :ok
   def evaluate_code(evaluator, code, ref, base_ref \\ nil, opts \\ []) when ref != nil do
@@ -319,16 +319,16 @@ defmodule Livebook.Runtime.Evaluator do
     context = put_in(context.env.file, file)
     start_time = System.monotonic_time()
 
-    {result_context, response, code_error} =
+    {result_context, result, code_error} =
       case eval(code, context.binding, context.env) do
-        {:ok, result, binding, env} ->
+        {:ok, value, binding, env} ->
           result_context = %{binding: binding, env: env, id: random_id()}
-          response = {:ok, result}
-          {result_context, response, nil}
+          result = {:ok, value}
+          {result_context, result, nil}
 
         {:error, kind, error, stacktrace, code_error} ->
-          response = {:error, kind, error, stacktrace}
-          {context, response, code_error}
+          result = {:error, kind, error, stacktrace}
+          {context, result, code_error}
       end
 
     evaluation_time_ms = get_execution_time_delta(start_time)
@@ -338,7 +338,7 @@ defmodule Livebook.Runtime.Evaluator do
     Evaluator.IOProxy.flush(state.io_proxy)
     Evaluator.IOProxy.clear_input_cache(state.io_proxy)
 
-    output = state.formatter.format_response(response)
+    output = state.formatter.format_result(result)
 
     metadata = %{
       evaluation_time_ms: evaluation_time_ms,
@@ -348,8 +348,8 @@ defmodule Livebook.Runtime.Evaluator do
 
     send(state.send_to, {:runtime_evaluation_response, ref, output, metadata})
 
-    if notify_to = opts[:notify_to] do
-      send(notify_to, {:evaluation_finished, self(), ref})
+    if on_finish = opts[:on_finish] do
+      on_finish.(result)
     end
 
     :erlang.garbage_collect(self())
@@ -429,11 +429,11 @@ defmodule Livebook.Runtime.Evaluator do
     try 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)
+      {value, 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}
+      {:ok, value, binding, env}
     catch
       kind, error ->
         stacktrace = prune_stacktrace(__STACKTRACE__)

lib/livebook/runtime/evaluator/default_formatter.ex

@@ -10,22 +10,22 @@ defmodule Livebook.Runtime.Evaluator.DefaultFormatter do
   require Logger
 
   @impl true
-  def format_response({:ok, :"do not show this result in output"}) do
+  def format_result({:ok, :"do not show this result in output"}) do
     # Functions in the `IEx.Helpers` module return this specific value
     # to indicate no result should be printed in the iex shell,
     # so we respect that as well.
     :ignored
   end
 
-  def format_response({:ok, {:module, _, _, _} = value}) do
+  def format_result({:ok, {:module, _, _, _} = value}) do
     to_inspect_output(value, limit: 10)
   end
 
-  def format_response({:ok, value}) do
+  def format_result({:ok, value}) do
     to_output(value)
   end
 
-  def format_response({:error, kind, error, stacktrace}) do
+  def format_result({:error, kind, error, stacktrace}) do
     formatted = format_error(kind, error, stacktrace)
     {:error, formatted, error_type(error)}
   end

lib/livebook/runtime/evaluator/formatter.ex

@@ -3,21 +3,20 @@ defmodule Livebook.Runtime.Evaluator.Formatter do
 
   # Behaviour defining how evaluation results are transformed.
   #
-  # The evaluation response is sent to the client as a message
-  # and it may potentially be huge. If the client eventually
-  # converts the result into some smaller representation,
-  # we would unnecessarily send a lot of data.
-  # By defining a custom formatter the client can instruct
-  # the `Evaluator` to send already transformed data.
+  # The evaluation result is sent to the client as a message and it
+  # may potentially be huge. If the client eventually converts the
+  # result into some smaller representation, we would unnecessarily
+  # send a lot of data. By defining a custom formatter the client can
+  # instruct the `Evaluator` to send already transformed data.
   #
-  # Additionally if the results rely on external package installed
-  # in the runtime node, then formatting anywhere else wouldn't be accurate,
-  # for example using `inspect` on an external struct.
+  # Additionally, if the results rely on external package installed
+  # in the runtime node, then formatting anywhere else wouldn't be
+  # accurate, for example using `inspect` on an external struct.
 
   alias Livebook.Runtime.Evaluator
 
   @doc """
-  Transforms the evaluation response.
+  Transforms the evaluation result.
   """
-  @callback format_response(Evaluator.evaluation_response()) :: term()
+  @callback format_result(Evaluator.evaluation_result()) :: term()
 end

lib/livebook/runtime/evaluator/identity_formatter.ex

@@ -6,5 +6,5 @@ defmodule Livebook.Runtime.Evaluator.IdentityFormatter do
   @behaviour Livebook.Runtime.Evaluator.Formatter
 
   @impl true
-  def format_response(evaluation_response), do: evaluation_response
+  def format_result(evaluation_response), do: evaluation_response
 end

lib/livebook/session.ex

@@ -1209,7 +1209,14 @@ defmodule Livebook.Session do
       end
 
     file = path <> "#cell"
-    opts = [file: file]
+
+    smart_cell_ref =
+      case cell do
+        %Cell.Smart{} -> cell.id
+        _ -> nil
+      end
+
+    opts = [file: file, smart_cell_ref: smart_cell_ref]
 
     locator = {container_ref_for_section(section), cell.id}
     base_locator = find_base_locator(state.data, cell, section)

lib/livebook/session/data.ex

@@ -395,8 +395,9 @@ defmodule Livebook.Session.Data do
   end
 
   def apply_operation(data, {:move_cell, _client_pid, id, offset}) do
-    with {:ok, cell, _section} <- Notebook.fetch_cell_and_section(data.notebook, id),
-         true <- offset != 0 do
+    with {:ok, cell, section} <- Notebook.fetch_cell_and_section(data.notebook, id),
+         true <- offset != 0,
+         true <- can_move_cell_by?(data, cell, section, offset) do
       data
       |> with_actions()
       |> move_cell(cell, offset)
@@ -843,6 +844,18 @@ defmodule Livebook.Session.Data do
     |> set!(bin_entries: List.delete(data.bin_entries, cell_bin_entry))
   end
 
+  defp can_move_cell_by?(data, cell, section, offset) do
+    case data.cell_infos[cell.id] do
+      %{eval: %{status: :evaluating}} ->
+        notebook = Notebook.move_cell(data.notebook, cell.id, offset)
+        {:ok, _cell, new_section} = Notebook.fetch_cell_and_section(notebook, cell.id)
+        section.id == new_section.id
+
+      _info ->
+        true
+    end
+  end
+
   defp move_cell({data, _} = data_actions, cell, offset) do
     updated_notebook = Notebook.move_cell(data.notebook, cell.id, offset)
 

test/livebook/runtime/erl_dist/runtime_server_test.exs

@@ -218,7 +218,8 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServerTest do
          %{
            js_view: %{ref: info.ref, pid: pid, assets: %{}},
            source: "source",
-           scan_binding: fn pid, _binding, _env -> send(pid, :scan_binding_result) end
+           scan_binding: fn pid, _binding, _env -> send(pid, :scan_binding_ping) end,
+           scan_eval_result: fn pid, _result -> send(pid, :scan_eval_result_ping) end
          }}
       end
 
@@ -251,24 +252,31 @@ defmodule Livebook.Runtime.ErlDist.RuntimeServerTest do
     @tag opts: @opts
     test "once started scans binding and sends the result to the cell server", %{pid: pid} do
       RuntimeServer.start_smart_cell(pid, "dumb", "ref", %{}, {:c1, nil})
-      assert_receive {:smart_cell_debug, "ref", :handle_info, :scan_binding_result}
+      assert_receive {:smart_cell_debug, "ref", :handle_info, :scan_binding_ping}
     end
 
     @tag opts: @opts
     test "scans binding when a new base locator is set", %{pid: pid} do
       RuntimeServer.start_smart_cell(pid, "dumb", "ref", %{}, {:c1, nil})
-      assert_receive {:smart_cell_debug, "ref", :handle_info, :scan_binding_result}
+      assert_receive {:smart_cell_debug, "ref", :handle_info, :scan_binding_ping}
       RuntimeServer.set_smart_cell_base_locator(pid, "ref", {:c2, nil})
-      assert_receive {:smart_cell_debug, "ref", :handle_info, :scan_binding_result}
+      assert_receive {:smart_cell_debug, "ref", :handle_info, :scan_binding_ping}
     end
 
     @tag opts: @opts
     test "scans binding when the base locator is evaluated", %{pid: pid} do
       RuntimeServer.evaluate_code(pid, "1 + 1", {:c1, :e1}, {:c1, nil})
       RuntimeServer.start_smart_cell(pid, "dumb", "ref", %{}, {:c1, :e1})
-      assert_receive {:smart_cell_debug, "ref", :handle_info, :scan_binding_result}
+      assert_receive {:smart_cell_debug, "ref", :handle_info, :scan_binding_ping}
       RuntimeServer.evaluate_code(pid, "1 + 1", {:c1, :e1}, {:c1, nil})
-      assert_receive {:smart_cell_debug, "ref", :handle_info, :scan_binding_result}
+      assert_receive {:smart_cell_debug, "ref", :handle_info, :scan_binding_ping}
+    end
+
+    @tag opts: @opts
+    test "scans evaluation result when the smart cell is evaluated", %{pid: pid} do
+      RuntimeServer.start_smart_cell(pid, "dumb", "ref", %{}, {:c1, nil})
+      RuntimeServer.evaluate_code(pid, "1 + 1", {:c1, :e1}, {:c1, nil}, smart_cell_ref: "ref")
+      assert_receive {:smart_cell_debug, "ref", :handle_info, :scan_eval_result_ping}
     end
   end
 end

test/livebook/runtime/evaluator/default_formatter_test.exs

@@ -5,12 +5,12 @@ defmodule Livebook.Runtime.Evaluator.DefaultFormatterTest do
 
   test "inspects successful results" do
     result = 10
-    assert {:text, "\e[34m10\e[0m"} = DefaultFormatter.format_response({:ok, result})
+    assert {:text, "\e[34m10\e[0m"} = DefaultFormatter.format_result({:ok, result})
   end
 
   test "gracefully handles errors in the inspect protocol" do
     result = %Livebook.TestModules.BadInspect{}
-    assert {:error, error, :other} = DefaultFormatter.format_response({:ok, result})
+    assert {:error, error, :other} = DefaultFormatter.format_result({:ok, result})
     assert error =~ ":bad_return"
   end
 end

test/livebook/session/data_test.exs

@@ -949,6 +949,24 @@ defmodule Livebook.Session.DataTest do
       assert :error = Data.apply_operation(data, operation)
     end
 
+    test "returns an error if the cell is evaluating and would move to a different section" do
+      # In practice we don't want evaluating cells to be moved between
+      # a section and a branching section, however for simplicity we
+      # do the same for other sections
+
+      data =
+        data_after_operations!([
+          {:insert_section, self(), 0, "s1"},
+          {:insert_section, self(), 1, "s2"},
+          {:insert_cell, self(), "s1", 0, :code, "c1", %{}},
+          {:set_runtime, self(), NoopRuntime.new()},
+          {:queue_cells_evaluation, self(), ["c1"]}
+        ])
+
+      operation = {:move_cell, self(), "c1", 1}
+      assert :error = Data.apply_operation(data, operation)
+    end
+
     test "given negative offset moves the cell and marks relevant cells as stale" do
       data =
         data_after_operations!([