Hydrate updated smart cell source digest before reevaluation (#1886)

lib/livebook/session.ex

@@ -1453,8 +1453,19 @@ defmodule Livebook.Session do
       {:ok, cell, _section} ->
         chunks = info[:chunks]
         delta = Livebook.JSInterop.diff(cell.source, source)
-        operation = {:update_smart_cell, @client_id, id, attrs, delta, chunks, info.reevaluate}
-        {:noreply, handle_operation(state, operation)}
+        operation = {:update_smart_cell, @client_id, id, attrs, delta, chunks}
+        state = handle_operation(state, operation)
+
+        # Note that we intentionally use a separate operation, so that
+        # the new source digest is already hydrated on the clients
+        state =
+          if info.reevaluate do
+            handle_operation(state, {:queue_smart_cell_reevaluation, @client_id, id})
+          else
+            state
+          end
+
+        {:noreply, state}
 
       :error ->
         {:noreply, state}
@@ -1873,7 +1884,7 @@ defmodule Livebook.Session do
   defp after_operation(
          state,
          _prev_state,
-         {:update_smart_cell, _client_id, cell_id, _attrs, _delta, _chunks, _reevaluate}
+         {:update_smart_cell, _client_id, cell_id, _attrs, _delta, _chunks}
        ) do
     hydrate_cell_source_digest(state, cell_id, :primary)
     state

lib/livebook/session/data.ex

@@ -193,7 +193,8 @@ defmodule Livebook.Session.Data do
           | {:smart_cell_started, client_id(), Cell.id(), Delta.t(), Runtime.chunks() | nil,
              Runtime.js_view(), Runtime.editor() | nil}
           | {:update_smart_cell, client_id(), Cell.id(), Cell.Smart.attrs(), Delta.t(),
-             Runtime.chunks() | nil, reevaluate :: boolean()}
+             Runtime.chunks() | nil}
+          | {:queue_smart_cell_reevaluation, client_id(), Cell.id()}
           | {:smart_cell_down, client_id(), Cell.id()}
           | {:recover_smart_cell, client_id(), Cell.id()}
           | {:erase_outputs, client_id()}
@@ -639,13 +640,12 @@ defmodule Livebook.Session.Data do
     end
   end
 
-  def apply_operation(data, {:update_smart_cell, client_id, id, attrs, delta, chunks, reevaluate}) do
-    with {:ok, %Cell.Smart{} = cell, section} <-
+  def apply_operation(data, {:update_smart_cell, client_id, id, attrs, delta, chunks}) do
+    with {:ok, %Cell.Smart{} = cell, _section} <-
            Notebook.fetch_cell_and_section(data.notebook, id) do
       data
       |> with_actions()
       |> update_smart_cell(cell, client_id, attrs, delta, chunks)
-      |> maybe_queue_updated_smart_cell(cell, section, reevaluate)
       |> set_dirty()
       |> wrap_ok()
     else
@@ -653,6 +653,23 @@ defmodule Livebook.Session.Data do
     end
   end
 
+  def apply_operation(data, {:queue_smart_cell_reevaluation, _client_id, id}) do
+    with {:ok, %Cell.Smart{} = cell, section} <-
+           Notebook.fetch_cell_and_section(data.notebook, id),
+         eval_info <- data.cell_infos[cell.id].eval,
+         :ready <- eval_info.status,
+         true <- eval_info.validity in [:evaluated, :stale] do
+      data
+      |> with_actions()
+      |> queue_prerequisite_cells_evaluation([cell.id])
+      |> queue_cell_evaluation(cell, section)
+      |> maybe_evaluate_queued()
+      |> wrap_ok()
+    else
+      _ -> :error
+    end
+  end
+
   def apply_operation(data, {:smart_cell_down, _client_id, id}) do
     with {:ok, %Cell.Smart{} = cell, _section} <-
            Notebook.fetch_cell_and_section(data.notebook, id) do
@@ -1601,21 +1618,6 @@ defmodule Livebook.Session.Data do
     |> update_cell_info!(cell.id, &%{&1 | status: :down})
   end
 
-  defp maybe_queue_updated_smart_cell({data, _} = data_actions, cell, section, reevaluate) do
-    info = data.cell_infos[cell.id]
-
-    evaluated? = info.eval.status == :ready and info.eval.validity in [:evaluated, :stale]
-
-    if evaluated? and reevaluate do
-      data_actions
-      |> queue_prerequisite_cells_evaluation([cell.id])
-      |> queue_cell_evaluation(cell, section)
-      |> maybe_evaluate_queued()
-    else
-      data_actions
-    end
-  end
-
   defp recover_smart_cell({data, _} = data_actions, cell, section) do
     if Runtime.connected?(data.runtime) do
       start_smart_cell(data_actions, cell, section)

test/livebook/session/data_test.exs

@@ -2856,7 +2856,7 @@ defmodule Livebook.Session.DataTest do
 
       attrs = %{"text" => "content!"}
       delta2 = Delta.new() |> Delta.retain(7) |> Delta.insert("!")
-      operation = {:update_smart_cell, client_id, "c1", attrs, delta2, nil, false}
+      operation = {:update_smart_cell, client_id, "c1", attrs, delta2, nil}
 
       assert {:ok,
               %{
@@ -2867,8 +2867,28 @@ defmodule Livebook.Session.DataTest do
               [{:broadcast_delta, ^client_id, _cell, :primary, ^delta2}]} =
                Data.apply_operation(data, operation)
     end
+  end
 
-    test "queues the cell when already evaluated and reevaluate is specified" do
+  describe "apply_operation/2 given :queue_smart_cell_reevaluation" do
+    test "returns error if the cell is fresh" do
+      client_id = "client1"
+
+      data =
+        data_after_operations!([
+          {:insert_section, @cid, 0, "s1"},
+          {:set_runtime, @cid, connected_noop_runtime()},
+          evaluate_cells_operations(["setup"]),
+          {:set_smart_cell_definitions, @cid, @smart_cell_definitions},
+          {:insert_cell, @cid, "s1", 0, :smart, "c1", %{kind: "text"}},
+          {:smart_cell_started, @cid, "c1", Delta.new(), nil, %{}, nil}
+        ])
+
+      operation = {:queue_smart_cell_reevaluation, client_id, "c1"}
+
+      assert :error = Data.apply_operation(data, operation)
+    end
+
+    test "queues the cell when already evaluated" do
       client_id = "client1"
 
       data =
@@ -2879,11 +2899,10 @@ defmodule Livebook.Session.DataTest do
           {:set_smart_cell_definitions, @cid, @smart_cell_definitions},
           {:insert_cell, @cid, "s1", 0, :smart, "c1", %{kind: "text"}},
           {:smart_cell_started, @cid, "c1", Delta.new(), nil, %{}, nil},
-          {:queue_cells_evaluation, @cid, ["c1"]},
-          {:add_cell_evaluation_response, @cid, "c1", @eval_resp, eval_meta()}
+          evaluate_cells_operations(["c1"])
         ])
 
-      operation = {:update_smart_cell, client_id, "c1", %{}, Delta.new(), nil, true}
+      operation = {:queue_smart_cell_reevaluation, client_id, "c1"}
 
       assert {:ok,
               %{