Fix frame updates containing inputs (#2219)

lib/livebook/notebook.ex

@@ -797,6 +797,12 @@ defmodule Livebook.Notebook do
     {{counter, %{output | outputs: outputs}}, counter + 1}
   end
 
+  defp index_output(%{type: :frame_update} = output, counter) do
+    {update_type, new_outputs} = output.update
+    {new_outputs, counter} = index_outputs(new_outputs, counter)
+    {{counter, %{output | update: {update_type, new_outputs}}}, counter + 1}
+  end
+
   defp index_output(output, counter) do
     {{counter, output}, counter + 1}
   end
@@ -804,13 +810,13 @@ defmodule Livebook.Notebook do
   @doc """
   Finds frame outputs matching the given ref.
   """
-  @spec find_frame_outputs(t(), String.t()) :: list(Cell.indexed_output())
+  @spec find_frame_outputs(t(), String.t()) :: list({Cell.indexed_output(), Cell.t()})
   def find_frame_outputs(notebook, frame_ref) do
     for section <- all_sections(notebook),
-        %{outputs: outputs} <- section.cells,
+        %{outputs: outputs} = cell <- section.cells,
         output <- outputs,
         frame_output <- do_find_frame_outputs(output, frame_ref),
-        do: frame_output
+        do: {frame_output, cell}
   end
 
   defp do_find_frame_outputs({_idx, %{type: :frame, ref: ref}} = output, ref) do

lib/livebook/notebook/cell.ex

@@ -64,6 +64,10 @@ defmodule Livebook.Notebook.Cell do
     Enum.flat_map(output.outputs, &find_inputs_in_output/1)
   end
 
+  def find_inputs_in_output({_idx, %{type: :frame_update, update: {_update_type, new_outputs}}}) do
+    Enum.flat_map(new_outputs, &find_inputs_in_output/1)
+  end
+
   def find_inputs_in_output(_output), do: []
 
   @doc """
@@ -78,6 +82,10 @@ defmodule Livebook.Notebook.Cell do
     Enum.flat_map(output.outputs, &find_assets_in_output/1)
   end
 
+  def find_assets_in_output(%{type: :frame_update, update: {_update_type, new_outputs}}) do
+    Enum.flat_map(new_outputs, &find_assets_in_output/1)
+  end
+
   def find_assets_in_output(_output), do: []
 
   @setup_cell_id "setup"

lib/livebook_web/live/app_session_live.ex

@@ -387,11 +387,14 @@ defmodule LivebookWeb.AppSessionLive do
       {:add_cell_evaluation_output, _client_id, _cell_id, %{type: :frame_update} = output} ->
         %{ref: ref, update: {update_type, _}} = output
 
-        for {idx, frame} <- Notebook.find_frame_outputs(data.notebook, ref) do
+        changed_input_ids = Session.Data.changed_input_ids(data)
+
+        for {{idx, frame} = output, _cell} <- Notebook.find_frame_outputs(data.notebook, ref) do
           send_update(LivebookWeb.Output.FrameComponent,
             id: "output-#{idx}",
             outputs: frame.outputs,
-            update_type: update_type
+            update_type: update_type,
+            input_views: input_views_for_output(output, data, changed_input_ids)
           )
         end
 

lib/livebook_web/live/session_live.ex

@@ -2784,13 +2784,19 @@ defmodule LivebookWeb.SessionLive do
       # to the corresponding component, so the DOM patch is isolated and fast.
       # This is important for intensive output updates
       {:add_cell_evaluation_output, _client_id, _cell_id, %{type: :frame_update} = output} ->
-        %{ref: ref, update: {update_type, _}} = output
+        %{ref: ref, update: {update_type, _new_outputs}} = output
 
-        for {idx, frame} <- Notebook.find_frame_outputs(data.notebook, ref) do
+        changed_input_ids = Session.Data.changed_input_ids(data)
+
+        for {{idx, frame}, cell} <- Notebook.find_frame_outputs(data.notebook, ref) do
           send_update(LivebookWeb.Output.FrameComponent,
             id: "output-#{idx}",
             outputs: frame.outputs,
-            update_type: update_type
+            update_type: update_type,
+            # Note that we are not updating data_view to avoid re-render,
+            # but any change that causes frame to re-render will update
+            # data_view first
+            input_views: input_views_for_cell(cell, data, changed_input_ids)
           )
         end
 

test/livebook/notebook_test.exs

@@ -533,24 +533,20 @@ defmodule Livebook.NotebookTest do
     test "returns frame outputs with matching ref" do
       frame_output = {0, %{type: :frame, ref: "1", outputs: [], placeholder: true}}
 
-      notebook = %{
+      cell = %{Cell.new(:code) | outputs: [frame_output]}
-        Notebook.new()
+      notebook = %{Notebook.new() | sections: [%{Section.new() | cells: [cell]}]}
-        | sections: [%{Section.new() | cells: [%{Cell.new(:code) | outputs: [frame_output]}]}]
-      }
 
-      assert [^frame_output] = Notebook.find_frame_outputs(notebook, "1")
+      assert [{^frame_output, ^cell}] = Notebook.find_frame_outputs(notebook, "1")
     end
 
     test "finds a nested frame" do
       nested_frame_output = {0, %{type: :frame, ref: "2", outputs: [], placeholder: true}}
       frame_output = {0, %{type: :frame, ref: "1", outputs: [nested_frame_output]}}
 
-      notebook = %{
+      cell = %{Cell.new(:code) | outputs: [frame_output]}
-        Notebook.new()
+      notebook = %{Notebook.new() | sections: [%{Section.new() | cells: [cell]}]}
-        | sections: [%{Section.new() | cells: [%{Cell.new(:code) | outputs: [frame_output]}]}]
-      }
 
-      assert [^nested_frame_output] = Notebook.find_frame_outputs(notebook, "2")
+      assert [{^nested_frame_output, ^cell}] = Notebook.find_frame_outputs(notebook, "2")
     end
   end
 end

test/livebook_web/live/session_live_test.exs

@@ -761,6 +761,46 @@ defmodule LivebookWeb.SessionLiveTest do
 
       assert render(view) =~ "This input has changed."
     end
+
+    test "frame output update with input", %{conn: conn, session: session, test: test} do
+      Session.subscribe(session.id)
+      evaluate_setup(session.pid)
+
+      section_id = insert_section(session.pid)
+      cell_id = insert_text_cell(session.pid, section_id, :code)
+
+      Session.queue_cell_evaluation(session.pid, cell_id)
+
+      frame = %{type: :frame, ref: "1", outputs: [], placeholder: true}
+      send(session.pid, {:runtime_evaluation_output, cell_id, frame})
+
+      {:ok, view, _} = live(conn, ~p"/sessions/#{session.id}")
+
+      input = %{
+        type: :input,
+        ref: "ref1",
+        id: "input1",
+        destination: test,
+        attrs: %{type: :number, default: 1, label: "Input inside frame"}
+      }
+
+      frame_update = %{
+        type: :frame_update,
+        ref: "1",
+        update: {:replace, [input]}
+      }
+
+      send(session.pid, {:runtime_evaluation_output, cell_id, frame_update})
+
+      wait_for_session_update(session.pid)
+
+      # Render once, so that frame send_update is processed
+      _ = render(view)
+
+      content = render(view)
+      assert content =~ "Input inside frame"
+      assert has_element?(view, ~s/input[value="1"]/)
+    end
   end
 
   describe "smart cells" do