|
16 | 16 | from rlm.environments.base_env import NonIsolatedEnv |
17 | 17 |
|
18 | 18 | # ============================================================================= |
19 | | -# Safe Builtins |
| 19 | +# Safe Builtins and Protected Names |
20 | 20 | # ============================================================================= |
21 | 21 |
|
| 22 | +# Protected scaffold names that cannot be overwritten by user code |
| 23 | +SCAFFOLD_NAMES = frozenset( |
| 24 | + { |
| 25 | + "context", |
| 26 | + "llm_query", |
| 27 | + "llm_query_batched", |
| 28 | + "FINAL_VAR", |
| 29 | + "SHOW_VARS", |
| 30 | + "finish", |
| 31 | + "finish_var", |
| 32 | + } |
| 33 | +) |
| 34 | + |
22 | 35 | # Safe builtins - blocks dangerous operations like eval/exec/input |
23 | 36 | _SAFE_BUILTINS = { |
24 | 37 | # Core types and functions |
@@ -159,11 +172,26 @@ def setup(self): |
159 | 172 | # Track LLM calls made during code execution |
160 | 173 | self._pending_llm_calls: list[RLMChatCompletion] = [] |
161 | 174 |
|
| 175 | + # Track finish state for programmatic completion (§3.1 fix) |
| 176 | + self._finish_answer: str | None = None |
| 177 | + |
162 | 178 | # Add helper functions |
163 | 179 | self.globals["FINAL_VAR"] = self._final_var |
164 | 180 | self.globals["SHOW_VARS"] = self._show_vars |
165 | 181 | self.globals["llm_query"] = self._llm_query |
166 | 182 | self.globals["llm_query_batched"] = self._llm_query_batched |
| 183 | + self.globals["finish"] = self._finish |
| 184 | + self.globals["finish_var"] = self._finish_var |
| 185 | + |
| 186 | + # Store original scaffold bindings for protection (§3.2 fix) |
| 187 | + self._scaffold_bindings = { |
| 188 | + "FINAL_VAR": self._final_var, |
| 189 | + "SHOW_VARS": self._show_vars, |
| 190 | + "llm_query": self._llm_query, |
| 191 | + "llm_query_batched": self._llm_query_batched, |
| 192 | + "finish": self._finish, |
| 193 | + "finish_var": self._finish_var, |
| 194 | + } |
167 | 195 |
|
168 | 196 | def _final_var(self, variable_name: str) -> str: |
169 | 197 | """Return the value of a variable as a final answer.""" |
@@ -192,6 +220,32 @@ def _show_vars(self) -> str: |
192 | 220 | return "No variables created yet. Use ```repl``` blocks to create variables." |
193 | 221 | return f"Available variables: {available}" |
194 | 222 |
|
| 223 | + def _finish(self, value: Any) -> str: |
| 224 | + """Signal completion with a final answer (programmatic FINAL).""" |
| 225 | + self._finish_answer = str(value) |
| 226 | + return f"Task completed with answer: {self._finish_answer}" |
| 227 | + |
| 228 | + def _finish_var(self, variable_name: str) -> str: |
| 229 | + """Signal completion with a variable's value (programmatic FINAL_VAR).""" |
| 230 | + variable_name = variable_name.strip().strip("\"'") |
| 231 | + if variable_name in self.locals: |
| 232 | + self._finish_answer = str(self.locals[variable_name]) |
| 233 | + return f"Task completed with answer from variable '{variable_name}': {self._finish_answer}" |
| 234 | + |
| 235 | + # Provide helpful error message with available variables |
| 236 | + available = [k for k in self.locals.keys() if not k.startswith("_")] |
| 237 | + if available: |
| 238 | + return ( |
| 239 | + f"Error: Variable '{variable_name}' not found. " |
| 240 | + f"Available variables: {available}. " |
| 241 | + f"You must create and assign a variable BEFORE calling finish_var on it." |
| 242 | + ) |
| 243 | + return ( |
| 244 | + f"Error: Variable '{variable_name}' not found. " |
| 245 | + f"No variables have been created yet. " |
| 246 | + f"You must create and assign a variable in a REPL block BEFORE calling finish_var on it." |
| 247 | + ) |
| 248 | + |
195 | 249 | def _llm_query(self, prompt: str, model: str | None = None) -> str: |
196 | 250 | """Query the LM via socket connection to the handler. |
197 | 251 |
|
@@ -287,6 +341,9 @@ def add_context( |
287 | 341 | # Alias context_0 as 'context' for backward compatibility |
288 | 342 | if context_index == 0: |
289 | 343 | self.execute_code(f"context = {var_name}") |
| 344 | + # Protect 'context' in scaffold bindings |
| 345 | + if hasattr(self, "_scaffold_bindings"): |
| 346 | + self._scaffold_bindings["context"] = self.locals.get("context") |
290 | 347 |
|
291 | 348 | self._context_count = max(self._context_count, context_index + 1) |
292 | 349 | return context_index |
@@ -370,6 +427,22 @@ def execute_code(self, code: str) -> REPLResult: |
370 | 427 | if key not in self.globals and not key.startswith("_"): |
371 | 428 | self.locals[key] = value |
372 | 429 |
|
| 430 | + # Restore protected scaffold bindings if overwritten (§3.2 fix) |
| 431 | + if hasattr(self, "_scaffold_bindings"): |
| 432 | + for name in SCAFFOLD_NAMES: |
| 433 | + if name in self._scaffold_bindings: |
| 434 | + # Check if binding was overwritten in globals |
| 435 | + if name in self.globals and self.globals[name] != self._scaffold_bindings[name]: |
| 436 | + self.globals[name] = self._scaffold_bindings[name] |
| 437 | + # Check if binding was overwritten in locals |
| 438 | + if name in self.locals and self.locals[name] != self._scaffold_bindings[name]: |
| 439 | + # For context, update the scaffold binding to track current value |
| 440 | + if name == "context": |
| 441 | + self._scaffold_bindings[name] = self.locals[name] |
| 442 | + else: |
| 443 | + # Restore other scaffold bindings |
| 444 | + self.locals[name] = self._scaffold_bindings[name] |
| 445 | + |
373 | 446 | stdout = stdout_buf.getvalue() |
374 | 447 | stderr = stderr_buf.getvalue() |
375 | 448 | except Exception as e: |
|
0 commit comments