refactor: share latest context lookup across REPL backends

Author: kalectory

rlm/environments/base_env.py

@@ -24,6 +24,18 @@
 )
 
 
+def get_latest_context(namespace: dict) -> object:
+    """Return the highest-numbered context stored in a REPL namespace."""
+    prefix = "context_"
+    names = (
+        name
+        for name in namespace
+        if name.startswith(prefix) and name.removeprefix(prefix).isdigit()
+    )
+    latest = max(names, key=lambda name: int(name.removeprefix(prefix)))
+    return namespace[latest]
+
+
 @dataclass
 class ToolInfo:
     """Parsed information about a custom tool."""

rlm/environments/docker_repl.py

@@ -19,6 +19,7 @@
 
 import base64
 import copy
+import inspect
 import json
 import os
 import shutil
@@ -37,6 +38,7 @@
 from rlm.environments.base_env import (
     NonIsolatedEnv,
     extract_tool_value,
+    get_latest_context,
     validate_custom_tools,
 )
 
@@ -281,6 +283,7 @@ def _build_exec_script(
 except ImportError:
     import pickle as dill
 
+{inspect.getsource(get_latest_context)}
 PROXY = "http://host.docker.internal:{proxy_port}"
 STATE = "/workspace/state.dill"
 DEPTH = {depth}
@@ -387,10 +390,7 @@ def SHOW_VARS():
 
 # Restore scaffold aliases if overwritten by executed code
 if "context_0" in _locals:
-    _locals["context"] = _locals[max(
-        (key for key in _locals if key.startswith("context_") and key[8:].isdigit()),
-        key=lambda key: int(key[8:]),
-    )]
+    _locals["context"] = get_latest_context(_locals)
 {history_alias}
 
 save_state(_locals)

rlm/environments/ipython_repl.py

@@ -50,6 +50,7 @@
     RESERVED_TOOL_NAMES,
     NonIsolatedEnv,
     extract_tool_value,
+    get_latest_context,
     validate_custom_tools,
 )
 from rlm.environments.local_repl import _AnswerDict
@@ -1398,12 +1399,7 @@ 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[
-                max(
-                    (key for key in ns if key.startswith("context_") and key[8:].isdigit()),
-                    key=lambda key: int(key[8:]),
-                )
-            ]
+            ns["context"] = get_latest_context(ns)
         if "history_0" in ns:
             ns["history"] = ns["history_0"]
         # Re-inject custom tools if overwritten

rlm/environments/local_repl.py

@@ -19,6 +19,7 @@
     RESERVED_TOOL_NAMES,
     NonIsolatedEnv,
     extract_tool_value,
+    get_latest_context,
     validate_custom_tools,
 )
 
@@ -537,16 +538,7 @@ 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[
-                    max(
-                        (
-                            key
-                            for key in self.locals
-                            if key.startswith("context_") and key[8:].isdigit()
-                        ),
-                        key=lambda key: int(key[8:]),
-                    )
-                ]
+                self.locals["context"] = get_latest_context(self.locals)
             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: