From 42d2386a2d4d3d115946fc6d9a98f2df6319f7a4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Tue, 15 Nov 2022 15:13:05 +0100
Subject: [PATCH] Fix reading inputs rendered in the same evaluation (#1531)

---
 lib/livebook/session/data.ex        | 34 +++++++++++++++++++++--------
 test/livebook/session/data_test.exs | 24 ++++++++++++++++++++
 2 files changed, 49 insertions(+), 9 deletions(-)

diff --git a/lib/livebook/session/data.ex b/lib/livebook/session/data.ex
index 0b5197d7d..9de2886d9 100644
--- a/lib/livebook/session/data.ex
+++ b/lib/livebook/session/data.ex
@@ -1018,15 +1018,31 @@ defmodule Livebook.Session.Data do
   end
 
   defp add_cell_output({data, _} = data_actions, cell, output) do
-    data_actions
-    |> set!(
-      notebook: Notebook.add_cell_output(data.notebook, cell.id, output),
-      input_values:
-        {0, output}
-        |> Cell.find_inputs_in_output()
-        |> Map.new(fn attrs -> {attrs.id, attrs.default} end)
-        |> Map.merge(data.input_values)
-    )
+    new_input_values =
+      {0, output}
+      |> Cell.find_inputs_in_output()
+      |> Map.new(fn attrs -> {attrs.id, attrs.default} end)
+
+    {data, _} =
+      data_actions =
+      data_actions
+      |> set!(
+        notebook: Notebook.add_cell_output(data.notebook, cell.id, output),
+        input_values: Map.merge(new_input_values, data.input_values)
+      )
+
+    if data.cell_infos[cell.id].eval.status == :evaluating do
+      # When a cell renders an input, it should also be able to read
+      # that input during the same evaluation, so we make it seem as
+      # if it already existed prior to the evaluation
+
+      data_actions
+      |> update_cell_eval_info!(cell.id, fn eval_info ->
+        update_in(eval_info.data.input_values, &Map.merge(new_input_values, &1))
+      end)
+    else
+      data_actions
+    end
   end
 
   defp finish_cell_evaluation(data_actions, cell, section, metadata) do
diff --git a/test/livebook/session/data_test.exs b/test/livebook/session/data_test.exs
index 4d351291d..b70f982bd 100644
--- a/test/livebook/session/data_test.exs
+++ b/test/livebook/session/data_test.exs
@@ -1819,6 +1819,30 @@ defmodule Livebook.Session.DataTest do
 
       assert {:ok, %{dirty: true}, []} = Data.apply_operation(data, operation)
     end
+
+    test "stores default values for new inputs in the pre-evaluation data" do
+      input = %{id: "i1", type: :text, label: "Text", default: "hey"}
+
+      data =
+        data_after_operations!([
+          {:insert_section, @cid, 0, "s1"},
+          {:insert_cell, @cid, "s1", 0, :code, "c1", %{}},
+          {:set_runtime, @cid, connected_noop_runtime()},
+          evaluate_cells_operations(["setup"]),
+          {:queue_cells_evaluation, @cid, ["c1"]}
+        ])
+
+      operation = {:add_cell_evaluation_output, @cid, "c1", {:input, input}}
+
+      assert {:ok,
+              %{
+                cell_infos: %{
+                  "c1" => %{
+                    eval: %{data: %{input_values: %{"i1" => "hey"}}}
+                  }
+                }
+              }, _} = Data.apply_operation(data, operation)
+    end
   end
 
   describe "apply_operation/2 given :add_cell_evaluation_response" do