fix: alias context to latest persistent input

Author: kalectory

rlm/environments/base_env.py

@@ -298,7 +298,7 @@ class SupportsPersistence(Protocol):
 
     ALIASING BEHAVIOR:
         The unversioned names always point to index 0:
-        - context  -> context_0 (first context)
+        - context  -> latest numbered context
         - history  -> history_0 (first history)
 
     EXAMPLE IMPLEMENTATION:
@@ -335,7 +335,7 @@ def add_context(
         Storage:
             Must store so executed code can access:
             - context_0, context_1, etc. (versioned)
-            - context (alias to context_0)
+            - context (alias to latest numbered context)
 
         Args:
             context_payload: The context data (string, dict, or list).

rlm/environments/docker_repl.py

@@ -387,7 +387,10 @@ def SHOW_VARS():
 
 # Restore scaffold aliases if overwritten by executed code
 if "context_0" in _locals:
-    _locals["context"] = _locals["context_0"]
+    _locals["context"] = _locals[max(
+        (key for key in _locals if key.startswith("context_") and key[8:].isdigit()),
+        key=lambda key: int(key[8:]),
+    )]
 {history_alias}
 
 save_state(_locals)
@@ -590,7 +593,7 @@ def add_context(
     ) -> int:
         """Add a context as ``context_N`` (auto-incrementing unless given).
 
-        ``context`` aliases ``context_0`` for backward compatibility, matching
+        ``context`` aliases the latest numbered context, matching
         the versioning behavior of the local environment.
         """
         if context_index is None:
@@ -611,8 +614,7 @@ def add_context(
                 f"with open('/workspace/{fname}', 'r') as _f:\n"
                 f"    {var_name} = json.load(_f)"
             )
-        if context_index == 0:
-            code += "\ncontext = context_0"
+        code += f"\ncontext = {var_name}"
         self.execute_code(code)
 
         self._context_count = max(self._context_count, context_index + 1)

rlm/environments/ipython_repl.py

@@ -1009,10 +1009,9 @@ def add_context(
                     f"    {var_name} = _rlm_json.load(_rlm_f)"
                 )
 
-            # Fold the ``context = context_0`` alias into the same cell to
+            # Fold the latest context alias into the same cell to
             # save a kernel round-trip in subprocess mode.
-            if context_index == 0:
-                code += f"\ncontext = {var_name}"
+            code += f"\ncontext = {var_name}"
 
             result = self.execute_code(code)
             if result.stderr:
@@ -1021,8 +1020,7 @@ def add_context(
             # Shadow for subprocess mode so self.locals can report context_N.
             if self.kernel_mode == "subprocess":
                 self._subprocess_shadow[var_name] = copy.deepcopy(context_payload)
-                if context_index == 0:
-                    self._subprocess_shadow["context"] = self._subprocess_shadow[var_name]
+                self._subprocess_shadow["context"] = self._subprocess_shadow[var_name]
 
             self._context_count = max(self._context_count, context_index + 1)
             return context_index
@@ -1400,7 +1398,12 @@ def _restore_scaffold_in_process(self) -> None:
                     self._last_final_answer = str(current.get("content", ""))
             ns["answer"] = replacement
         if "context_0" in ns:
-            ns["context"] = ns["context_0"]
+            ns["context"] = ns[
+                max(
+                    (key for key in ns if key.startswith("context_") and key[8:].isdigit()),
+                    key=lambda key: int(key[8:]),
+                )
+            ]
         if "history_0" in ns:
             ns["history"] = ns["history_0"]
         # Re-inject custom tools if overwritten

rlm/environments/local_repl.py

@@ -431,9 +431,8 @@ def add_context(
                 f"import json\nwith open(r'{context_path}', 'r') as f:\n    {var_name} = json.load(f)"
             )
 
-        # Alias context_0 as 'context' for backward compatibility
-        if context_index == 0:
-            self.execute_code(f"context = {var_name}")
+        # Alias the latest context as 'context'.
+        self.execute_code(f"context = {var_name}")
 
         self._context_count = max(self._context_count, context_index + 1)
         return context_index
@@ -538,7 +537,16 @@ def _restore_scaffold(self) -> None:
                             self._last_final_answer = str(current.get("content", ""))
                     self.locals["answer"] = replacement
             elif name == "context" and "context_0" in self.locals:
-                self.locals["context"] = self.locals["context_0"]
+                self.locals["context"] = self.locals[
+                    max(
+                        (
+                            key
+                            for key in self.locals
+                            if key.startswith("context_") and key[8:].isdigit()
+                        ),
+                        key=lambda key: int(key[8:]),
+                    )
+                ]
             elif name == "history" and "history_0" in self.locals and not self.compaction:
                 self.locals["history"] = self.locals["history_0"]
             elif name == "history" and self.compaction: