Ensure tool calls and reserved names aren't overwritten + add logger metadata

README.md

@@ -135,17 +135,20 @@ If you use this code or repository in your research, please cite:
 }
 ```
 
+## Optional: Trajectory metadata and logging
+`RLMChatCompletion` has an optional `metadata` field (default empty) that can hold the full trajectory (run config + all iterations and sub-calls) so you can reconstruct the run. Pass an `RLMLogger` to capture it:
+
+- **In-memory only** (trajectory on `completion.metadata`): `logger=RLMLogger()` (no `log_dir`).
+- **Also save to disk** (JSONL for the visualizer): `logger=RLMLogger(log_dir="./logs")`.
+
 ## Optional Debugging: Visualizing RLM Trajectories
-We additionally provide a simple visualizer tool to examine and view the code, sub-LM, and root-LM calls of an RLM trajectory. To save log files (`.jsonl`) on every completion call that can be viewed in the visualizer, initialize the `RLMLogger` object and pass it into the `RLM` on initialization:
+We provide a simple visualizer to inspect code, sub-LM, and root-LM calls. Use `RLMLogger(log_dir="./logs")` so each completion writes a `.jsonl` file:
 ```python
 from rlm.logger import RLMLogger
 from rlm import RLM
 
 logger = RLMLogger(log_dir="./logs")
