Make failed evaluations preserve evaluation context (#145)

lib/livebook/evaluator.ex

@@ -121,19 +121,22 @@ defmodule Livebook.Evaluator do
     file = Keyword.get(opts, :file, "nofile")
     context = put_in(context.env.file, file)
 
-    case eval(code, context.binding, context.env) do
+    {result_context, response} =
-      {:ok, result, binding, env} ->
+      case eval(code, context.binding, context.env) do
-        result_context = %{binding: binding, env: env}
+        {:ok, result, binding, env} ->
-        new_contexts = Map.put(state.contexts, ref, result_context)
+          result_context = %{binding: binding, env: env}
-        new_state = %{state | contexts: new_contexts}
+          response = {:ok, result}
+          {result_context, response}
 
-        send_evaluation_response(send_to, ref, {:ok, result}, state.formatter)
+        {:error, kind, error, stacktrace} ->
-        {:noreply, new_state}
+          response = {:error, kind, error, stacktrace}
+          {context, response}
+      end
 
-      {:error, kind, error, stacktrace} ->
+    send_evaluation_response(send_to, ref, response, state.formatter)
-        send_evaluation_response(send_to, ref, {:error, kind, error, stacktrace}, state.formatter)
+
-        {:noreply, state}
+    new_state = put_in(state.contexts[ref], result_context)
-    end
+    {:noreply, new_state}
   end
 
   def handle_cast({:forget_evaluation, ref}, state) do

test/livebook/evaluator_test.exs

@@ -109,6 +109,30 @@ defmodule Livebook.EvaluatorTest do
       end)
     end
 
+    test "in case of an error uses own evaluation context as the resulting context",
+         %{evaluator: evaluator} do
+      code1 = """
+      x = 2
+      """
+
+      code2 = """
+      raise ":<"
+      """
+
+      code3 = """
+      x * x
+      """
+
+      Evaluator.evaluate_code(evaluator, self(), code1, :code_1)
+      assert_receive {:evaluation_response, :code_1, {:ok, _}}
+
+      Evaluator.evaluate_code(evaluator, self(), code2, :code_2, :code_1)
+      assert_receive {:evaluation_response, :code_2, {:error, _, _, _}}
+
+      Evaluator.evaluate_code(evaluator, self(), code3, :code_3, :code_2)
+      assert_receive {:evaluation_response, :code_3, {:ok, 4}}
+    end
+
     test "given file option sets it in evaluation environment", %{evaluator: evaluator} do
       code = """
       __DIR__