From f2b5382ae72ea478dbcc570dfa3921661efab48a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jonatan=20K=C5=82osko?= <jonatanklosko@gmail.com>
Date: Mon, 7 Oct 2024 11:55:53 +0200
Subject: [PATCH] Fix imports and process dictionary being erased after errored
 evaluation (#2822)

---
 lib/livebook/runtime/evaluator.ex        | 16 ++++-
 test/livebook/runtime/evaluator_test.exs | 90 +++++++++++++++++++++++-
 2 files changed, 103 insertions(+), 3 deletions(-)

diff --git a/lib/livebook/runtime/evaluator.ex b/lib/livebook/runtime/evaluator.ex
index 2343b8479..e1b952e4d 100644
--- a/lib/livebook/runtime/evaluator.ex
+++ b/lib/livebook/runtime/evaluator.ex
@@ -468,8 +468,20 @@ defmodule Livebook.Runtime.Evaluator do
           identifiers_used = :unknown
           identifiers_defined = %{}
           identifier_definitions = []
-          # Empty context
-          new_context = initial_context()
+
+          # Mostly empty context, however we keep imports and process
+          # dictionary from the previous context, since these are not
+          # diffed
+          new_context = %{
+            id: random_long_id(),
+            binding: [],
+            env:
+              context.env
+              |> prune_env(%Evaluator.Tracer{})
+              |> Map.replace!(:versioned_vars, %{}),
+            pdict: context.pdict
+          }
+
           {new_context, result, identifiers_used, identifiers_defined, identifier_definitions}
       end
 
diff --git a/test/livebook/runtime/evaluator_test.exs b/test/livebook/runtime/evaluator_test.exs
index 6d2fe89f0..9e0a4f63e 100644
--- a/test/livebook/runtime/evaluator_test.exs
+++ b/test/livebook/runtime/evaluator_test.exs
@@ -811,7 +811,7 @@ defmodule Livebook.Runtime.EvaluatorTest do
       ref = eval_idx
       parent_refs = Enum.to_list((eval_idx - 1)..0//-1)
       Evaluator.evaluate_code(evaluator, :elixir, code, ref, parent_refs)
-      assert_receive {:runtime_evaluation_response, ^ref, terminal_text(_), metadata}
+      assert_receive {:runtime_evaluation_response, ^ref, _output, metadata}
       %{used: metadata.identifiers_used, defined: metadata.identifiers_defined}
     end
 
@@ -1191,6 +1191,94 @@ defmodule Livebook.Runtime.EvaluatorTest do
 
       assert context.pdict == %{x: 1, y: 2}
     end
+
+    test "context merging with errored evaluation", %{evaluator: evaluator} do
+      # This test is similar to the above, but the second evaluation
+      # fails with an error
+
+      """
+      x = 1
+      y = 1
+
+      alias Enum, as: E
+      alias Map, as: M
+
+      require Enum
+
+      import Integer, only: [is_odd: 1, is_even: 1, to_string: 2, to_charlist: 2]
+
+      Process.put(:x, 1)
+      Process.put(:y, 1)
+      Process.put(:z, 1)
+
+      defmodule Livebook.Runtime.EvaluatorTest.Identifiers.ErrorContextMergingOne do
+      end
+      """
+      |> eval(evaluator, 0)
+
+      """
+      y = 2
+
+      alias MapSet, as: M
+
+      require Map
+
+      import Integer, except: [is_even: 1, to_string: 2]
+
+      Process.put(:y, 2)
+      Process.delete(:z)
+
+      defmodule Livebook.Runtime.EvaluatorTest.Identifiers.ErrorContextMergingTwo do
+      end
+
+      raise "oops"
+      """
+      |> eval(evaluator, 1)
+
+      # Evaluation 1 context
+      #
+      # Should be empty, except for imports and process dictionary,
+      # which are not diffed and should be kept from evaluation 0.
+
+      context = Evaluator.get_evaluation_context(evaluator, [1])
+
+      assert Enum.sort(context.binding) == []
+
+      assert Enum.sort(context.env.aliases) == []
+
+      assert Map not in context.env.requires
+      assert Enum not in context.env.requires
+
+      assert [_, _ | _] = context.env.functions[Integer]
+      assert [_, _ | _] = context.env.macros[Integer]
+
+      assert context.env.versioned_vars == %{}
+
+      assert context.env.context_modules == []
+
+      assert context.pdict == %{x: 1, y: 1, z: 1}
+
+      # Merged context
+
+      context = Evaluator.get_evaluation_context(evaluator, [1, 0])
+
+      assert Enum.sort(context.binding) == [x: 1, y: 1]
+
+      assert Enum.sort(context.env.aliases) == [{E, Enum}, {M, Map}]
+
+      assert Enum in context.env.requires
+
+      assert [_, _ | _] = context.env.functions[Integer]
+      assert [_, _ | _] = context.env.macros[Integer]
+
+      assert context.env.versioned_vars == %{{:x, nil} => 0, {:y, nil} => 1}
+
+      assert context.env.context_modules == [
+               Livebook.Runtime.EvaluatorTest.Identifiers.ErrorContextMergingOne
+             ]
+
+      assert context.pdict == %{x: 1, y: 1, z: 1}
+    end
   end
 
   describe "forget_evaluation/2" do