-rlm = RLM(
-    ...
-    logger=logger
-)
+rlm = RLM(..., logger=logger)
 ```
 
 To run the visualizer locally, we use Node.js and shadcn/ui:

examples/logger_example.py

@@ -0,0 +1,46 @@
+"""
+Logger example: same structure as quickstart, with a different prompt.
+Shows trajectory capture in completion.metadata (and optional save to disk).
+"""
+
+import os
+
+from dotenv import load_dotenv
+
+from rlm import RLM
+from rlm.logger import RLMLogger
+
+load_dotenv()
+
+# In-memory only: trajectory on completion.metadata (no disk write)
+logger = RLMLogger()
+# To also save JSONL for the visualizer: logger = RLMLogger(log_dir="./logs")
+
+rlm = RLM(
+    backend="portkey",
+    backend_kwargs={
+        "model_name": "@openai/gpt-5-nano",
+        "api_key": os.getenv("PORTKEY_API_KEY"),
+    },
+    environment="local",
+    environment_kwargs={},
+    max_depth=1,
+    logger=logger,
+    verbose=True,
+)
+
+prompt = "Compute the sum 1 + 2 + 3 + 4 + 5 and print the result using a REPL block, then return it with FINAL_VAR."
+result = rlm.completion(prompt)
+
+print("Response:", result.response)
+if result.metadata:
+    traj = result.metadata
+    print("Trajectory: run_metadata +", len(traj.get("iterations", [])), "iterations")
+    print("Trajectory:", traj)
+    if traj.get("iterations"):
+        first = traj["iterations"][0]
+        print("  First iteration keys:", list(first.keys()))
+else:
+    print("Trajectory: (none — no logger or metadata not captured)")
+
+print("Full response:", result)

rlm/core/rlm.py

@@ -229,6 +229,9 @@ def completion(
         if self.depth >= self.max_depth:
             return self._fallback_answer(prompt)
 
+        if self.logger:
+            self.logger.clear_iterations()
+
         with self._spawn_completion_context(prompt) as (lm_handler, environment):
             message_history = self._setup_prompt(prompt)
 
@@ -283,6 +286,7 @@ def completion(
                         response=final_answer,
                         usage_summary=usage,
                         execution_time=time_end - time_start,
+                        metadata=self.logger.get_trajectory() if self.logger else None,
                     )
 
                 # Format the iteration for the next prompt.
@@ -310,6 +314,7 @@ def completion(
                 response=final_answer,
                 usage_summary=usage,
                 execution_time=time_end - time_start,
+                metadata=self.logger.get_trajectory() if self.logger else None,
             )
 
     def _completion_turn(

rlm/core/types.py

@@ -96,15 +96,21 @@ class RLMChatCompletion:
     response: str
     usage_summary: UsageSummary
     execution_time: float
+    metadata: dict | None = (
+        None  # Full trajectory (run_metadata + iterations) when logger captures it
+    )
 
     def to_dict(self):
-        return {
+        out = {
             "root_model": self.root_model,
             "prompt": self.prompt,
             "response": self.response,
             "usage_summary": self.usage_summary.to_dict(),
             "execution_time": self.execution_time,
         }
+        if self.metadata is not None:
+            out["metadata"] = self.metadata
+        return out
 
     @classmethod
     def from_dict(cls, data: dict) -> "RLMChatCompletion":
@@ -114,6 +120,7 @@ def from_dict(cls, data: dict) -> "RLMChatCompletion":
             response=data.get("response"),
             usage_summary=UsageSummary.from_dict(data.get("usage_summary")),
             execution_time=data.get("execution_time"),
+            metadata=data.get("metadata"),
         )
 
 

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 = {{

rlm/logger/rlm_logger.py

@@ -1,7 +1,8 @@
 """
 Logger for RLM iterations.
 
-Writes RLMIteration data to JSON-lines files for analysis and debugging.
+Captures run metadata and iterations in memory so they can be attached to
+RLMChatCompletion.metadata. Optionally writes the same data to JSON-lines files.
 """
 
 import json
@@ -13,50 +14,77 @@
 
 
 class RLMLogger:
-    """Logger that writes RLMIteration data to a JSON-lines file."""
+    """
+    Captures trajectory (run metadata + iterations) for each completion.
+    By default only captures in memory; set log_dir to also save to disk.
 
-    def __init__(self, log_dir: str, file_name: str = "rlm"):
-        self.log_dir = log_dir
-        os.makedirs(log_dir, exist_ok=True)
-
-        timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-        run_id = str(uuid.uuid4())[:8]
-        self.log_file_path = os.path.join(log_dir, f"{file_name}_{timestamp}_{run_id}.jsonl")
+    - log_dir=None: trajectory is available via get_trajectory() and can be
+      attached to RLMChatCompletion.metadata (no disk write).
+    - log_dir="path": same capture plus appends to a JSONL file per run.
+    """
 
+    def __init__(self, log_dir: str | None = None, file_name: str = "rlm"):
+        self._save_to_disk = log_dir is not None
+        self.log_dir = log_dir
+        self.log_file_path: str | None = None
+        if self._save_to_disk and log_dir:
+            os.makedirs(log_dir, exist_ok=True)
+            timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+            run_id = str(uuid.uuid4())[:8]
+            self.log_file_path = os.path.join(log_dir, f"{file_name}_{timestamp}_{run_id}.jsonl")
+
+        self._run_metadata: dict | None = None
+        self._iterations: list[dict] = []
         self._iteration_count = 0
         self._metadata_logged = False
 
-    def log_metadata(self, metadata: RLMMetadata):
-        """Log RLM metadata as the first entry in the file."""
+    def log_metadata(self, metadata: RLMMetadata) -> None:
+        """Capture run metadata (and optionally write to file)."""
         if self._metadata_logged:
             return
 
-        entry = {
-            "type": "metadata",
-            "timestamp": datetime.now().isoformat(),
-            **metadata.to_dict(),
-        }
-
-        with open(self.log_file_path, "a") as f:
-            json.dump(entry, f)
-            f.write("\n")
-
+        self._run_metadata = metadata.to_dict()
         self._metadata_logged = True
 
-    def log(self, iteration: RLMIteration):
-        """Log an RLMIteration to the file."""
+        if self._save_to_disk and self.log_file_path:
+            entry = {
+                "type": "metadata",
+                "timestamp": datetime.now().isoformat(),
+                **self._run_metadata,
+            }
+            with open(self.log_file_path, "a") as f:
+                json.dump(entry, f)
+                f.write("\n")
+
+    def log(self, iteration: RLMIteration) -> None:
+        """Capture one iteration (and optionally append to file)."""
         self._iteration_count += 1
-
         entry = {
             "type": "iteration",
             "iteration": self._iteration_count,
             "timestamp": datetime.now().isoformat(),
             **iteration.to_dict(),
         }
+        self._iterations.append(entry)
+
+        if self._save_to_disk and self.log_file_path:
+            with open(self.log_file_path, "a") as f:
+                json.dump(entry, f)
+                f.write("\n")
 
-        with open(self.log_file_path, "a") as f:
-            json.dump(entry, f)
-            f.write("\n")
+    def clear_iterations(self) -> None:
+        """Reset iterations for the next completion (trajectory is per completion)."""
+        self._iterations = []
+        self._iteration_count = 0
+
+    def get_trajectory(self) -> dict | None:
+        """Return captured run_metadata + iterations for the current completion, or None if no metadata yet."""
+        if self._run_metadata is None:
+            return None
+        return {
+            "run_metadata": self._run_metadata,
+            "iterations": list(self._iterations),
+        }
 
     @property
     def iteration_count(self) -> int: