[REPL_ENV] update RLM env to support recursion - #437
Conversation
Greptile SummaryThis PR refactors the Key issues found:
What is fixed in this PR vs. prior threads:
Confidence Score: 1/5
Important Files Changed
Sequence DiagramsequenceDiagram
participant User
participant LocalRLMRunner
participant LocalREPLEnv
participant REPLEnvironment
participant Backend as LocalChildRLMBackend
participant ChildRunner as LocalRLMRunner (child)
User->>LocalRLMRunner: run(context, task_prompt)
LocalRLMRunner->>Backend: backend_factory(llm_chat_fn, ...)
Backend-->>LocalRLMRunner: backend instance (shared BackendLimits)
LocalRLMRunner->>LocalREPLEnv: __init__(llm_query_fn, subcall_fn, ...)
LocalRLMRunner->>REPLEnvironment: reset(context, task_prompt, max_iterations)
REPLEnvironment-->>LocalRLMRunner: REPLObservation (done=False)
loop up to max_iterations
LocalRLMRunner->>LocalRLMRunner: _chat(messages, model)
LocalRLMRunner->>LocalREPLEnv: execute(code)
LocalREPLEnv->>REPLEnvironment: step(REPLAction)
alt code calls rlm_query(prompt)
REPLEnvironment->>Backend: recursive_query(prompt)
Backend->>Backend: check _children_lock / _children_spawned
Backend->>ChildRunner: runner_factory(..., depth=N+1)
ChildRunner->>ChildRunner: run(prompt, prompt, timeout_s=per_child_timeout_s)
ChildRunner-->>Backend: RLMRunResult.final_answer
Backend-->>REPLEnvironment: child answer string
end
REPLEnvironment-->>LocalREPLEnv: REPLObservation
LocalREPLEnv-->>LocalRLMRunner: StepResult
alt done == True
LocalRLMRunner-->>User: RLMRunResult
end
end
LocalRLMRunner->>LocalRLMRunner: _default_answer(messages) [fallback]
LocalRLMRunner-->>User: RLMRunResult
|
| on_subcall_complete=self.on_subcall_complete, | ||
| ) | ||
| if self.limits.per_child_timeout_s is None: | ||
| result = child.run(prompt, prompt, model=model) |
There was a problem hiding this comment.
Same string used as both context and task_prompt
child.run(prompt, prompt, model=model) passes the user's prompt as both the context variable and the task_prompt. This means when a child RLM runner is spawned, it loads the question text itself as the context object inside the REPL. While functional (the child can still answer by reasoning on it), it is likely to confuse models: the system prompt tells them to "check the content of the context variable" which contains their task, not a separate document.
Consider whether a child should receive an empty context ("") and the prompt as only the task, or whether this is intentional for self-referential tasks:
| result = child.run(prompt, prompt, model=model) | |
| result = child.run("", prompt, model=model) |
Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/recursive_backends.py
Line: 179
Comment:
**Same string used as both `context` and `task_prompt`**
`child.run(prompt, prompt, model=model)` passes the user's `prompt` as both the `context` variable and the `task_prompt`. This means when a child RLM runner is spawned, it loads the question text itself as the `context` object inside the REPL. While functional (the child can still answer by reasoning on it), it is likely to confuse models: the system prompt tells them to "check the content of the `context` variable" which contains their task, not a separate document.
Consider whether a child should receive an empty context (`""`) and the prompt as only the task, or whether this is intentional for self-referential tasks:
```suggestion
result = child.run("", prompt, model=model)
```
How can I resolve this? If you propose a fix, please make it concise.Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com>
| self.rlm_max_depth = runtime_rlm_max_depth | ||
| self.rlm_max_iterations = runtime_rlm_max_iterations | ||
|
|
||
| # Create LLM functions if not already provided at init. | ||
| # Match example/client behavior by allowing anonymous routed inference too. | ||
| if not self.llm_query_fn: | ||
| effective_token = hf_token or os.environ.get("HF_TOKEN") | ||
| if effective_token: | ||
| self._create_llm_functions(effective_token, llm_model) | ||
| effective_token = ( | ||
| hf_token if hf_token is not None else os.environ.get("HF_TOKEN") | ||
| ) | ||
| self._create_llm_functions(effective_token, llm_model) |
There was a problem hiding this comment.
Recursive controller not rebuilt when rlm_max_depth changes on re-reset
self.rlm_max_depth is updated on lines 239–240 when a caller passes new rlm_max_depth/rlm_max_iterations kwargs. However, the guard on line 244 — if not self.llm_query_fn — is already True after the first successful reset (because self.llm_query_fn was set by _create_llm_functions). This means any updated depth/iteration values are written to the instance attributes but the underlying RecursiveController (and its LocalChildRLMBackend or DirectLMBackend) is never re-created to reflect them.
Concrete failure scenario:
- Client resets with no kwargs → controller is built with
rlm_max_depth=1(default),rlm_query_fn=None. - Client resets again with
rlm_max_depth=2→self.rlm_max_depthbecomes2, but the old depth-1 controller is still active, sorlm_queryremains absent from the REPL namespace.
The fix is to also trigger _create_llm_functions when the runtime depth or iteration config changes:
depth_changed = (
runtime_rlm_max_depth != self.rlm_max_depth
or runtime_rlm_max_iterations != self.rlm_max_iterations
)
if not self.llm_query_fn or depth_changed:
...
self._create_llm_functions(effective_token, llm_model)(Store the depth/iterations before updating self so the comparison is meaningful.)
Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/server/repl_environment.py
Line: 239-248
Comment:
**Recursive controller not rebuilt when `rlm_max_depth` changes on re-reset**
`self.rlm_max_depth` is updated on lines 239–240 when a caller passes new `rlm_max_depth`/`rlm_max_iterations` kwargs. However, the guard on line 244 — `if not self.llm_query_fn` — is already `True` after the *first* successful reset (because `self.llm_query_fn` was set by `_create_llm_functions`). This means any updated depth/iteration values are written to the instance attributes but the underlying `RecursiveController` (and its `LocalChildRLMBackend` or `DirectLMBackend`) is **never re-created** to reflect them.
Concrete failure scenario:
1. Client resets with no kwargs → controller is built with `rlm_max_depth=1` (default), `rlm_query_fn=None`.
2. Client resets again with `rlm_max_depth=2` → `self.rlm_max_depth` becomes `2`, but the old depth-1 controller is still active, so `rlm_query` remains absent from the REPL namespace.
The fix is to also trigger `_create_llm_functions` when the runtime depth or iteration config changes:
```python
depth_changed = (
runtime_rlm_max_depth != self.rlm_max_depth
or runtime_rlm_max_iterations != self.rlm_max_iterations
)
if not self.llm_query_fn or depth_changed:
...
self._create_llm_functions(effective_token, llm_model)
```
(Store the depth/iterations *before* updating `self` so the comparison is meaningful.)
How can I resolve this? If you propose a fix, please make it concise.| """Create LLM/subcall functions dynamically using client-provided token.""" | ||
| try: | ||
| chat_fn = self._build_hf_chat_fn(hf_token, llm_model) | ||
| except RuntimeError: | ||
| return |
There was a problem hiding this comment.
Silent failure when huggingface_hub is not installed
_create_llm_functions catches RuntimeError from _build_hf_chat_fn (raised when huggingface_hub is missing) and returns silently. A caller that passed hf_token expecting LLM support gets none, with no log line or warning. The LLM functions won't be injected into the REPL namespace, and the model will see llm_query/rlm_query as undefined — an inscrutable failure.
At minimum, emit a warning:
except RuntimeError as exc:
import warnings
warnings.warn(
f"[repl_env] Could not initialize LLM functions: {exc}. "
"llm_query/rlm_query will not be available in the REPL.",
RuntimeWarning,
stacklevel=2,
)
returnPrompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/server/repl_environment.py
Line: 176-180
Comment:
**Silent failure when `huggingface_hub` is not installed**
`_create_llm_functions` catches `RuntimeError` from `_build_hf_chat_fn` (raised when `huggingface_hub` is missing) and returns silently. A caller that passed `hf_token` expecting LLM support gets none, with no log line or warning. The LLM functions won't be injected into the REPL namespace, and the model will see `llm_query`/`rlm_query` as undefined — an inscrutable failure.
At minimum, emit a warning:
```python
except RuntimeError as exc:
import warnings
warnings.warn(
f"[repl_env] Could not initialize LLM functions: {exc}. "
"llm_query/rlm_query will not be available in the REPL.",
RuntimeWarning,
stacklevel=2,
)
return
```
How can I resolve this? If you propose a fix, please make it concise.| if not self.llm_query_fn: | ||
| effective_token = hf_token or os.environ.get("HF_TOKEN") | ||
| if effective_token: | ||
| self._create_llm_functions(effective_token, llm_model) | ||
| effective_token = ( | ||
| hf_token if hf_token is not None else os.environ.get("HF_TOKEN") | ||
| ) | ||
| self._create_llm_functions(effective_token, llm_model) |
There was a problem hiding this comment.
LLM client created even when no auth token is available
_create_llm_functions is called whenever self.llm_query_fn is unset, including when effective_token is None (no token was supplied by the client or server environment). _build_hf_chat_fn accepts None for hf_token despite the str annotation, passes it to InferenceClient, and the resulting client object silently accepts all construction calls.
For gated models (including the server default), every inference call then fails at runtime with an authentication error. No diagnostic appears at reset() time — the failure surfaces mid-episode as an opaque error string returned from the injected REPL function.
The guard on this line should also require that effective_token is non-None before constructing the client, so that the fast-fail happens at initialization rather than during inference. Additionally, the hf_token parameter type annotation in _build_hf_chat_fn should be Optional[str].
Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/server/repl_environment.py
Line: 244-248
Comment:
**LLM client created even when no auth token is available**
`_create_llm_functions` is called whenever `self.llm_query_fn` is unset, including when `effective_token` is `None` (no token was supplied by the client or server environment). `_build_hf_chat_fn` accepts `None` for `hf_token` despite the `str` annotation, passes it to `InferenceClient`, and the resulting client object silently accepts all construction calls.
For gated models (including the server default), every inference call then fails at runtime with an authentication error. No diagnostic appears at `reset()` time — the failure surfaces mid-episode as an opaque error string returned from the injected REPL function.
The guard on this line should also require that `effective_token` is non-`None` before constructing the client, so that the fast-fail happens at initialization rather than during inference. Additionally, the `hf_token` parameter type annotation in `_build_hf_chat_fn` should be `Optional[str]`.
How can I resolve this? If you propose a fix, please make it concise.| self._env.max_iterations = max_iterations | ||
| if rlm_max_depth is not None: | ||
| self._env.rlm_max_depth = rlm_max_depth | ||
| if rlm_max_iterations is not None: | ||
| self._env.rlm_max_iterations = rlm_max_iterations | ||
| reset_kwargs = {} | ||
| if rlm_max_depth is not None: | ||
| reset_kwargs["rlm_max_depth"] = rlm_max_depth | ||
| if rlm_max_iterations is not None: | ||
| reset_kwargs["rlm_max_iterations"] = rlm_max_iterations | ||
| if expected_answer is not None: | ||
| reset_kwargs["expected_answer"] = expected_answer | ||
| obs = self._env.reset( | ||
| seed=seed, | ||
| episode_id=episode_id, |
There was a problem hiding this comment.
depth_changed detection always evaluates false via LocalREPLEnv
LocalREPLEnv.reset() writes self._env.rlm_max_depth = rlm_max_depth on lines 77–78 before calling self._env.reset(...). Inside REPLEnvironment.reset(), the depth-change guard compares runtime_rlm_max_depth (taken from the kwarg) against self.rlm_max_depth (already updated by the pre-assignment). The comparison is therefore new_depth != new_depth, which is always False.
Concrete failure scenario:
LocalREPLEnv()is constructed withoutllm_query_fn. The firstreset(rlm_max_depth=1)builds a depth-1 controller.reset(rlm_max_depth=3)is called next —LocalREPLEnvpre-setsself._env.rlm_max_depth = 3, then callsself._env.reset(…, rlm_max_depth=3). InsideREPLEnvironment.reset(),depth_changedevaluates to3 != 3 → False, so the old depth-1 controller is silently reused.
Fix: remove the redundant pre-assignments (lines 77–80) from LocalREPLEnv.reset() and rely solely on the kwarg-based path in REPLEnvironment.reset(), which was specifically designed to detect and handle the change. The values will still be updated correctly via the depth_changed path in REPLEnvironment.reset().
Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/local.py
Line: 76-90
Comment:
**`depth_changed` detection always evaluates false via `LocalREPLEnv`**
`LocalREPLEnv.reset()` writes `self._env.rlm_max_depth = rlm_max_depth` on lines 77–78 **before** calling `self._env.reset(...)`. Inside `REPLEnvironment.reset()`, the depth-change guard compares `runtime_rlm_max_depth` (taken from the kwarg) against `self.rlm_max_depth` (already updated by the pre-assignment). The comparison is therefore `new_depth != new_depth`, which is always `False`.
Concrete failure scenario:
1. `LocalREPLEnv()` is constructed without `llm_query_fn`. The first `reset(rlm_max_depth=1)` builds a depth-1 controller.
2. `reset(rlm_max_depth=3)` is called next — `LocalREPLEnv` pre-sets `self._env.rlm_max_depth = 3`, then calls `self._env.reset(…, rlm_max_depth=3)`. Inside `REPLEnvironment.reset()`, `depth_changed` evaluates to `3 != 3 → False`, so the old depth-1 controller is silently reused.
**Fix**: remove the redundant pre-assignments (lines 77–80) from `LocalREPLEnv.reset()` and rely solely on the kwarg-based path in `REPLEnvironment.reset()`, which was specifically designed to detect and handle the change. The values will still be updated correctly via the `depth_changed` path in `REPLEnvironment.reset()`.
How can I resolve this? If you propose a fix, please make it concise.| # Token resolution: explicit hf_token > HF_TOKEN env var > cached HF login. | ||
| if not self.llm_query_fn or depth_changed: | ||
| effective_token = ( | ||
| hf_token if hf_token is not None else os.environ.get("HF_TOKEN") | ||
| ) | ||
| self._create_llm_functions(effective_token, llm_model) | ||
|
|
||
| # Initialize state |
There was a problem hiding this comment.
Depth-change rebuild can clobber a valid LLM setup with a no-credential client
The original code only called _create_llm_functions when a valid API credential was available. This PR removed that guard when rebuilding due to depth_changed.
The consequence: if a caller provides credentials on the first reset() (valid controller is built), then calls reset(rlm_max_depth=3) on the next episode without re-providing credentials — depth_changed=True and no credential is resolved, so _create_llm_functions(None, …) is invoked. The old valid controller is replaced by one backed by an InferenceClient constructed without auth. All subsequent llm_query / rlm_query calls then fail at inference time with auth errors, with no diagnostic at reset time.
The fix is to restore the credential guard so _create_llm_functions is only called when a valid credential is present, regardless of whether the trigger is a missing llm_query_fn or the depth_changed flag.
Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/server/repl_environment.py
Line: 249-256
Comment:
**Depth-change rebuild can clobber a valid LLM setup with a no-credential client**
The original code only called `_create_llm_functions` when a valid API credential was available. This PR removed that guard when rebuilding due to `depth_changed`.
The consequence: if a caller provides credentials on the first `reset()` (valid controller is built), then calls `reset(rlm_max_depth=3)` on the next episode without re-providing credentials — `depth_changed=True` and no credential is resolved, so `_create_llm_functions(None, …)` is invoked. The old valid controller is **replaced** by one backed by an `InferenceClient` constructed without auth. All subsequent `llm_query` / `rlm_query` calls then fail at inference time with auth errors, with no diagnostic at reset time.
The fix is to restore the credential guard so `_create_llm_functions` is only called when a valid credential is present, regardless of whether the trigger is a missing `llm_query_fn` or the `depth_changed` flag.
How can I resolve this? If you propose a fix, please make it concise.| message="Final answer submitted.", | ||
| reward=self.reward_on_success, | ||
| ) | ||
| obs.reward = self._apply_rubric(action, obs) |
There was a problem hiding this comment.
_apply_rubric method undefined — every step() call will raise AttributeError
self._apply_rubric(action, obs) is called on lines 439, 452, and 492, but the method is never defined anywhere in REPLEnvironment. At runtime, every call to env.step() — and therefore every call to env.execute() via LocalREPLEnv or LocalRLMRunner — will raise:
AttributeError: 'REPLEnvironment' object has no attribute '_apply_rubric'
The rubric is stored at self.rubric (a REPLRubric instance). The method body should be:
def _apply_rubric(self, action: REPLAction, obs: REPLObservation) -> float:
return self.rubric(action, obs)Why this wasn't caught: the entire test file starts with pytest.importorskip("smolagents"). If smolagents is absent in the CI environment, every test — including test_rubric_reward_on_success, test_rubric_reward_on_error, and the full TestLocalRLMRunner suite — is silently skipped. The pre-submit checks would show "0 failures" while the environment is completely broken for any caller that actually uses step().
Prompt To Fix With AI
This is a comment left during a code review.
Path: envs/repl_env/server/repl_environment.py
Line: 439
Comment:
**`_apply_rubric` method undefined — every `step()` call will raise `AttributeError`**
`self._apply_rubric(action, obs)` is called on lines 439, 452, and 492, but the method is never defined anywhere in `REPLEnvironment`. At runtime, every call to `env.step()` — and therefore every call to `env.execute()` via `LocalREPLEnv` or `LocalRLMRunner` — will raise:
```
AttributeError: 'REPLEnvironment' object has no attribute '_apply_rubric'
```
The rubric is stored at `self.rubric` (a `REPLRubric` instance). The method body should be:
```python
def _apply_rubric(self, action: REPLAction, obs: REPLObservation) -> float:
return self.rubric(action, obs)
```
**Why this wasn't caught:** the entire test file starts with `pytest.importorskip("smolagents")`. If `smolagents` is absent in the CI environment, every test — including `test_rubric_reward_on_success`, `test_rubric_reward_on_error`, and the full `TestLocalRLMRunner` suite — is silently skipped. The pre-submit checks would show "0 failures" while the environment is completely broken for any caller that actually uses `step()`.
How can I resolve this? If you propose a fix, please make it concise.
Summary
Added support for recursion for RLMs
Type of Change
Alignment Checklist
Before submitting, verify:
.claude/docs/PRINCIPLES.mdand this PR aligns with our principles.claude/docs/INVARIANTS.mdand no invariants are violated/pre-submit-pr(orbash .claude/hooks/lint.shand tests) and addressed all issuesRFC Status
Test Plan
Claude Code Review