diff --git a/lib/livebook/runtime/evaluator/io_proxy.ex b/lib/livebook/runtime/evaluator/io_proxy.ex index e5ed12865..d3b85e8e5 100644 --- a/lib/livebook/runtime/evaluator/io_proxy.ex +++ b/lib/livebook/runtime/evaluator/io_proxy.ex @@ -262,6 +262,7 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do {:ok, state} end + # This request is used by Kino <= 0.13.1 defp io_request({:livebook_get_input_value, input_id}, state) do input_cache = Map.put_new_lazy(state.input_cache, input_id, fn -> @@ -271,6 +272,25 @@ defmodule Livebook.Runtime.Evaluator.IOProxy do {input_cache[input_id], %{state | input_cache: input_cache}} end + defp io_request({:livebook_get_input_value, input_id, pid}, state) do + # We only allow reading the input value from the evaluator process, + # to make sure the cell is still evaluating. A different process + # could attempt to read the value after the cell evaluation was + # finished, in which case we could return a cached value from a + # different evaluation and we would not track the read properly. + + if pid == state.evaluator do + input_cache = + Map.put_new_lazy(state.input_cache, input_id, fn -> + request_input_value(input_id, state) + end) + + {input_cache[input_id], %{state | input_cache: input_cache}} + else + {{:error, :bad_process}, state} + end + end + defp io_request({:livebook_get_file_path, file_id}, state) do # We could cache forever, but we don't want the cache to pile up # indefinitely, so we just reuse the input cache which is cleared diff --git a/test/livebook/runtime/evaluator/io_proxy_test.exs b/test/livebook/runtime/evaluator/io_proxy_test.exs index 16b3e8e2d..acc5eef52 100644 --- a/test/livebook/runtime/evaluator/io_proxy_test.exs +++ b/test/livebook/runtime/evaluator/io_proxy_test.exs @@ -18,7 +18,7 @@ defmodule Livebook.Runtime.Evaluator.IOProxyTest do io = Process.info(evaluator.pid)[:group_leader] IOProxy.before_evaluation(io, :ref, "cell") - %{io: io} + %{io: io, evaluator: evaluator} end describe ":stdio interoperability" do @@ -47,7 +47,19 @@ defmodule Livebook.Runtime.Evaluator.IOProxyTest do end describe "input" do - test "responds to Livebook input request", %{io: io} do + test "responds to Livebook input request", %{io: io, evaluator: evaluator} do + pid = + spawn(fn -> + pid = evaluator.pid + assert livebook_get_input_value(io, "input1", pid) == {:ok, :value} + end) + + reply_to_input_request(:ref, "input1", {:ok, :value}, 1) + + await_termination(pid) + end + + test "responds to Livebook input request (legacy)", %{io: io} do pid = spawn(fn -> assert livebook_get_input_value(io, "input1") == {:ok, :value} @@ -58,11 +70,12 @@ defmodule Livebook.Runtime.Evaluator.IOProxyTest do await_termination(pid) end - test "responds to subsequent requests with the same value", %{io: io} do + test "responds to subsequent requests with the same value", %{io: io, evaluator: evaluator} do pid = spawn(fn -> - assert livebook_get_input_value(io, "input1") == {:ok, :value} - assert livebook_get_input_value(io, "input1") == {:ok, :value} + pid = evaluator.pid + assert livebook_get_input_value(io, "input1", pid) == {:ok, :value} + assert livebook_get_input_value(io, "input1", pid) == {:ok, :value} end) reply_to_input_request(:ref, "input1", {:ok, :value}, 1) @@ -70,13 +83,25 @@ defmodule Livebook.Runtime.Evaluator.IOProxyTest do await_termination(pid) end - test "before_evaluation/3 clears all cached input information", %{io: io} do + test "responds with error, when request does not come from the evaluator process", %{io: io} do + pid = + spawn(fn -> + pid = self() + assert livebook_get_input_value(io, "input1", pid) == {:error, :bad_process} + end) + + await_termination(pid) + end + + test "before_evaluation/3 clears all cached input information", + %{io: io, evaluator: evaluator} do pid = spawn_link(fn -> + pid = evaluator.pid IOProxy.before_evaluation(io, :ref, "cell") - assert livebook_get_input_value(io, "input1") == {:ok, :value1} + assert livebook_get_input_value(io, "input1", pid) == {:ok, :value1} IOProxy.before_evaluation(io, :ref, "cell") - assert livebook_get_input_value(io, "input1") == {:ok, :value2} + assert livebook_get_input_value(io, "input1", pid) == {:ok, :value2} end) reply_to_input_request(:ref, "input1", {:ok, :value1}, 1) @@ -177,6 +202,10 @@ defmodule Livebook.Runtime.Evaluator.IOProxyTest do io_request(io, {:livebook_get_input_value, input_id}) end + defp livebook_get_input_value(io, input_id, pid) do + io_request(io, {:livebook_get_input_value, input_id, pid}) + end + defp livebook_generate_token(io) do io_request(io, :livebook_generate_token) end