ensure tool calls and reserved names aren't overwritten

Author: alexzhang13

rlm/environments/base_env.py

@@ -8,15 +8,16 @@
 # Custom Tools Support
 # =============================================================================
 
-# Reserved names that cannot be overridden by custom tools.
-# These are the built-in REPL functions that all environments must provide.
+# Reserved names: cannot be overridden by custom tools, and are restored after each
+# code execution to prevent namespace corruption (e.g. context = "...", llm_query = ...).
 RESERVED_TOOL_NAMES: frozenset[str] = frozenset(
     {
         "llm_query",
         "llm_query_batched",
         "FINAL_VAR",
         "SHOW_VARS",
         "context",
+        "history",
     }
 )
 
@@ -173,7 +174,7 @@ class SupportsCustomTools(Protocol):
         The following names cannot be used as custom tool names:
         - llm_query, llm_query_batched: Single LM completion functions (no tool access)
         - FINAL_VAR, SHOW_VARS: Built-in helper functions
-        - context: The input context variable
+        - context, history: The input context and conversation history variables
 
     EXAMPLE:
         custom_tools = {

rlm/environments/daytona_repl.py

@@ -349,6 +349,12 @@ def SHOW_VARS():
     sys.stdout = old_stdout
     sys.stderr = old_stderr
 
+# Restore scaffold aliases if overwritten by executed code
+if "context_0" in _locals:
+    _locals["context"] = _locals["context_0"]
+if "history_0" in _locals:
+    _locals["history"] = _locals["history_0"]
+
 save_state(_locals)
 
 result = {{

rlm/environments/docker_repl.py

@@ -173,6 +173,12 @@ def SHOW_VARS():
 finally:
     sys.stdout, sys.stderr = old_stdout, old_stderr
 
+# Restore scaffold aliases if overwritten by executed code
+if "context_0" in _locals:
+    _locals["context"] = _locals["context_0"]
+if "history_0" in _locals:
+    _locals["history"] = _locals["history_0"]
+
 save_state(_locals)
 print(json.dumps({{"stdout": stdout_buf.getvalue(), "stderr": stderr_buf.getvalue(), "locals": {{k: repr(v) for k, v in _locals.items() if not k.startswith("_")}}}}, ensure_ascii=False))
 '''

rlm/environments/e2b_repl.py

@@ -241,6 +241,12 @@ def FINAL_VAR(variable_name):
     sys.stdout = old_stdout
     sys.stderr = old_stderr
 
+# Restore scaffold aliases if overwritten by executed code
+if "context_0" in _locals:
+    _locals["context"] = _locals["context_0"]
+if "history_0" in _locals:
+    _locals["history"] = _locals["history_0"]
+
 save_state(_locals)
 
 result = {{

rlm/environments/local_repl.py

@@ -13,7 +13,12 @@
 
 from rlm.core.comms_utils import LMRequest, send_lm_request, send_lm_request_batched
 from rlm.core.types import REPLResult, RLMChatCompletion
-from rlm.environments.base_env import NonIsolatedEnv, extract_tool_value, validate_custom_tools
+from rlm.environments.base_env import (
+    RESERVED_TOOL_NAMES,
+    NonIsolatedEnv,
+    extract_tool_value,
+    validate_custom_tools,
+)
 
 # =============================================================================
 # Safe Builtins
@@ -375,6 +380,22 @@ def _temp_cwd(self):
         finally:
             os.chdir(old_cwd)
 
+    def _restore_scaffold(self) -> None:
+        """Restore scaffold names after execution so overwrites (e.g. context = 'x') don't persist."""
+        for name in RESERVED_TOOL_NAMES:
+            if name == "llm_query":
+                self.globals["llm_query"] = self._llm_query
+            elif name == "llm_query_batched":
+                self.globals["llm_query_batched"] = self._llm_query_batched
+            elif name == "FINAL_VAR":
+                self.globals["FINAL_VAR"] = self._final_var
+            elif name == "SHOW_VARS":
+                self.globals["SHOW_VARS"] = self._show_vars
+            elif name == "context" and "context_0" in self.locals:
+                self.locals["context"] = self.locals["context_0"]
+            elif name == "history" and "history_0" in self.locals:
+                self.locals["history"] = self.locals["history_0"]
+
     def execute_code(self, code: str) -> REPLResult:
         """Execute code in the persistent namespace and return result."""
         start_time = time.perf_counter()
@@ -392,6 +413,9 @@ def execute_code(self, code: str) -> REPLResult:
                     if key not in self.globals and not key.startswith("_"):
                         self.locals[key] = value
 
+                # Restore scaffold so model overwrites (context = ..., llm_query = ...) don't persist
+                self._restore_scaffold()
+
                 stdout = stdout_buf.getvalue()
                 stderr = stderr_buf.getvalue()
             except Exception as e:

rlm/environments/modal_repl.py

@@ -261,6 +261,12 @@ def SHOW_VARS():
     sys.stdout = old_stdout
     sys.stderr = old_stderr
 
+# Restore scaffold aliases if overwritten by executed code
+if "context_0" in _locals:
+    _locals["context"] = _locals["context_0"]
+if "history_0" in _locals:
+    _locals["history"] = _locals["history_0"]
+
 save_state(_locals)
 
 result = {{

rlm/environments/prime_repl.py

@@ -260,6 +260,12 @@ def SHOW_VARS():
     sys.stdout = old_stdout
     sys.stderr = old_stderr
 
+# Restore scaffold aliases if overwritten by executed code
+if "context_0" in _locals:
+    _locals["context"] = _locals["context_0"]
+if "history_0" in _locals:
+    _locals["history"] = _locals["history_0"]
+
 save_state(_locals)
 
 result = {{

tests/test_local_repl.py

@@ -175,6 +175,41 @@ def test_list_context(self):
         repl.cleanup()
 
 
+class TestLocalREPLScaffoldRestoration:
+    """Tests that overwriting scaffold names (context, llm_query, etc.) is reverted after each execution."""
+
+    def test_context_restored_after_overwrite(self):
+        """If the model does context = 'something', the next execution still sees the real context."""
+        repl = LocalREPL(context_payload="original context content")
+        assert repl.locals["context"] == "original context content"
+
+        repl.execute_code('context = "hijacked"')
+        assert repl.locals["context"] == "original context content"
+
+        out = repl.execute_code("print(context)")
+        assert "original context content" in out.stdout
+        repl.cleanup()
+
+    def test_llm_query_restored_after_overwrite(self):
+        """If the model does llm_query = lambda x: 'hijacked', the next execution still has real llm_query."""
+        repl = LocalREPL()
+        repl.execute_code("llm_query = lambda x: 'hijacked'")
+
+        repl.execute_code("r = llm_query('test')")
+        assert "Error" in repl.locals["r"]
+        repl.cleanup()
+
+    def test_final_var_restored_after_overwrite(self):
+        """If the model overwrites FINAL_VAR, the next execution still has the real FINAL_VAR."""
+        repl = LocalREPL()
+        repl.execute_code("answer = 42")
+        repl.execute_code("FINAL_VAR = lambda x: 'overwritten'")
+
+        repl.execute_code("result = FINAL_VAR('answer')")
+        assert repl.locals["result"] == "42"
+        repl.cleanup()
+
+
 class TestLocalREPLCleanup:
     """Tests for cleanup behavior."""