diff --git a/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/ContextIsEnabledNode.java b/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/ContextIsEnabledNode.java
index 9e19228090bd..b2e74dbf2d72 100644
--- a/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/ContextIsEnabledNode.java
+++ b/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/ContextIsEnabledNode.java
@@ -7,7 +7,6 @@
 import org.enso.interpreter.runtime.data.atom.Atom;
 import org.enso.interpreter.runtime.error.PanicException;
 import org.enso.interpreter.runtime.state.ExecutionEnvironment;
-import org.enso.interpreter.runtime.state.State;
 
 @BuiltinMethod(
     type = "Context",
@@ -16,9 +15,9 @@
 public class ContextIsEnabledNode extends Node {
   private @Child ExpectStringNode expectStringNode = ExpectStringNode.build();
 
-  Object execute(State state, Atom self, Object environmentName) {
+  Object execute(Atom self, Object environmentName) {
     String envName = expectStringNode.execute(environmentName);
-    ExecutionEnvironment currentEnv = state.currentEnvironment();
+    ExecutionEnvironment currentEnv = EnsoContext.get(this).getExecutionEnvironment();
     if (!currentEnv.getName().equals(envName)) {
       Atom error =
           EnsoContext.get(this)
diff --git a/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/RuntimeCurrentExecutionEnvironmentNode.java b/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/RuntimeCurrentExecutionEnvironmentNode.java
index a881e75af174..7fd13a913790 100644
--- a/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/RuntimeCurrentExecutionEnvironmentNode.java
+++ b/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/RuntimeCurrentExecutionEnvironmentNode.java
@@ -2,8 +2,8 @@
 
 import com.oracle.truffle.api.nodes.Node;
 import org.enso.interpreter.dsl.BuiltinMethod;
+import org.enso.interpreter.runtime.EnsoContext;
 import org.enso.interpreter.runtime.data.text.Text;
-import org.enso.interpreter.runtime.state.State;
 
 @BuiltinMethod(
     type = "Runtime",
@@ -11,7 +11,7 @@
     description = "Returns the name of the current execution environment.",
     autoRegister = false)
 public class RuntimeCurrentExecutionEnvironmentNode extends Node {
-  Object execute(State state) {
-    return Text.create(state.currentEnvironment().getName());
+  Object execute() {
+    return Text.create(EnsoContext.get(this).getExecutionEnvironment().getName());
   }
 }
diff --git a/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/RuntimeWithDisabledContextNode.java b/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/RuntimeWithDisabledContextNode.java
index b4a3a8889f6f..5c486b47c1e4 100644
--- a/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/RuntimeWithDisabledContextNode.java
+++ b/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/RuntimeWithDisabledContextNode.java
@@ -7,7 +7,9 @@
 import org.enso.interpreter.node.BaseNode;
 import org.enso.interpreter.node.callable.thunk.ThunkExecutorNode;
 import org.enso.interpreter.node.expression.builtin.text.util.ExpectStringNode;
+import org.enso.interpreter.runtime.EnsoContext;
 import org.enso.interpreter.runtime.data.atom.Atom;
+import org.enso.interpreter.runtime.state.ExecutionEnvironment;
 import org.enso.interpreter.runtime.state.State;
 
 @BuiltinMethod(
@@ -23,7 +25,12 @@ public class RuntimeWithDisabledContextNode extends Node {
   Object execute(
       VirtualFrame frame, State state, Atom context, Object env_name, @Suspend Object action) {
     String envName = expectStringNode.execute(env_name);
-    return thunkExecutorNode.executeThunk(
-        frame, action, state.withContextDisabledIn(context, envName), BaseNode.TailStatus.NOT_TAIL);
+    ExecutionEnvironment original =
+        EnsoContext.get(this).disableExecutionEnvironment(context, envName);
+    try {
+      return thunkExecutorNode.executeThunk(frame, action, state, BaseNode.TailStatus.NOT_TAIL);
+    } finally {
+      EnsoContext.get(this).setExecutionEnvironment(original);
+    }
   }
 }
diff --git a/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/RuntimeWithEnabledContextNode.java b/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/RuntimeWithEnabledContextNode.java
index d280a343f05a..d4fcd4380860 100644
--- a/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/RuntimeWithEnabledContextNode.java
+++ b/engine/runtime/src/main/java/org/enso/interpreter/node/expression/builtin/runtime/RuntimeWithEnabledContextNode.java
@@ -7,7 +7,9 @@
 import org.enso.interpreter.node.BaseNode;
 import org.enso.interpreter.node.callable.thunk.ThunkExecutorNode;
 import org.enso.interpreter.node.expression.builtin.text.util.ExpectStringNode;
+import org.enso.interpreter.runtime.EnsoContext;
 import org.enso.interpreter.runtime.data.atom.Atom;
+import org.enso.interpreter.runtime.state.ExecutionEnvironment;
 import org.enso.interpreter.runtime.state.State;
 
 @BuiltinMethod(
@@ -23,7 +25,12 @@ public class RuntimeWithEnabledContextNode extends Node {
   Object execute(
       VirtualFrame frame, State state, Atom context, Object env_name, @Suspend Object action) {
     String envName = expectStringNode.execute(env_name);
-    return thunkExecutorNode.executeThunk(
-        frame, action, state.withContextEnabledIn(context, envName), BaseNode.TailStatus.NOT_TAIL);
+    ExecutionEnvironment original =
+        EnsoContext.get(this).enableExecutionEnvironment(context, envName);
+    try {
+      return thunkExecutorNode.executeThunk(frame, action, state, BaseNode.TailStatus.NOT_TAIL);
+    } finally {
+      EnsoContext.get(this).setExecutionEnvironment(original);
+    }
   }
 }
diff --git a/engine/runtime/src/main/java/org/enso/interpreter/runtime/EnsoContext.java b/engine/runtime/src/main/java/org/enso/interpreter/runtime/EnsoContext.java
index db7ce623b893..dc4dc3f4a465 100644
--- a/engine/runtime/src/main/java/org/enso/interpreter/runtime/EnsoContext.java
+++ b/engine/runtime/src/main/java/org/enso/interpreter/runtime/EnsoContext.java
@@ -51,6 +51,7 @@
 import org.enso.interpreter.OptionsHelper;
 import org.enso.interpreter.runtime.builtin.Builtins;
 import org.enso.interpreter.runtime.data.Type;
+import org.enso.interpreter.runtime.data.atom.Atom;
 import org.enso.interpreter.runtime.error.DataflowError;
 import org.enso.interpreter.runtime.error.PanicException;
 import org.enso.interpreter.runtime.instrument.NotificationHandler;
@@ -875,6 +876,34 @@ public void setExecutionEnvironment(ExecutionEnvironment executionEnvironment) {
     this.executionEnvironment = executionEnvironment;
   }
 
+  /**
+   * Enable execution context in the execution environment.
+   *
+   * @param context the execution context
+   * @param environmentName the execution environment name
+   */
+  public ExecutionEnvironment enableExecutionEnvironment(Atom context, String environmentName) {
+    ExecutionEnvironment original = executionEnvironment;
+    if (executionEnvironment.getName().equals(environmentName)) {
+      executionEnvironment = executionEnvironment.withContextEnabled(context);
+    }
+    return original;
+  }
+
+  /**
+   * Enable execution context in the execution environment.
+   *
+   * @param context the execution context
+   * @param environmentName the execution environment name
+   */
+  public ExecutionEnvironment disableExecutionEnvironment(Atom context, String environmentName) {
+    ExecutionEnvironment original = executionEnvironment;
+    if (executionEnvironment.getName().equals(environmentName)) {
+      executionEnvironment = executionEnvironment.withContextDisabled(context);
+    }
+    return original;
+  }
+
   /** Returns a maximal number of warnings that can be attached to a value */
   public int getWarningsLimit() {
     return this.warningsLimit;
diff --git a/engine/runtime/src/main/java/org/enso/interpreter/runtime/error/DataflowError.java b/engine/runtime/src/main/java/org/enso/interpreter/runtime/error/DataflowError.java
index c87e494e00b0..1d65eac65161 100644
--- a/engine/runtime/src/main/java/org/enso/interpreter/runtime/error/DataflowError.java
+++ b/engine/runtime/src/main/java/org/enso/interpreter/runtime/error/DataflowError.java
@@ -54,7 +54,10 @@ public final class DataflowError extends AbstractTruffleException implements Ens
   public static DataflowError withDefaultTrace(State state, Object payload, Node location) {
     assert payload != null;
     boolean attachFullStackTrace =
-        state == null ? true : state.currentEnvironment().hasContextEnabled("Dataflow_Stack_Trace");
+        state == null
+            || EnsoContext.get(location)
+                .getExecutionEnvironment()
+                .hasContextEnabled("Dataflow_Stack_Trace");
     if (attachFullStackTrace) {
       var result = new DataflowError(payload, UNLIMITED_STACK_TRACE, location);
       TruffleStackTrace.fillIn(result);
diff --git a/engine/runtime/src/main/java/org/enso/interpreter/runtime/state/State.java b/engine/runtime/src/main/java/org/enso/interpreter/runtime/state/State.java
index bd8687b60b31..20699b105332 100644
--- a/engine/runtime/src/main/java/org/enso/interpreter/runtime/state/State.java
+++ b/engine/runtime/src/main/java/org/enso/interpreter/runtime/state/State.java
@@ -3,44 +3,20 @@
 import com.oracle.truffle.api.object.DynamicObject;
 import com.oracle.truffle.api.object.Shape;
 import org.enso.interpreter.runtime.EnsoContext;
-import org.enso.interpreter.runtime.data.atom.Atom;
 
 public class State {
   private final Container container;
 
-  private final ExecutionEnvironment executionEnvironment;
-
-  public State(Container container, ExecutionEnvironment executionEnvironment) {
+  public State(Container container) {
     this.container = container;
-    this.executionEnvironment = executionEnvironment;
   }
 
   public Container getContainer() {
     return container;
   }
 
-  public ExecutionEnvironment currentEnvironment() {
-    return executionEnvironment;
-  }
-
   public static State create(EnsoContext context) {
-    return new State(Container.create(context), context.getExecutionEnvironment());
-  }
-
-  public State withContextEnabledIn(Atom context, String environmentName) {
-    if (executionEnvironment.getName().equals(environmentName)) {
-      return new State(container, executionEnvironment.withContextEnabled(context));
-    } else {
-      return this;
-    }
-  }
-
-  public State withContextDisabledIn(Atom context, String environmentName) {
-    if (executionEnvironment.getName().equals(environmentName)) {
-      return new State(container, executionEnvironment.withContextDisabled(context));
-    } else {
-      return this;
-    }
+    return new State(Container.create(context));
   }
 
   public static class Container extends DynamicObject {