Upload folder using huggingface_hub

Dockerfile

@@ -10,8 +10,8 @@
 # - Standalone environments (with openenv from pip)
 # The build script (openenv build) handles context detection and sets appropriate build args.
 
-ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
-FROM ghcr.io/meta-pytorch/openenv-base:latest AS builder
+ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest
+FROM ghcr.io/huggingface/openenv-base:latest AS builder
 
 WORKDIR /app
 
@@ -22,8 +22,8 @@ ARG ENV_NAME=repl_env
 # Copy environment code (always at root of build context)
 COPY . /app/env
 
-# For in-repo builds, openenv-core is already in the pyproject.toml dependencies
-# For standalone builds, openenv-core will be installed from pip via pyproject.toml
+# For in-repo builds, openenv is already in the pyproject.toml dependencies
+# For standalone builds, openenv will be installed from pip via pyproject.toml
 WORKDIR /app/env
 
 # Ensure uv is available (for local builds where base image lacks it)
@@ -59,7 +59,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
     fi
 
 # Final runtime stage
-FROM ghcr.io/meta-pytorch/openenv-base:latest
+FROM ghcr.io/huggingface/openenv-base:latest
 
 WORKDIR /app
 

README.md

@@ -8,124 +8,34 @@ pinned: false
 app_port: 8000
 base_path: /web
 tags:
-  - openenv-0.2.3
+  - openenv-0.10.0
   - openenv
 ---
 
-## Hugging Face Space Deployment
-
-This Space is built from OpenEnv environment `repl_env`.
-
-- Space URL: `https://huggingface.co/spaces/openenv/repl`
-- OpenEnv pinned ref: `0.2.3`
-- Hub tag: `openenv`
-
-### Connecting from Code
-
-```python
-from envs.repl_env import Env
-
-env = Env(base_url="https://huggingface.co/spaces/openenv/repl")
-```
-
 # REPL Environment for OpenEnv
 
-`repl_env` is an OpenEnv-native Python REPL environment for Recursive Language Model style execution. It now follows the current OpenEnv client/server conventions:
-
-- `REPLEnv` is the remote async `EnvClient`
-- `.sync()` is the sync wrapper for remote usage
-- `LocalREPLEnv` is the explicit in-process helper
-- `LocalRLMRunner` is the higher-level orchestration loop for local recursive RLM runs
-
-The architecture is intentionally split the same way the official `rlm` and DSPy implementations split things:
-
-- the environment executes code and exposes tools
-- the runner owns the iterative prompting loop
-- recursive behavior lives in backend/controller modules, not in the executor
-
-## Overview
-
-Inside the REPL, the model can:
+`repl_env` is a Python REPL environment for [Recursive Language Model](https://huggingface.co/papers/2512.24601) (RLM) style execution. The model writes code that runs in a persistent namespace and can:
 
 - inspect `context`
-- execute Python code across multiple turns with persistent state
-- call `llm_query(...)` and `llm_query_batched(...)`
-- call `rlm_query(...)` and `rlm_query_batched(...)` for recursive child runs when configured
+- execute Python across multiple turns with persistent state
+- call `llm_query(...)` and `llm_query_batched(...)` to query a language model
+- call `rlm_query(...)` and `rlm_query_batched(...)` for recursive child runs, when configured
 - finish with `FINAL(...)`, `FINAL_VAR(...)`, or `answer = {"content": ..., "ready": True}`
 
-## Current Architecture
-
-Main modules:
-
-- [`client.py`](client.py): remote async OpenEnv client
-- [`local.py`](local.py): explicit in-process local env helper
-- [`runner.py`](runner.py): local RLM orchestration loop
-- [`recursive_backends.py`](recursive_backends.py): direct and recursive backend implementations
-- [`recursive_controller.py`](recursive_controller.py): server-side backend/broker composition
-- [`rubrics.py`](rubrics.py): reward rubrics (OpenEnv RFC 004)
-- [`server/repl_environment.py`](server/repl_environment.py): server-side execution environment
-- [`server/app.py`](server/app.py): OpenEnv HTTP server app and env factory
-
-## What Works Today
-
-- Standard remote OpenEnv usage through `REPLEnv`
-- Local in-process execution through `LocalREPLEnv`
-- Local recursive RLM runs through `LocalRLMRunner`
-- Server-backed recursive calls through the current controller/broker path
-- Explicit recursion controls:
-  - `max_depth`
-  - `max_children_total`
-  - `max_children_per_batch`
-  - `per_child_timeout_s`
-  - `result_truncation_limit`
-- Lightweight child trace metadata on local runner results
-- Rubric-based rewards (OpenEnv RFC 004):
-  - `ExactMatchRubric`: binary outcome reward against ground truth
-  - `FuzzyMatchRubric`: partial credit for containment matches
-  - `CustomMetricRubric`: user-provided `metric(expected, predicted) -> float`
-  - `CodeExecutionRubric`: per-step process reward for code errors
-  - `REPLRubric`: composite rubric combining outcome + process
-  - Ground truth injectable at reset via `expected_answer`
-
-## Rewards
-
-Rewards follow the OpenEnv Rubric system (RFC 004). The environment uses
-`REPLRubric` by default, which combines:
-
-- **Outcome reward** (on terminal steps): compares `final_answer` against
-  `expected_answer` if provided. Returns 1.0 for match, 0.0 otherwise.
-- **Process reward** (on non-terminal steps): returns -0.05 for code
-  execution errors, 0.0 for successful steps.
-- **Failure reward**: returns -0.1 when max iterations exhausted without an answer.
-
-For RL training (GRPO, etc.), pass `expected_answer` at reset time:
-
-```python
-with LocalREPLEnv() as env:
-    env.reset(
-        context="...",
-        task_prompt="...",
-        expected_answer="42",  # ground truth for rubric scoring
-    )
-    result = env.execute("print(FINAL(42))")
-    print(result.reward)  # 1.0 (correct)
-```
+The package provides:
 
-Custom rubrics can be injected at construction:
+- `REPLEnv`: the async client for a remote server (`.sync()` for synchronous code), with `execute(code)`, `submit_final_answer(answer)`, `get_variable(name)` and `list_variables()` on top of `reset()`/`step()`
+- `LocalREPLEnv`: the same environment, in process
+- `LocalRLMRunner`: a local RLM loop that prompts a model, runs its code and handles recursion
 
-```python
-from repl_env import LocalREPLEnv, CustomMetricRubric, REPLRubric
+## Quick Start
 
-def my_metric(expected, predicted):
-    return 1.0 if expected.strip() == predicted.strip() else 0.0
+Start a server:
 
-env = LocalREPLEnv(rubric=REPLRubric(outcome=CustomMetricRubric(my_metric)))
+```bash
+PYTHONPATH=src:envs uvicorn envs.repl_env.server.app:app --host 127.0.0.1 --port 8000
 ```
 
-## Quick Start
-
-### Remote Server Usage
-
 Async:
 
 ```python
@@ -162,7 +72,7 @@ with REPLEnv(base_url="http://127.0.0.1:8000").sync() as env:
     print(result.observation.result.stdout)
 ```
 
-### Local Environment Usage
+### In Process
 
 ```python
 from repl_env import LocalREPLEnv
@@ -177,7 +87,56 @@ with LocalREPLEnv() as env:
     print(env.state().final_answer)
 ```
 
-### Local Recursive RLM Usage
+## Server Configuration
+
+Environment variables read by [`server/app.py`](server/app.py):
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `HF_TOKEN` | unset | Enables `llm_query` on the server. Without it, a client can pass `hf_token` to `reset()` |
+| `LLM_MODEL` | `Qwen/Qwen3.5-9B` | Default model for `llm_query`. A client can pass `llm_model` to `reset()` |
+| `REPL_MAX_ITERATIONS` | `30` | Maximum steps per episode |
+| `REPL_MAX_OUTPUT_LENGTH` | `8192` | Maximum captured output per step |
+| `REPL_CONTEXT_PREVIEW_LENGTH` | `500` | Length of `context_preview` in observations |
+| `REPL_RLM_MAX_DEPTH` | `2` | Maximum recursion depth for `rlm_query` |
+| `REPL_RLM_MAX_ITERATIONS` | `30` | Maximum iterations for recursive child runs |
+| `MAX_CONCURRENT_ENVS` | `8` | Maximum concurrent sessions |
+
+## Reward
+
+Rewards use the OpenEnv [rubric system](https://huggingface.co/docs/openenv/guides/rewards). The default `REPLRubric` combines:
+
+- **Outcome reward** (on terminal steps): compares `final_answer` against
+  `expected_answer` if provided. Returns 1.0 for match, 0.0 otherwise.
+- **Process reward** (on non-terminal steps): returns -0.05 for code
+  execution errors, 0.0 for successful steps.
+- **Failure reward**: returns -0.1 when max iterations exhausted without an answer.
+
+For RL training (GRPO, etc.), pass `expected_answer` at reset time:
+
+```python
+with LocalREPLEnv() as env:
+    env.reset(
+        context="...",
+        task_prompt="...",
+        expected_answer="42",  # ground truth for rubric scoring
+    )
+    result = env.execute("print(FINAL(42))")
+    print(result.reward)  # 1.0 (correct)
+```
+
+Other rubrics: `ExactMatchRubric` (binary match), `FuzzyMatchRubric` (1.0 for an exact match, 0.5 when the expected answer is contained in the final answer), `CustomMetricRubric` (your `metric(expected, predicted) -> float`) and `CodeExecutionRubric` (per-step error penalty). Pass one at construction:
+
+```python
+from repl_env import LocalREPLEnv, CustomMetricRubric, REPLRubric
+
+def my_metric(expected, predicted):
+    return 1.0 if expected.strip() == predicted.strip() else 0.0
+
+env = LocalREPLEnv(rubric=REPLRubric(outcome=CustomMetricRubric(my_metric)))
+```
+
+## Running an RLM Locally
 
 `LocalRLMRunner` takes any `chat_fn(messages, model=None) -> str`. It works
 with HF Inference API, vLLM, SGLang, Ollama, or any OpenAI-compatible server.
@@ -226,7 +185,7 @@ runner = LocalRLMRunner(chat_fn, max_iterations=30, max_depth=2)
 result = runner.run(context, task)
 ```
 
-### Using Different Models for Outer and Inner Loops
+### Different Models for Outer and Inner Loops
 
 The outer loop (code generation) can use a large model while inner
 `llm_query`/`rlm_query` calls use a smaller, faster model. Pass a
@@ -277,45 +236,9 @@ runner = LocalRLMRunner(
 result = runner.run(context, task)
 ```
 
-## Server
-
-Run the local server:
-
-```bash
-PYTHONPATH=src:envs uvicorn envs.repl_env.server.app:app --host 127.0.0.1 --port 8000
-```
-
-The server uses a proper OpenEnv environment factory in [`server/app.py`](server/app.py).
+`LocalRLMRunner` also takes recursion limits (`max_children_total`, `max_children_per_batch`, `per_child_timeout_s`, `result_truncation_limit`) and lifecycle callbacks (`on_subcall_start(depth, model, prompt_preview)`, `on_subcall_complete(depth, model, duration, error_or_none)`). Its results carry lightweight child trace metadata.
 
-## API Surface
-
-### Remote Client
-
-```python
-class REPLEnv(EnvClient[REPLAction, REPLObservation, REPLState]):
-    async def reset(...)
-    async def execute(code: str)
-    async def submit_final_answer(answer: str)
-    async def state()
-```
-
-Use `.sync()` for synchronous code.
-
-### Local Helpers
-
-```python
-class LocalREPLEnv:
-    def reset(...)
-    def execute(code: str)
-    def state()
-```
-
-```python
-class LocalRLMRunner:
-    def run(context: str, task_prompt: str, *, model: str | None = None) -> RLMRunResult
-```
-
-### Actions and Observations
+## Actions and Observations
 
 `REPLAction`
 
@@ -339,25 +262,13 @@ reward: float | None
 metadata: dict
 ```
 
-## Injected REPL Helpers
+## REPL Helpers
 
 When configured, the REPL namespace exposes:
 
-- `llm_query(prompt, model=None)`
-- `llm_query_batched(prompts, model=None)`
-- `rlm_query(prompt, model=None)`
-- `rlm_query_batched(prompts, model=None)`
-- `FINAL(value)`
-- `FINAL_VAR(name)`
-- `SHOW_VARS()`
-
-Notes:
-
-- `rlm_query` is the recursive child-run surface.
-- At max recursion depth, recursion falls back to direct LM calls rather than spawning more children.
-- Lifecycle callbacks follow the official `rlm` pattern:
-  - `on_subcall_start(depth, model, prompt_preview)`
-  - `on_subcall_complete(depth, model, duration, error_or_none)`
+- `llm_query(prompt, model=None)` and `llm_query_batched(prompts, model=None)`
+- `rlm_query(prompt, model=None)` and `rlm_query_batched(prompts, model=None)`: recursive child runs. At the maximum depth they fall back to direct model calls.
+- `FINAL(value)`, `FINAL_VAR(name)` and `SHOW_VARS()`
 
 ## Finalization Patterns
 
@@ -382,44 +293,39 @@ result = env.execute("answer['content'] = '42'")
 result = env.execute("answer['ready'] = True")
 ```
 
-## Prompt Utilities
-
-[`prompts.py`](prompts.py) contains the current message-building and parsing helpers used by the examples and runner.
+## Prompts and Examples
 
-Important exports:
+[`prompts.py`](prompts.py) has the system prompts and helpers used by the runner: `RLM_SYSTEM_PROMPT`, `RLM_SYSTEM_PROMPT_QWEN`, `QueryMetadata`, `build_rlm_system_prompt(...)`, `build_user_prompt(...)`, `extract_code_blocks(...)` and `format_observations(...)`.
 
-- `RLM_SYSTEM_PROMPT`
-- `RLM_SYSTEM_PROMPT_QWEN`
-- `QueryMetadata`
-- `build_rlm_system_prompt(...)`
-- `build_user_prompt(...)`
-- `extract_code_blocks(...)`
-- `format_observations(...)`
-
-These prompts were updated to reflect the actual helper surface the environment provides, rather than documenting tools that do not exist.
-
-## Examples
+Examples, which default to `Qwen/Qwen3.5-9B` through Hugging Face inference (needs `HF_TOKEN`):
 
 - [`examples/repl_with_llm.py`](../../examples/repl_with_llm.py)
 - [`examples/repl_oolong_simple.py`](../../examples/repl_oolong_simple.py)
 
-Default hosted model in the examples is currently `Qwen/Qwen3.5-9B`, but real hosted inference still depends on provider availability and token access.
-
-## Environment Variables
-
-Server-side configuration in [`server/app.py`](server/app.py):
-
-- `LLM_MODEL`
-- `HF_TOKEN`
-- `REPL_MAX_ITERATIONS`
-- `REPL_MAX_OUTPUT_LENGTH`
-- `REPL_CONTEXT_PREVIEW_LENGTH`
-- `REPL_RLM_MAX_DEPTH`
-- `REPL_RLM_MAX_ITERATIONS`
-
 ## References
 
 - [RLM Paper (arXiv:2512.24601)](https://huggingface.co/papers/2512.24601)
 - [RLM Implementation](https://github.com/alexzhang13/rlm)
 - [Alex Zhang's RLM Blog](https://alexzhang13.github.io/blog/2025/rlm/)
 - [Prime Intellect RLM Blog](https://www.primeintellect.ai/blog/rlm)
+
+## Hugging Face Space Deployment
+
+This Space is built from OpenEnv environment `repl_env`.
+
+- Space URL: `https://huggingface.co/spaces/openenv/repl`
+- OpenEnv pinned ref: `0.10.0`
+- Hub tag: `openenv`
+
+### Connecting from Code
+
+```bash
+pip install git+https://huggingface.co/spaces/openenv/repl
+```
+
+```python
+from repl_env import REPLEnv
+
+with REPLEnv(base_url="https://openenv-repl.hf.space").sync() as env:
+    result = env.reset()
+```

client.py

@@ -76,13 +76,14 @@ class REPLEnv(EnvClient[REPLAction, REPLObservation, REPLState]):
             max_iterations=obs_data.get("max_iterations", 30),
             done=payload.get("done", False),
             reward=payload.get("reward"),
-            metadata=obs_data.get("metadata", {}),
+            metadata=payload.get("metadata", obs_data.get("metadata", {})),
         )
 
         return StepResult(
             observation=observation,
             reward=payload.get("reward"),
             done=payload.get("done", False),
+            metadata=payload.get("metadata"),
         )
 
     def _parse_state(self, payload: dict[str, Any]) -> REPLState:

envs/repl_env/README.md

@@ -13,102 +13,28 @@ tags:
 
 # REPL Environment for OpenEnv
 
-`repl_env` is an OpenEnv-native Python REPL environment for Recursive Language Model style execution. It now follows the current OpenEnv client/server conventions:
-
-- `REPLEnv` is the remote async `EnvClient`
-- `.sync()` is the sync wrapper for remote usage
-- `LocalREPLEnv` is the explicit in-process helper
-- `LocalRLMRunner` is the higher-level orchestration loop for local recursive RLM runs
-
-The architecture is intentionally split the same way the official `rlm` and DSPy implementations split things:
-
-- the environment executes code and exposes tools
-- the runner owns the iterative prompting loop
-- recursive behavior lives in backend/controller modules, not in the executor
-
-## Overview
-
-Inside the REPL, the model can:
+`repl_env` is a Python REPL environment for [Recursive Language Model](https://huggingface.co/papers/2512.24601) (RLM) style execution. The model writes code that runs in a persistent namespace and can:
 
 - inspect `context`
-- execute Python code across multiple turns with persistent state
-- call `llm_query(...)` and `llm_query_batched(...)`
-- call `rlm_query(...)` and `rlm_query_batched(...)` for recursive child runs when configured
+- execute Python across multiple turns with persistent state
+- call `llm_query(...)` and `llm_query_batched(...)` to query a language model
+- call `rlm_query(...)` and `rlm_query_batched(...)` for recursive child runs, when configured
 - finish with `FINAL(...)`, `FINAL_VAR(...)`, or `answer = {"content": ..., "ready": True}`
 
-## Current Architecture
-
-Main modules:
-
-- [`client.py`](client.py): remote async OpenEnv client
-- [`local.py`](local.py): explicit in-process local env helper
-- [`runner.py`](runner.py): local RLM orchestration loop
-- [`recursive_backends.py`](recursive_backends.py): direct and recursive backend implementations
-- [`recursive_controller.py`](recursive_controller.py): server-side backend/broker composition
-- [`rubrics.py`](rubrics.py): reward rubrics (OpenEnv RFC 004)
-- [`server/repl_environment.py`](server/repl_environment.py): server-side execution environment
-- [`server/app.py`](server/app.py): OpenEnv HTTP server app and env factory
-
-## What Works Today
-
-- Standard remote OpenEnv usage through `REPLEnv`
-- Local in-process execution through `LocalREPLEnv`
-- Local recursive RLM runs through `LocalRLMRunner`
-- Server-backed recursive calls through the current controller/broker path
-- Explicit recursion controls:
-  - `max_depth`
-  - `max_children_total`
-  - `max_children_per_batch`
-  - `per_child_timeout_s`
-  - `result_truncation_limit`
-- Lightweight child trace metadata on local runner results
-- Rubric-based rewards (OpenEnv RFC 004):
-  - `ExactMatchRubric`: binary outcome reward against ground truth
-  - `FuzzyMatchRubric`: partial credit for containment matches
-  - `CustomMetricRubric`: user-provided `metric(expected, predicted) -> float`
-  - `CodeExecutionRubric`: per-step process reward for code errors
-  - `REPLRubric`: composite rubric combining outcome + process
-  - Ground truth injectable at reset via `expected_answer`
-
-## Rewards
-
-Rewards follow the OpenEnv Rubric system (RFC 004). The environment uses
-`REPLRubric` by default, which combines:
+The package provides:
 
-- **Outcome reward** (on terminal steps): compares `final_answer` against
-  `expected_answer` if provided. Returns 1.0 for match, 0.0 otherwise.
-- **Process reward** (on non-terminal steps): returns -0.05 for code
-  execution errors, 0.0 for successful steps.
-- **Failure reward**: returns -0.1 when max iterations exhausted without an answer.
+- `REPLEnv`: the async client for a remote server (`.sync()` for synchronous code), with `execute(code)`, `submit_final_answer(answer)`, `get_variable(name)` and `list_variables()` on top of `reset()`/`step()`
+- `LocalREPLEnv`: the same environment, in process
+- `LocalRLMRunner`: a local RLM loop that prompts a model, runs its code and handles recursion
 
-For RL training (GRPO, etc.), pass `expected_answer` at reset time:
-
-```python
-with LocalREPLEnv() as env:
-    env.reset(
-        context="...",
-        task_prompt="...",
-        expected_answer="42",  # ground truth for rubric scoring
-    )
-    result = env.execute("print(FINAL(42))")
-    print(result.reward)  # 1.0 (correct)
-```
-
-Custom rubrics can be injected at construction:
-
-```python
-from repl_env import LocalREPLEnv, CustomMetricRubric, REPLRubric
+## Quick Start
 
-def my_metric(expected, predicted):
-    return 1.0 if expected.strip() == predicted.strip() else 0.0
+Start a server:
 
-env = LocalREPLEnv(rubric=REPLRubric(outcome=CustomMetricRubric(my_metric)))
+```bash
+PYTHONPATH=src:envs uvicorn envs.repl_env.server.app:app --host 127.0.0.1 --port 8000
 ```
 
-## Quick Start
-
-### Remote Server Usage
-
 Async:
 
 ```python
@@ -145,7 +71,7 @@ with REPLEnv(base_url="http://127.0.0.1:8000").sync() as env:
     print(result.observation.result.stdout)
 ```
 
-### Local Environment Usage
+### In Process
 
 ```python
 from repl_env import LocalREPLEnv
@@ -160,7 +86,56 @@ with LocalREPLEnv() as env:
     print(env.state().final_answer)
 ```
 
-### Local Recursive RLM Usage
+## Server Configuration
+
+Environment variables read by [`server/app.py`](server/app.py):
+
+| Variable | Default | Description |
+|----------|---------|-------------|
+| `HF_TOKEN` | unset | Enables `llm_query` on the server. Without it, a client can pass `hf_token` to `reset()` |
+| `LLM_MODEL` | `Qwen/Qwen3.5-9B` | Default model for `llm_query`. A client can pass `llm_model` to `reset()` |
+| `REPL_MAX_ITERATIONS` | `30` | Maximum steps per episode |
+| `REPL_MAX_OUTPUT_LENGTH` | `8192` | Maximum captured output per step |
+| `REPL_CONTEXT_PREVIEW_LENGTH` | `500` | Length of `context_preview` in observations |
+| `REPL_RLM_MAX_DEPTH` | `2` | Maximum recursion depth for `rlm_query` |
+| `REPL_RLM_MAX_ITERATIONS` | `30` | Maximum iterations for recursive child runs |
+| `MAX_CONCURRENT_ENVS` | `8` | Maximum concurrent sessions |
+
+## Reward
+
+Rewards use the OpenEnv [rubric system](https://huggingface.co/docs/openenv/guides/rewards). The default `REPLRubric` combines:
+
+- **Outcome reward** (on terminal steps): compares `final_answer` against
+  `expected_answer` if provided. Returns 1.0 for match, 0.0 otherwise.
+- **Process reward** (on non-terminal steps): returns -0.05 for code
+  execution errors, 0.0 for successful steps.
+- **Failure reward**: returns -0.1 when max iterations exhausted without an answer.
+
+For RL training (GRPO, etc.), pass `expected_answer` at reset time:
+
+```python
+with LocalREPLEnv() as env:
+    env.reset(
+        context="...",
+        task_prompt="...",
+        expected_answer="42",  # ground truth for rubric scoring
+    )
+    result = env.execute("print(FINAL(42))")
+    print(result.reward)  # 1.0 (correct)
+```
+
+Other rubrics: `ExactMatchRubric` (binary match), `FuzzyMatchRubric` (1.0 for an exact match, 0.5 when the expected answer is contained in the final answer), `CustomMetricRubric` (your `metric(expected, predicted) -> float`) and `CodeExecutionRubric` (per-step error penalty). Pass one at construction:
+
+```python
+from repl_env import LocalREPLEnv, CustomMetricRubric, REPLRubric
+
+def my_metric(expected, predicted):
+    return 1.0 if expected.strip() == predicted.strip() else 0.0
+
+env = LocalREPLEnv(rubric=REPLRubric(outcome=CustomMetricRubric(my_metric)))
+```
+
+## Running an RLM Locally
 
 `LocalRLMRunner` takes any `chat_fn(messages, model=None) -> str`. It works
 with HF Inference API, vLLM, SGLang, Ollama, or any OpenAI-compatible server.
@@ -209,7 +184,7 @@ runner = LocalRLMRunner(chat_fn, max_iterations=30, max_depth=2)
 result = runner.run(context, task)
 ```
 
-### Using Different Models for Outer and Inner Loops
+### Different Models for Outer and Inner Loops
 
 The outer loop (code generation) can use a large model while inner
 `llm_query`/`rlm_query` calls use a smaller, faster model. Pass a
@@ -260,45 +235,9 @@ runner = LocalRLMRunner(
 result = runner.run(context, task)
 ```
 
-## Server
-
-Run the local server:
-
-```bash
-PYTHONPATH=src:envs uvicorn envs.repl_env.server.app:app --host 127.0.0.1 --port 8000
-```
-
-The server uses a proper OpenEnv environment factory in [`server/app.py`](server/app.py).
-
-## API Surface
-
-### Remote Client
-
-```python
-class REPLEnv(EnvClient[REPLAction, REPLObservation, REPLState]):
-    async def reset(...)
-    async def execute(code: str)
-    async def submit_final_answer(answer: str)
-    async def state()
-```
-
-Use `.sync()` for synchronous code.
-
-### Local Helpers
-
-```python
-class LocalREPLEnv:
-    def reset(...)
-    def execute(code: str)
-    def state()
-```
-
-```python
-class LocalRLMRunner:
-    def run(context: str, task_prompt: str, *, model: str | None = None) -> RLMRunResult
-```
+`LocalRLMRunner` also takes recursion limits (`max_children_total`, `max_children_per_batch`, `per_child_timeout_s`, `result_truncation_limit`) and lifecycle callbacks (`on_subcall_start(depth, model, prompt_preview)`, `on_subcall_complete(depth, model, duration, error_or_none)`). Its results carry lightweight child trace metadata.
 
-### Actions and Observations
+## Actions and Observations
 
 `REPLAction`
 
@@ -322,25 +261,13 @@ reward: float | None
 metadata: dict
 ```
 
-## Injected REPL Helpers
+## REPL Helpers
 
 When configured, the REPL namespace exposes:
 
-- `llm_query(prompt, model=None)`
-- `llm_query_batched(prompts, model=None)`
-- `rlm_query(prompt, model=None)`
-- `rlm_query_batched(prompts, model=None)`
-- `FINAL(value)`
-- `FINAL_VAR(name)`
-- `SHOW_VARS()`
-
-Notes:
-
-- `rlm_query` is the recursive child-run surface.
-- At max recursion depth, recursion falls back to direct LM calls rather than spawning more children.
-- Lifecycle callbacks follow the official `rlm` pattern:
-  - `on_subcall_start(depth, model, prompt_preview)`
-  - `on_subcall_complete(depth, model, duration, error_or_none)`
+- `llm_query(prompt, model=None)` and `llm_query_batched(prompts, model=None)`
+- `rlm_query(prompt, model=None)` and `rlm_query_batched(prompts, model=None)`: recursive child runs. At the maximum depth they fall back to direct model calls.
+- `FINAL(value)`, `FINAL_VAR(name)` and `SHOW_VARS()`
 
 ## Finalization Patterns
 
@@ -365,41 +292,15 @@ result = env.execute("answer['content'] = '42'")
 result = env.execute("answer['ready'] = True")
 ```
 
-## Prompt Utilities
-
-[`prompts.py`](prompts.py) contains the current message-building and parsing helpers used by the examples and runner.
+## Prompts and Examples
 
-Important exports:
+[`prompts.py`](prompts.py) has the system prompts and helpers used by the runner: `RLM_SYSTEM_PROMPT`, `RLM_SYSTEM_PROMPT_QWEN`, `QueryMetadata`, `build_rlm_system_prompt(...)`, `build_user_prompt(...)`, `extract_code_blocks(...)` and `format_observations(...)`.
 
-- `RLM_SYSTEM_PROMPT`
-- `RLM_SYSTEM_PROMPT_QWEN`
-- `QueryMetadata`
-- `build_rlm_system_prompt(...)`
-- `build_user_prompt(...)`
-- `extract_code_blocks(...)`
-- `format_observations(...)`
-
-These prompts were updated to reflect the actual helper surface the environment provides, rather than documenting tools that do not exist.
-
-## Examples
+Examples, which default to `Qwen/Qwen3.5-9B` through Hugging Face inference (needs `HF_TOKEN`):
 
 - [`examples/repl_with_llm.py`](../../examples/repl_with_llm.py)
 - [`examples/repl_oolong_simple.py`](../../examples/repl_oolong_simple.py)
 
-Default hosted model in the examples is currently `Qwen/Qwen3.5-9B`, but real hosted inference still depends on provider availability and token access.
-
-## Environment Variables
-
-Server-side configuration in [`server/app.py`](server/app.py):
-
-- `LLM_MODEL`
-- `HF_TOKEN`
-- `REPL_MAX_ITERATIONS`
-- `REPL_MAX_OUTPUT_LENGTH`
-- `REPL_CONTEXT_PREVIEW_LENGTH`
-- `REPL_RLM_MAX_DEPTH`
-- `REPL_RLM_MAX_ITERATIONS`
-
 ## References
 
 - [RLM Paper (arXiv:2512.24601)](https://huggingface.co/papers/2512.24601)

envs/repl_env/client.py

@@ -76,13 +76,14 @@ class REPLEnv(EnvClient[REPLAction, REPLObservation, REPLState]):
             max_iterations=obs_data.get("max_iterations", 30),
             done=payload.get("done", False),
             reward=payload.get("reward"),
-            metadata=obs_data.get("metadata", {}),
+            metadata=payload.get("metadata", obs_data.get("metadata", {})),
         )
 
         return StepResult(
             observation=observation,
             reward=payload.get("reward"),
             done=payload.get("done", False),
+            metadata=payload.get("metadata"),
         )
 
     def _parse_state(self, payload: dict[str, Any]) -> REPLState:

envs/repl_env/pyproject.toml

@@ -15,7 +15,7 @@ description = "Recursive Language Model REPL Environment for OpenEnv"
 requires-python = ">=3.10"
 dependencies = [
     # Core OpenEnv dependencies (required for server functionality)
-    "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git@v0.2.3",
+    "openenv @ git+https://github.com/huggingface/OpenEnv.git@v0.10.0",
     "fastapi>=0.115.0",
     "pydantic>=2.0.0",
     "uvicorn>=0.24.0",
@@ -26,11 +26,13 @@ dependencies = [
     "huggingface_hub>=0.20.0",
     # REPL custom Gradio tab on /web
     "gradio>=4.0.0",
+    # Document upload support in the Gradio tab (server-side only, not in the sandbox)
+    "pypdf>=6.16.1",
 ]
 
 [project.optional-dependencies]
 dev = [
-    "pytest>=8.0.0",
+    "pytest>=9.0.3",
     "pytest-cov>=4.0.0",
 ]
 

envs/repl_env/runner.py

@@ -243,10 +243,18 @@ class LocalRLMRunner:
         ]
         try:
             response = self._chat(final_prompt, model)
-            # Try to extract FINAL(...) from the response
-            match = re.search(r"FINAL\((.*?)\)", response, re.DOTALL)
-            if match:
-                return match.group(1).strip()
+            # Try to extract FINAL(...) from the response. Paren-counting handles
+            # nested parens and mid-sentence FINAL without regex flag trade-offs.
+            idx = response.find("FINAL(")
+            if idx != -1:
+                depth, start = 0, idx + len("FINAL")
+                for i, ch in enumerate(response[idx + len("FINAL"):], start=idx + len("FINAL")):
+                    if ch == "(":
+                        depth += 1
+                    elif ch == ")":
+                        depth -= 1
+                        if depth == 0:
+                            return response[start + 1 : i].strip()
             # If no FINAL pattern, return the raw response as best-effort
             return response.strip() if response.strip() else None
         except Exception:

envs/repl_env/server/Dockerfile

@@ -10,7 +10,7 @@
 # - Standalone environments (with openenv from pip)
 # The build script (openenv build) handles context detection and sets appropriate build args.
 
-ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
+ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest
 FROM ${BASE_IMAGE} AS builder
 
 WORKDIR /app
@@ -22,8 +22,8 @@ ARG ENV_NAME=repl_env
 # Copy environment code (always at root of build context)
 COPY . /app/env
 
-# For in-repo builds, openenv-core is already in the pyproject.toml dependencies
-# For standalone builds, openenv-core will be installed from pip via pyproject.toml
+# For in-repo builds, openenv is already in the pyproject.toml dependencies
+# For standalone builds, openenv will be installed from pip via pyproject.toml
 WORKDIR /app/env
 
 # Ensure uv is available (for local builds where base image lacks it)

envs/repl_env/server/app.py

@@ -33,6 +33,7 @@ Usage:
 Environment Variables:
     HF_TOKEN: Fallback HuggingFace API token (client token takes priority)
     LLM_MODEL: Model to use for llm_query/llm_query_batched (default: Qwen/Qwen3.5-9B)
+    MAX_CONCURRENT_ENVS: Maximum concurrent WebSocket sessions (default: 8)
 """
 
 import inspect
@@ -42,11 +43,11 @@ import os
 try:
     from openenv.core.env_server.http_server import create_app
 
-    from ..models import REPLAction, REPLObservation
+    from ..models import REPLAction, REPLObservation, REPLState
     from .gradio_ui import build_repl_gradio_app
     from .repl_environment import REPLEnvironment
 except ImportError:
-    from models import REPLAction, REPLObservation
+    from models import REPLAction, REPLObservation, REPLState
     from openenv.core.env_server.http_server import create_app
     from server.gradio_ui import build_repl_gradio_app
     from server.repl_environment import REPLEnvironment
@@ -60,6 +61,7 @@ REPL_MAX_OUTPUT_LENGTH = int(os.environ.get("REPL_MAX_OUTPUT_LENGTH", "8192"))
 REPL_CONTEXT_PREVIEW_LENGTH = int(os.environ.get("REPL_CONTEXT_PREVIEW_LENGTH", "500"))
 REPL_RLM_MAX_DEPTH = int(os.environ.get("REPL_RLM_MAX_DEPTH", "2"))
 REPL_RLM_MAX_ITERATIONS = int(os.environ.get("REPL_RLM_MAX_ITERATIONS", "30"))
+MAX_CONCURRENT_ENVS = int(os.environ.get("MAX_CONCURRENT_ENVS", "8"))
 # ==========================================
 
 _logger = logging.getLogger(__name__)
@@ -94,17 +96,34 @@ def create_repl_environment() -> REPLEnvironment:
 # Create the app with web interface and README integration.
 _sig = inspect.signature(create_app)
 if "gradio_builder" in _sig.parameters:
+    # Each kwarg is guarded by inspect.signature so older openenv
+    # releases that predate the param still boot this env.
+    create_app_kwargs: dict = {
+        "env_name": "repl_env",
+        "max_concurrent_envs": MAX_CONCURRENT_ENVS,
+        "gradio_builder": build_repl_gradio_app,
+    }
+    if "custom_tab_name" in _sig.parameters:
+        create_app_kwargs["custom_tab_name"] = "REPL"
+    if "custom_tab_primary" in _sig.parameters:
+        create_app_kwargs["custom_tab_primary"] = True
+    if "show_default_tab" in _sig.parameters:
+        create_app_kwargs["show_default_tab"] = False
+    if "title_override" in _sig.parameters:
+        create_app_kwargs["title_override"] = (
+            "OpenEnv REPL — Recursive Language Model playground"
+        )
+    if "state_cls" in _sig.parameters:
+        create_app_kwargs["state_cls"] = REPLState
     app = create_app(
         create_repl_environment,
         REPLAction,
         REPLObservation,
-        env_name="repl_env",
-        max_concurrent_envs=8,
-        gradio_builder=build_repl_gradio_app,
+        **create_app_kwargs,
     )
 else:
     _logger.warning(
-        "Installed openenv-core does not support gradio_builder; "
+        "Installed openenv does not support gradio_builder; "
         "custom REPL Gradio tab will not be available."
     )
     app = create_app(
@@ -112,7 +131,7 @@ else:
         REPLAction,
         REPLObservation,
         env_name="repl_env",
-        max_concurrent_envs=8,
+        max_concurrent_envs=MAX_CONCURRENT_ENVS,
     )
 
 

envs/repl_env/server/gradio_ui.py

@@ -8,57 +8,210 @@
 
 from __future__ import annotations
 
+import io
 import json
+import re
+import time
+from pathlib import Path
 from typing import Any, Dict, List, Optional
 
 import gradio as gr
 from openenv.core.env_server.types import EnvironmentMetadata
 
 
+_FINAL_RE = re.compile(r"FINAL\(([^\n]+?)\)\s*$", re.MULTILINE)
+
+
+def _extract_final(stdout: str) -> Optional[str]:
+    """Pull the final answer from the last `FINAL(...)` line of stdout."""
+    matches = _FINAL_RE.findall(stdout or "")
+    return matches[-1].strip() if matches else None
+
+
+try:
+    import pypdf
+except ImportError:  # pragma: no cover - optional dep
+    pypdf = None  # type: ignore[assignment]
+
+
+_EX1_NEEDLE_CONTEXT = """ACME Robotics — Q1 internal changelog, build notes, and incident log.
+
+Week 1. Fleet rollout reached 412 units across warehouses A and B. Latency regression on the manipulator firmware was traced back to the new servo driver; hotfix 1.4.2 shipped Wednesday. No customer impact.
+
+Week 2. Vision team rotated the training dataset to include low-light scenes. Offline eval improved 3.1% on the occluded-grasping benchmark but regressed 0.8% on the clean-table benchmark. Product signed off on the trade-off.
+
+Week 3. Security audit flagged credential handling in the pairing flow. The legacy MQTT broker was retired. Any robot still on firmware < 1.3.0 will refuse to pair after May 1 — document this in the migration note. The activation code is BANANA-747. Do not share externally. Re-keying ceremony for the HSM happens next sprint.
+
+Week 4. Onboarding four new SREs. Runbook for the arm-joint recalibration was rewritten from scratch after the Tuesday incident. New runbook lives under ops/arm/joint-recal/v2.md. The Tuesday incident: a collision with a pallet rack during autonomous pickup. Root cause: stale depth-map cache. Mitigation: cache TTL dropped from 5 minutes to 20 seconds.
+
+Week 5. Starting the unification work on the control plane. Two clusters are being collapsed into one, with staged traffic shifting. Rollback criteria are defined in the RFC. Early signal looks good: no increase in p99 pick latency. Next milestone: full cutover.
+
+Week 6. Offsite in Porto. Platform team proposed a new interface between the planner and the low-level controller. Discussion deferred to the RFC process. Interesting prior art from the MIT paper on whole-body MPC. Follow-up reading list shared in the team channel."""
+
+_EX1_TASK = "Find the activation code hidden somewhere in the changelog."
+
+_EX1_CODE = """# Needle in a haystack: fan out direct LM calls over chunks in parallel.
+chunks = [context[i:i+300] for i in range(0, len(context), 300)]
+prompts = [
+    "In the text below, find any hyphenated identifier of the form LETTERS-DIGITS "
+    "(for example: ABC-123, FOO-42). Reply with just that identifier. "
+    f"If there is none, reply with the single word NONE.\\n\\n---\\n{c}"
+    for c in chunks
+]
+answers = llm_query_batched(prompts)
+hit = next((a.strip() for a in answers if "NONE" not in a.upper()), "not found")
+print(f"spawned {len(chunks)} children")
+print(FINAL(hit))"""
+
+
+_EX2_CONTEXT = ""
+_EX2_TASK = "Show the recursive primitive: the child runs its own REPL agent."
+_EX2_CODE = """# rlm_query spawns a full recursive REPL agent as a child.
+# The child can reason, run code, and finalize with FINAL(...).
+answer = rlm_query("What is 15 * 23? Work it out step by step and respond with only the final number.")
+print(f"Child returned: {answer!r}")
+print(FINAL(answer.strip()))"""
+
+
+_EXAMPLES: Dict[str, tuple[str, str, str, str]] = {
+    # 4th element is `expected_answer`: RFC 004 ground truth, matched
+    # against FINAL(...) for reward scoring.
+    "Needle in a haystack (llm_query_batched)": (
+        _EX1_NEEDLE_CONTEXT,
+        _EX1_TASK,
+        _EX1_CODE,
+        "BANANA-747",
+    ),
+    "Recursive child agent (rlm_query)": (
+        _EX2_CONTEXT,
+        _EX2_TASK,
+        _EX2_CODE,
+        "345",
+    ),
+}
+
+
+_HELPERS_MD = """
+### REPL helpers
+
+These names are injected into the Python namespace once you Reset:
+
+- `context` — the string you passed in the *Context* field, available as a Python variable.
+- `llm_query(prompt, model=None)` — single direct call to the configured LLM.
+- `llm_query_batched(prompts, model=None)` — fan out N direct LLM calls in parallel.
+- `rlm_query(prompt)` / `rlm_query_batched(prompts)` — each child runs a full recursive REPL loop (deeper RLM pattern).
+- `FINAL(value)` — finalize the episode with `value` as the answer.
+- `FINAL_VAR("name")` — finalize with the value of the named variable.
+- `answer = {"content": ..., "ready": True}` — dict-based finalization.
+
+### Rewards (RFC 004)
+
+Pass **Expected answer** at Reset to enable rubric-based scoring: the observation's `reward` field becomes `1.0` when `FINAL(...)` matches (exact match, case-insensitive strip) and `0.0` otherwise. Leave blank to skip scoring.
+
+### Typical flow
+
+1. Paste your **Hugging Face Token** — required for any example that calls `llm_query` / `rlm_query`.
+2. Write Context + Task Prompt (or pick a demo from **💡 Load example**).
+3. Click **🔁 Reset episode** — loads your context and wires up the helpers above.
+4. Write Python in *Python Code* that uses those helpers and ends with `FINAL(...)`.
+5. Click **▶ Run**. Inspect *Stdout* and *Raw JSON response* (the latter contains child-call metadata for recursive runs).
+"""
+
+
+_MAX_UPLOAD_BYTES = 20 * 1024 * 1024  # 20 MB — comfortably fits any arXiv PDF
+
+
+def _extract_text_from_upload(file_path: str) -> str:
+    """Read a .txt or .pdf upload from disk and return its text content.
+
+    Runs in the Gradio server process (outside the smolagents sandbox),
+    so it can use `pypdf` and the filesystem without triggering the
+    REPL's import restrictions.
+    """
+    path = Path(file_path)
+    size = path.stat().st_size
+    if size > _MAX_UPLOAD_BYTES:
+        return (
+            f"[File too large ({size // (1024 * 1024)} MB). "
+            f"Cap is {_MAX_UPLOAD_BYTES // (1024 * 1024)} MB to protect "
+            "the shared Space process from OOM.]"
+        )
+    suffix = path.suffix.lower()
+    if suffix == ".txt":
+        return path.read_text(encoding="utf-8", errors="replace")
+    if suffix == ".pdf":
+        if pypdf is None:
+            return "[pypdf is not installed in this deployment; cannot parse PDFs.]"
+        with path.open("rb") as fh:
+            reader = pypdf.PdfReader(io.BytesIO(fh.read()))
+        pages = []
+        for page in reader.pages:
+            try:
+                pages.append(page.extract_text() or "")
+            except Exception as exc:  # pypdf can raise on malformed PDFs
+                pages.append(f"[page extraction failed: {exc}]")
+        text = "\n\n".join(pages).strip()
+        return text or "[No extractable text — PDF may be image-only.]"
+    return f"[Unsupported file type: {suffix}. Use .txt or .pdf.]"
+
+
 def _code_block(title: str, content: str) -> str:
     if not content:
         return ""
     return f"**{title}:**\n```text\n{content}\n```"
 
 
-def _format_repl_response(data: Dict[str, Any]) -> str:
-    """Render REPL observations in a compact Markdown view."""
+def _format_repl_response(
+    data: Dict[str, Any],
+    elapsed_s: Optional[float] = None,
+) -> str:
+    """Render REPL observations as a compact, screenshot-friendly view.
+
+    The FINAL answer is surfaced at the top as a quoted block, followed by
+    the task prompt echo, then stdout / stderr, and a one-line metadata
+    strip. Context preview / locals snapshot / available variables are
+    still in the raw JSON pane for debugging.
+    """
     observation = data.get("observation", {})
     result = observation.get("result", {})
-    sections: List[str] = ["# REPL Session"]
 
-    context_preview = observation.get("context_preview")
-    if context_preview:
-        sections.append(_code_block("Context Preview", context_preview))
+    sections: List[str] = []
 
-    task_prompt = observation.get("task_prompt")
-    if task_prompt:
-        sections.append(_code_block("Task Prompt", task_prompt))
+    stdout = (result.get("stdout") or "").rstrip()
+    final_answer = _extract_final(stdout) if stdout else None
+    if final_answer:
+        sections.append(f"### 🎯 Answer\n\n> **{final_answer}**")
 
-    available_variables = observation.get("available_variables") or []
-    if available_variables:
-        sections.append(
-            "**Available Variables:** " + ", ".join(f"`{name}`" for name in available_variables)
-        )
+    task_prompt = observation.get("metadata", {}).get("task_prompt")
+    if task_prompt:
+        sections.append(f"**Task:** {task_prompt}")
 
-    if result.get("locals_snapshot"):
-        sections.append(
-            "**Locals Snapshot:**\n```json\n"
-            + json.dumps(result["locals_snapshot"], indent=2, sort_keys=True)
-            + "\n```"
-        )
+    if stdout:
+        sections.append(_code_block("Stdout", stdout))
 
-    stdout = result.get("stdout", "")
-    stderr = result.get("stderr", "")
-    sections.append(_code_block("Stdout", stdout))
-    sections.append(_code_block("Stderr", stderr))
+    stderr = (result.get("stderr") or "").rstrip()
+    if stderr:
+        sections.append(_code_block("Stderr", stderr))
 
+    meta: List[str] = []
     reward = data.get("reward")
+    if reward is not None:
+        meta.append(f"Reward `{reward}`")
     done = data.get("done")
-    sections.append(f"**Reward:** `{reward}`")
-    sections.append(f"**Done:** `{done}`")
+    if done is not None:
+        meta.append(f"Done `{done}`")
+    iteration = observation.get("iteration")
+    max_iter = observation.get("max_iterations")
+    if iteration is not None and max_iter is not None:
+        meta.append(f"Step `{iteration}/{max_iter}`")
+    if elapsed_s is not None:
+        meta.append(f"Elapsed `{elapsed_s:.1f}s`")
+    if meta:
+        sections.append(" · ".join(meta))
 
-    return "\n\n".join(section for section in sections if section)
+    if not sections:
+        return "_No output yet. Reset and Run._"
+    return "\n\n".join(sections)
 
 
 def build_repl_gradio_app(
@@ -70,48 +223,62 @@ def build_repl_gradio_app(
     quick_start_md: str,
 ) -> gr.Blocks:
     """Build the REPL-specific Gradio tab."""
-    del action_fields, is_chat_env, metadata, quick_start_md
+    del action_fields, is_chat_env
+
+    # Skip `metadata.description`: the core default is a generic
+    # "<name> environment" placeholder that adds noise without content.
+    readme_content = (
+        metadata.readme_content if metadata and metadata.readme_content else ""
+    )
 
     async def reset_repl(
-        context: str,
-        task_prompt: str,
-        hf_token: str,
-        llm_model: str,
+        context: Optional[str],
+        task_prompt: Optional[str],
+        expected_answer: Optional[str],
+        hf_token: Optional[str],
+        llm_model: Optional[str],
     ):
+        # Gradio can pass None for untouched password textboxes. Normalize here.
         reset_kwargs: Dict[str, Any] = {}
-        if context.strip():
+        if context and context.strip():
             reset_kwargs["context"] = context
-        if task_prompt.strip():
+        if task_prompt and task_prompt.strip():
             reset_kwargs["task_prompt"] = task_prompt
-        if hf_token.strip():
-            reset_kwargs["hf_token"] = hf_token
-        if llm_model.strip():
+        if expected_answer and expected_answer.strip():
+            reset_kwargs["expected_answer"] = expected_answer.strip()
+        if hf_token and hf_token.strip():
+            reset_kwargs["hf_token"] = hf_token.strip()
+        if llm_model and llm_model.strip():
             reset_kwargs["llm_model"] = llm_model
 
         try:
+            start = time.monotonic()
             data = await web_manager.reset_environment(reset_kwargs)
+            elapsed_s = time.monotonic() - start
             state = web_manager.get_state()
             return (
-                _format_repl_response(data),
+                _format_repl_response(data, elapsed_s=elapsed_s),
                 json.dumps(data, indent=2, sort_keys=True),
                 json.dumps(state, indent=2, sort_keys=True),
-                "REPL reset complete.",
+                f"Reset complete in {elapsed_s:.1f}s.",
             )
         except Exception as exc:
             return ("", "", "", f"Error: {exc}")
 
-    async def run_code(code: str):
-        if not code.strip():
+    async def run_code(code: Optional[str]):
+        if not code or not code.strip():
             return ("", "", "", "Enter Python code to run.")
 
         try:
+            start = time.monotonic()
             data = await web_manager.step_environment({"code": code})
+            elapsed_s = time.monotonic() - start
             state = web_manager.get_state()
             return (
-                _format_repl_response(data),
+                _format_repl_response(data, elapsed_s=elapsed_s),
                 json.dumps(data, indent=2, sort_keys=True),
                 json.dumps(state, indent=2, sort_keys=True),
-                "Code executed.",
+                f"Code executed in {elapsed_s:.1f}s.",
             )
         except Exception as exc:
             return ("", "", "", f"Error: {exc}")
@@ -122,62 +289,132 @@ def build_repl_gradio_app(
         except Exception as exc:
             return f"Error: {exc}"
 
-    with gr.Blocks(title=f"{title} - REPL") as blocks:
+    def load_example(label: str):
+        if not label or label not in _EXAMPLES:
+            return (
+                gr.update(),
+                gr.update(),
+                gr.update(),
+                gr.update(),
+                gr.update(),
+            )
+        ctx, task, code, expected = _EXAMPLES[label]
+        # Auto-collapse the Context accordion once populated to keep the
+        # viewport focused on Task + Code + Output.
+        return ctx, task, code, expected, gr.update(open=False)
+
+    def load_uploaded_document(file_obj):
+        if file_obj is None:
+            return gr.update(), gr.update(), gr.update()
+        file_path = (
+            file_obj if isinstance(file_obj, str) else getattr(file_obj, "name", None)
+        )
+        if not file_path:
+            return gr.update(), "Upload error: could not read file path.", gr.update()
+        try:
+            text = _extract_text_from_upload(file_path)
+        except Exception as exc:
+            return gr.update(), f"Upload error: {exc}", gr.update()
+        return (
+            text,
+            f"Loaded {Path(file_path).name} ({len(text)} chars) into Context.",
+            gr.update(open=False),
+        )
+
+    with gr.Blocks(title=title) as blocks:
         gr.Markdown(
-            "# REPL Control Panel\n\n"
-            "Load a problem into the REPL, execute Python, and inspect state without "
-            "leaving the Space."
+            "# REPL Control Panel\n"
+            "*Recursive Language Model REPL — run agentic Python with recursive LM calls.*"
         )
+
         with gr.Row():
             with gr.Column(scale=2):
-                context = gr.Textbox(
-                    label="Context",
-                    placeholder="Problem context or source text...",
-                    lines=8,
-                )
+                with gr.Row():
+                    example_dropdown = gr.Dropdown(
+                        label="💡 Load example",
+                        choices=list(_EXAMPLES.keys()),
+                        value=None,
+                    )
+                    upload = gr.File(
+                        label="📎 Or drop a .txt / .pdf",
+                        file_types=[".txt", ".pdf"],
+                        type="filepath",
+                    )
                 task_prompt = gr.Textbox(
                     label="Task Prompt",
                     placeholder="What should the agent solve?",
-                    lines=3,
+                    lines=2,
                 )
-                with gr.Accordion("Optional Model Settings", open=False):
+                with gr.Accordion("Context", open=True) as context_accordion:
+                    context = gr.Textbox(
+                        label="",
+                        placeholder="Problem context or source text (auto-fills from Load example / Drop)...",
+                        lines=3,
+                        show_label=False,
+                    )
+                with gr.Accordion("Advanced options", open=False):
                     hf_token = gr.Textbox(
-                        label="Hugging Face Token",
-                        placeholder="Used only for this reset; not persisted in state",
+                        label="Hugging Face Token (required for llm_query / rlm_query)",
+                        placeholder="hf_...  — used only for this reset; not persisted",
                         type="password",
                     )
+                    expected_answer = gr.Textbox(
+                        label="Expected answer (activates reward scoring)",
+                        placeholder="e.g. BANANA-747  —  leave blank to skip",
+                    )
                     llm_model = gr.Textbox(
                         label="LLM Model",
-                        placeholder="Optional override for llm_query / rlm_query",
+                        placeholder="e.g. Qwen/Qwen2.5-7B-Instruct",
+                        info=(
+                            "Blank = server default (Qwen/Qwen3.5-9B, tuned for the "
+                            "`enable_thinking=False` chat template)."
+                        ),
                     )
+                with gr.Row():
+                    reset_btn = gr.Button("🔁 Reset episode", variant="secondary")
+                    status = gr.Textbox(
+                        label="Status",
+                        interactive=False,
+                        lines=1,
+                        scale=1,
+                    )
+            with gr.Column(scale=3):
                 code = gr.Textbox(
                     label="Python Code",
                     placeholder="count = len(context.split())",
-                    lines=10,
+                    lines=8,
                 )
                 with gr.Row():
-                    reset_btn = gr.Button("Reset", variant="secondary")
-                    run_btn = gr.Button("Run", variant="primary")
-                    state_btn = gr.Button("Get state", variant="secondary")
-                status = gr.Textbox(label="Status", interactive=False)
-            with gr.Column(scale=3):
-                session_view = gr.Markdown(
-                    value="# REPL Session\n\nReset the environment to start."
-                )
-                raw_json = gr.Code(
-                    label="Raw JSON response",
-                    language="json",
-                    interactive=False,
-                )
-                state_json = gr.Code(
-                    label="Session state",
-                    language="json",
-                    interactive=False,
-                )
+                    run_btn = gr.Button("▶ Run", variant="primary")
+                    state_btn = gr.Button("📋 Get state", variant="secondary")
+                session_view = gr.Markdown(value="_Reset and Run to see output._")
+                with gr.Accordion("Raw JSON response", open=False):
+                    raw_json = gr.Code(
+                        label="",
+                        language="json",
+                        interactive=False,
+                    )
+                with gr.Accordion("Session state", open=False):
+                    state_json = gr.Code(
+                        label="",
+                        language="json",
+                        interactive=False,
+                    )
+
+        with gr.Accordion("REPL helpers reference & typical flow", open=False):
+            gr.Markdown(_HELPERS_MD)
+
+        if quick_start_md:
+            with gr.Accordion("API quick start (Python client snippet)", open=False):
+                gr.Markdown(quick_start_md)
+
+        if readme_content:
+            with gr.Accordion("Full README", open=False):
+                gr.Markdown(readme_content)
 
         reset_btn.click(
             fn=reset_repl,
-            inputs=[context, task_prompt, hf_token, llm_model],
+            inputs=[context, task_prompt, expected_answer, hf_token, llm_model],
             outputs=[session_view, raw_json, state_json, status],
         )
         run_btn.click(
@@ -191,5 +428,21 @@ def build_repl_gradio_app(
             outputs=[session_view, raw_json, state_json, status],
         )
         state_btn.click(fn=get_state_sync, outputs=[state_json])
+        example_dropdown.change(
+            fn=load_example,
+            inputs=[example_dropdown],
+            outputs=[
+                context,
+                task_prompt,
+                code,
+                expected_answer,
+                context_accordion,
+            ],
+        )
+        upload.change(
+            fn=load_uploaded_document,
+            inputs=[upload],
+            outputs=[context, status, context_accordion],
+        )
 
     return blocks

envs/repl_env/server/repl_environment.py

@@ -33,7 +33,12 @@ try:
     from ..models import CodeBlockResult, REPLAction, REPLObservation, REPLState
 except ImportError:
     try:
-        from repl_env.models import CodeBlockResult, REPLAction, REPLObservation, REPLState
+        from repl_env.models import (
+            CodeBlockResult,
+            REPLAction,
+            REPLObservation,
+            REPLState,
+        )
     except ImportError:
         from models import CodeBlockResult, REPLAction, REPLObservation, REPLState
 
@@ -47,10 +52,11 @@ except ImportError:
         from repl_env.rubrics import REPLRubric
         from .python_executor import PythonExecutor
     except ImportError:
-        from .python_executor import PythonExecutor
         from recursive_controller import create_server_recursive_controller
         from rubrics import REPLRubric
 
+        from .python_executor import PythonExecutor
+
 
 class REPLEnvironment(Environment):
     """
@@ -138,6 +144,12 @@ class REPLEnvironment(Environment):
         self._executor: Optional[PythonExecutor] = None
         self._runtime_controller = None
         self._runtime_controller_chat_fn: Optional[Callable[..., str]] = None
+        self._current_llm_model: Optional[str] = None
+
+    @staticmethod
+    def _resolve_model(llm_model: Optional[str]) -> str:
+        """Resolve the effective model name: explicit arg > env var > hard default."""
+        return llm_model or os.environ.get("LLM_MODEL", "Qwen/Qwen3.5-9B")
 
     @staticmethod
     def _build_hf_chat_fn(
@@ -149,7 +161,7 @@ class REPLEnvironment(Environment):
         except ImportError:
             raise RuntimeError("huggingface_hub is required for HF-backed recursion")
 
-        default_model = llm_model or os.environ.get("LLM_MODEL", "Qwen/Qwen3.5-9B")
+        default_model = REPLEnvironment._resolve_model(llm_model)
         client = InferenceClient(model=default_model, token=hf_token, timeout=300)
 
         def chat_fn(messages: list[dict[str, str]], model: str | None = None) -> str:
@@ -255,13 +267,23 @@ class REPLEnvironment(Environment):
         self.rlm_max_depth = runtime_rlm_max_depth
         self.rlm_max_iterations = runtime_rlm_max_iterations
 
-        # Create or rebuild LLM functions when needed.
         # Token resolution: explicit hf_token > HF_TOKEN env var > cached HF login.
-        if not self.llm_query_fn:
+        # Compare resolved model names so reset(llm_model=<default>) and
+        # reset() are treated as equal and don't trigger a redundant rebuild.
+        resolved_model = self._resolve_model(llm_model)
+        has_runtime_llm = self._runtime_controller is not None
+        model_changed = (
+            has_runtime_llm and resolved_model != self._current_llm_model
+        )
+        token_provided = hf_token is not None
+        if not self.llm_query_fn or model_changed or token_provided:
             effective_token = (
                 hf_token if hf_token is not None else os.environ.get("HF_TOKEN")
             )
+            # Pass raw llm_model; _build_hf_chat_fn resolves it via
+            # _resolve_model so the client matches the cache key below.
             self._create_llm_functions(effective_token, llm_model)
+            self._current_llm_model = resolved_model
         elif depth_changed and self._runtime_controller is not None:
             # Rebuild controller with new depth/iteration config but reuse
             # the existing chat_fn — don't require re-providing credentials.
@@ -514,10 +536,18 @@ class REPLEnvironment(Environment):
         Returns:
             Final answer string or None if not found
         """
-        # Pattern 1: RLM-style FINAL(answer)
-        final_match = re.search(r"FINAL\((.*?)\)", stdout, re.DOTALL)
-        if final_match:
-            return final_match.group(1).strip()
+        # Pattern 1: RLM-style FINAL(answer). Paren-counting handles nested
+        # parens (e.g. FINAL(f(x))) and multi-line values without regex flag trade-offs.
+        idx = stdout.find("FINAL(")
+        if idx != -1:
+            depth, start = 0, idx + len("FINAL")
+            for i, ch in enumerate(stdout[idx + len("FINAL"):], start=idx + len("FINAL")):
+                if ch == "(":
+                    depth += 1
+                elif ch == ")":
+                    depth -= 1
+                    if depth == 0:
+                        return stdout[start + 1 : i].strip()
 
         # Pattern 2: RLM-style FINAL_VAR(variable_name)
         final_var_match = re.search(r"FINAL_VAR\((\w+)\)", stdout)

pyproject.toml

@@ -15,7 +15,7 @@ description = "Recursive Language Model REPL Environment for OpenEnv"
 requires-python = ">=3.10"
 dependencies = [
     # Core OpenEnv dependencies (required for server functionality)
-    "openenv-core[core] @ git+https://github.com/meta-pytorch/OpenEnv.git@v0.2.3",
+    "openenv @ git+https://github.com/huggingface/OpenEnv.git@v0.10.0",
     "fastapi>=0.115.0",
     "pydantic>=2.0.0",
     "uvicorn>=0.24.0",
@@ -26,11 +26,13 @@ dependencies = [
     "huggingface_hub>=0.20.0",
     # REPL custom Gradio tab on /web
     "gradio>=4.0.0",
+    # Document upload support in the Gradio tab (server-side only, not in the sandbox)
+    "pypdf>=6.16.1",
 ]
 
 [project.optional-dependencies]
 dev = [
-    "pytest>=8.0.0",
+    "pytest>=9.0.3",
     "pytest-cov>=4.0.0",
 ]
 

runner.py

@@ -243,10 +243,18 @@ class LocalRLMRunner:
         ]
         try:
             response = self._chat(final_prompt, model)
-            # Try to extract FINAL(...) from the response
-            match = re.search(r"FINAL\((.*?)\)", response, re.DOTALL)
-            if match:
-                return match.group(1).strip()
+            # Try to extract FINAL(...) from the response. Paren-counting handles
+            # nested parens and mid-sentence FINAL without regex flag trade-offs.
+            idx = response.find("FINAL(")
+            if idx != -1:
+                depth, start = 0, idx + len("FINAL")
+                for i, ch in enumerate(response[idx + len("FINAL"):], start=idx + len("FINAL")):
+                    if ch == "(":
+                        depth += 1
+                    elif ch == ")":
+                        depth -= 1
+                        if depth == 0:
+                            return response[start + 1 : i].strip()
             # If no FINAL pattern, return the raw response as best-effort
             return response.strip() if response.strip() else None
         except Exception:

server/Dockerfile

@@ -10,7 +10,7 @@
 # - Standalone environments (with openenv from pip)
 # The build script (openenv build) handles context detection and sets appropriate build args.
 
-ARG BASE_IMAGE=ghcr.io/meta-pytorch/openenv-base:latest
+ARG BASE_IMAGE=ghcr.io/huggingface/openenv-base:latest
 FROM ${BASE_IMAGE} AS builder
 
 WORKDIR /app
@@ -22,8 +22,8 @@ ARG ENV_NAME=repl_env
 # Copy environment code (always at root of build context)
 COPY . /app/env
 
-# For in-repo builds, openenv-core is already in the pyproject.toml dependencies
-# For standalone builds, openenv-core will be installed from pip via pyproject.toml
+# For in-repo builds, openenv is already in the pyproject.toml dependencies
+# For standalone builds, openenv will be installed from pip via pyproject.toml
 WORKDIR /app/env
 
 # Ensure uv is available (for local builds where base image lacks it)

server/app.py

@@ -33,6 +33,7 @@ Usage:
 Environment Variables:
     HF_TOKEN: Fallback HuggingFace API token (client token takes priority)
     LLM_MODEL: Model to use for llm_query/llm_query_batched (default: Qwen/Qwen3.5-9B)
+    MAX_CONCURRENT_ENVS: Maximum concurrent WebSocket sessions (default: 8)
 """
 
 import inspect
@@ -42,11 +43,11 @@ import os
 try:
     from openenv.core.env_server.http_server import create_app
 
-    from ..models import REPLAction, REPLObservation
+    from ..models import REPLAction, REPLObservation, REPLState
     from .gradio_ui import build_repl_gradio_app
     from .repl_environment import REPLEnvironment
 except ImportError:
-    from models import REPLAction, REPLObservation
+    from models import REPLAction, REPLObservation, REPLState
     from openenv.core.env_server.http_server import create_app
     from server.gradio_ui import build_repl_gradio_app
     from server.repl_environment import REPLEnvironment
@@ -60,6 +61,7 @@ REPL_MAX_OUTPUT_LENGTH = int(os.environ.get("REPL_MAX_OUTPUT_LENGTH", "8192"))
 REPL_CONTEXT_PREVIEW_LENGTH = int(os.environ.get("REPL_CONTEXT_PREVIEW_LENGTH", "500"))
 REPL_RLM_MAX_DEPTH = int(os.environ.get("REPL_RLM_MAX_DEPTH", "2"))
 REPL_RLM_MAX_ITERATIONS = int(os.environ.get("REPL_RLM_MAX_ITERATIONS", "30"))
+MAX_CONCURRENT_ENVS = int(os.environ.get("MAX_CONCURRENT_ENVS", "8"))
 # ==========================================
 
 _logger = logging.getLogger(__name__)
@@ -94,17 +96,34 @@ def create_repl_environment() -> REPLEnvironment:
 # Create the app with web interface and README integration.
 _sig = inspect.signature(create_app)
 if "gradio_builder" in _sig.parameters:
+    # Each kwarg is guarded by inspect.signature so older openenv
+    # releases that predate the param still boot this env.
+    create_app_kwargs: dict = {
+        "env_name": "repl_env",
+        "max_concurrent_envs": MAX_CONCURRENT_ENVS,
+        "gradio_builder": build_repl_gradio_app,
+    }
+    if "custom_tab_name" in _sig.parameters:
+        create_app_kwargs["custom_tab_name"] = "REPL"
+    if "custom_tab_primary" in _sig.parameters:
+        create_app_kwargs["custom_tab_primary"] = True
+    if "show_default_tab" in _sig.parameters:
+        create_app_kwargs["show_default_tab"] = False
+    if "title_override" in _sig.parameters:
+        create_app_kwargs["title_override"] = (
+            "OpenEnv REPL — Recursive Language Model playground"
+        )
+    if "state_cls" in _sig.parameters:
+        create_app_kwargs["state_cls"] = REPLState
     app = create_app(
         create_repl_environment,
         REPLAction,
         REPLObservation,
-        env_name="repl_env",
-        max_concurrent_envs=8,
-        gradio_builder=build_repl_gradio_app,
+        **create_app_kwargs,
     )
 else:
     _logger.warning(
-        "Installed openenv-core does not support gradio_builder; "
+        "Installed openenv does not support gradio_builder; "
         "custom REPL Gradio tab will not be available."
     )
     app = create_app(
@@ -112,7 +131,7 @@ else:
         REPLAction,
         REPLObservation,
         env_name="repl_env",
-        max_concurrent_envs=8,
+        max_concurrent_envs=MAX_CONCURRENT_ENVS,
     )
 
 

server/gradio_ui.py

@@ -8,57 +8,210 @@
 
 from __future__ import annotations
 
+import io
 import json
+import re
+import time
+from pathlib import Path
 from typing import Any, Dict, List, Optional
 
 import gradio as gr
 from openenv.core.env_server.types import EnvironmentMetadata
 
 
+_FINAL_RE = re.compile(r"FINAL\(([^\n]+?)\)\s*$", re.MULTILINE)
+
+
+def _extract_final(stdout: str) -> Optional[str]:
+    """Pull the final answer from the last `FINAL(...)` line of stdout."""
+    matches = _FINAL_RE.findall(stdout or "")
+    return matches[-1].strip() if matches else None
+
+
+try:
+    import pypdf
+except ImportError:  # pragma: no cover - optional dep
+    pypdf = None  # type: ignore[assignment]
+
+
+_EX1_NEEDLE_CONTEXT = """ACME Robotics — Q1 internal changelog, build notes, and incident log.
+
+Week 1. Fleet rollout reached 412 units across warehouses A and B. Latency regression on the manipulator firmware was traced back to the new servo driver; hotfix 1.4.2 shipped Wednesday. No customer impact.
+
+Week 2. Vision team rotated the training dataset to include low-light scenes. Offline eval improved 3.1% on the occluded-grasping benchmark but regressed 0.8% on the clean-table benchmark. Product signed off on the trade-off.
+
+Week 3. Security audit flagged credential handling in the pairing flow. The legacy MQTT broker was retired. Any robot still on firmware < 1.3.0 will refuse to pair after May 1 — document this in the migration note. The activation code is BANANA-747. Do not share externally. Re-keying ceremony for the HSM happens next sprint.
+
+Week 4. Onboarding four new SREs. Runbook for the arm-joint recalibration was rewritten from scratch after the Tuesday incident. New runbook lives under ops/arm/joint-recal/v2.md. The Tuesday incident: a collision with a pallet rack during autonomous pickup. Root cause: stale depth-map cache. Mitigation: cache TTL dropped from 5 minutes to 20 seconds.
+
+Week 5. Starting the unification work on the control plane. Two clusters are being collapsed into one, with staged traffic shifting. Rollback criteria are defined in the RFC. Early signal looks good: no increase in p99 pick latency. Next milestone: full cutover.
+
+Week 6. Offsite in Porto. Platform team proposed a new interface between the planner and the low-level controller. Discussion deferred to the RFC process. Interesting prior art from the MIT paper on whole-body MPC. Follow-up reading list shared in the team channel."""
+
+_EX1_TASK = "Find the activation code hidden somewhere in the changelog."
+
+_EX1_CODE = """# Needle in a haystack: fan out direct LM calls over chunks in parallel.
+chunks = [context[i:i+300] for i in range(0, len(context), 300)]
+prompts = [
+    "In the text below, find any hyphenated identifier of the form LETTERS-DIGITS "
+    "(for example: ABC-123, FOO-42). Reply with just that identifier. "
+    f"If there is none, reply with the single word NONE.\\n\\n---\\n{c}"
+    for c in chunks
+]
+answers = llm_query_batched(prompts)
+hit = next((a.strip() for a in answers if "NONE" not in a.upper()), "not found")
+print(f"spawned {len(chunks)} children")
+print(FINAL(hit))"""
+
+
+_EX2_CONTEXT = ""
+_EX2_TASK = "Show the recursive primitive: the child runs its own REPL agent."
+_EX2_CODE = """# rlm_query spawns a full recursive REPL agent as a child.
+# The child can reason, run code, and finalize with FINAL(...).
+answer = rlm_query("What is 15 * 23? Work it out step by step and respond with only the final number.")
+print(f"Child returned: {answer!r}")
+print(FINAL(answer.strip()))"""
+
+
+_EXAMPLES: Dict[str, tuple[str, str, str, str]] = {
+    # 4th element is `expected_answer`: RFC 004 ground truth, matched
+    # against FINAL(...) for reward scoring.
+    "Needle in a haystack (llm_query_batched)": (
+        _EX1_NEEDLE_CONTEXT,
+        _EX1_TASK,
+        _EX1_CODE,
+        "BANANA-747",
+    ),
+    "Recursive child agent (rlm_query)": (
+        _EX2_CONTEXT,
+        _EX2_TASK,
+        _EX2_CODE,
+        "345",
+    ),
+}
+
+
+_HELPERS_MD = """
+### REPL helpers
+
+These names are injected into the Python namespace once you Reset:
+
+- `context` — the string you passed in the *Context* field, available as a Python variable.
+- `llm_query(prompt, model=None)` — single direct call to the configured LLM.
+- `llm_query_batched(prompts, model=None)` — fan out N direct LLM calls in parallel.
+- `rlm_query(prompt)` / `rlm_query_batched(prompts)` — each child runs a full recursive REPL loop (deeper RLM pattern).
+- `FINAL(value)` — finalize the episode with `value` as the answer.
+- `FINAL_VAR("name")` — finalize with the value of the named variable.
+- `answer = {"content": ..., "ready": True}` — dict-based finalization.
+
+### Rewards (RFC 004)
+
+Pass **Expected answer** at Reset to enable rubric-based scoring: the observation's `reward` field becomes `1.0` when `FINAL(...)` matches (exact match, case-insensitive strip) and `0.0` otherwise. Leave blank to skip scoring.
+
+### Typical flow
+
+1. Paste your **Hugging Face Token** — required for any example that calls `llm_query` / `rlm_query`.
+2. Write Context + Task Prompt (or pick a demo from **💡 Load example**).
+3. Click **🔁 Reset episode** — loads your context and wires up the helpers above.
+4. Write Python in *Python Code* that uses those helpers and ends with `FINAL(...)`.
+5. Click **▶ Run**. Inspect *Stdout* and *Raw JSON response* (the latter contains child-call metadata for recursive runs).
+"""
+
+
+_MAX_UPLOAD_BYTES = 20 * 1024 * 1024  # 20 MB — comfortably fits any arXiv PDF
+
+
+def _extract_text_from_upload(file_path: str) -> str:
+    """Read a .txt or .pdf upload from disk and return its text content.
+
+    Runs in the Gradio server process (outside the smolagents sandbox),
+    so it can use `pypdf` and the filesystem without triggering the
+    REPL's import restrictions.
+    """
+    path = Path(file_path)
+    size = path.stat().st_size
+    if size > _MAX_UPLOAD_BYTES:
+        return (
+            f"[File too large ({size // (1024 * 1024)} MB). "
+            f"Cap is {_MAX_UPLOAD_BYTES // (1024 * 1024)} MB to protect "
+            "the shared Space process from OOM.]"
+        )
+    suffix = path.suffix.lower()
+    if suffix == ".txt":
+        return path.read_text(encoding="utf-8", errors="replace")
+    if suffix == ".pdf":
+        if pypdf is None:
+            return "[pypdf is not installed in this deployment; cannot parse PDFs.]"
+        with path.open("rb") as fh:
+            reader = pypdf.PdfReader(io.BytesIO(fh.read()))
+        pages = []
+        for page in reader.pages:
+            try:
+                pages.append(page.extract_text() or "")
+            except Exception as exc:  # pypdf can raise on malformed PDFs
+                pages.append(f"[page extraction failed: {exc}]")
+        text = "\n\n".join(pages).strip()
+        return text or "[No extractable text — PDF may be image-only.]"
+    return f"[Unsupported file type: {suffix}. Use .txt or .pdf.]"
+
+
 def _code_block(title: str, content: str) -> str:
     if not content:
         return ""
     return f"**{title}:**\n```text\n{content}\n```"
 
 
-def _format_repl_response(data: Dict[str, Any]) -> str:
-    """Render REPL observations in a compact Markdown view."""
+def _format_repl_response(
+    data: Dict[str, Any],
+    elapsed_s: Optional[float] = None,
+) -> str:
+    """Render REPL observations as a compact, screenshot-friendly view.
+
+    The FINAL answer is surfaced at the top as a quoted block, followed by
+    the task prompt echo, then stdout / stderr, and a one-line metadata
+    strip. Context preview / locals snapshot / available variables are
+    still in the raw JSON pane for debugging.
+    """
     observation = data.get("observation", {})
     result = observation.get("result", {})
-    sections: List[str] = ["# REPL Session"]
 
-    context_preview = observation.get("context_preview")
-    if context_preview:
-        sections.append(_code_block("Context Preview", context_preview))
+    sections: List[str] = []
 
-    task_prompt = observation.get("task_prompt")
-    if task_prompt:
-        sections.append(_code_block("Task Prompt", task_prompt))
+    stdout = (result.get("stdout") or "").rstrip()
+    final_answer = _extract_final(stdout) if stdout else None
+    if final_answer:
+        sections.append(f"### 🎯 Answer\n\n> **{final_answer}**")
 
-    available_variables = observation.get("available_variables") or []
-    if available_variables:
-        sections.append(
-            "**Available Variables:** " + ", ".join(f"`{name}`" for name in available_variables)
-        )
+    task_prompt = observation.get("metadata", {}).get("task_prompt")
+    if task_prompt:
+        sections.append(f"**Task:** {task_prompt}")
 
-    if result.get("locals_snapshot"):
-        sections.append(
-            "**Locals Snapshot:**\n```json\n"
-            + json.dumps(result["locals_snapshot"], indent=2, sort_keys=True)
-            + "\n```"
-        )
+    if stdout:
+        sections.append(_code_block("Stdout", stdout))
 
-    stdout = result.get("stdout", "")
-    stderr = result.get("stderr", "")
-    sections.append(_code_block("Stdout", stdout))
-    sections.append(_code_block("Stderr", stderr))
+    stderr = (result.get("stderr") or "").rstrip()
+    if stderr:
+        sections.append(_code_block("Stderr", stderr))
 
+    meta: List[str] = []
     reward = data.get("reward")
+    if reward is not None:
+        meta.append(f"Reward `{reward}`")
     done = data.get("done")
-    sections.append(f"**Reward:** `{reward}`")
-    sections.append(f"**Done:** `{done}`")
+    if done is not None:
+        meta.append(f"Done `{done}`")
+    iteration = observation.get("iteration")
+    max_iter = observation.get("max_iterations")
+    if iteration is not None and max_iter is not None:
+        meta.append(f"Step `{iteration}/{max_iter}`")
+    if elapsed_s is not None:
+        meta.append(f"Elapsed `{elapsed_s:.1f}s`")
+    if meta:
+        sections.append(" · ".join(meta))
 
-    return "\n\n".join(section for section in sections if section)
+    if not sections:
+        return "_No output yet. Reset and Run._"
+    return "\n\n".join(sections)
 
 
 def build_repl_gradio_app(
@@ -70,48 +223,62 @@ def build_repl_gradio_app(
     quick_start_md: str,
 ) -> gr.Blocks:
     """Build the REPL-specific Gradio tab."""
-    del action_fields, is_chat_env, metadata, quick_start_md
+    del action_fields, is_chat_env
+
+    # Skip `metadata.description`: the core default is a generic
+    # "<name> environment" placeholder that adds noise without content.
+    readme_content = (
+        metadata.readme_content if metadata and metadata.readme_content else ""
+    )
 
     async def reset_repl(
-        context: str,
-        task_prompt: str,
-        hf_token: str,
-        llm_model: str,
+        context: Optional[str],
+        task_prompt: Optional[str],
+        expected_answer: Optional[str],
+        hf_token: Optional[str],
+        llm_model: Optional[str],
     ):
+        # Gradio can pass None for untouched password textboxes. Normalize here.
         reset_kwargs: Dict[str, Any] = {}
-        if context.strip():
+        if context and context.strip():
             reset_kwargs["context"] = context
-        if task_prompt.strip():
+        if task_prompt and task_prompt.strip():
             reset_kwargs["task_prompt"] = task_prompt
-        if hf_token.strip():
-            reset_kwargs["hf_token"] = hf_token
-        if llm_model.strip():
+        if expected_answer and expected_answer.strip():
+            reset_kwargs["expected_answer"] = expected_answer.strip()
+        if hf_token and hf_token.strip():
+            reset_kwargs["hf_token"] = hf_token.strip()
+        if llm_model and llm_model.strip():
             reset_kwargs["llm_model"] = llm_model
 
         try:
+            start = time.monotonic()
             data = await web_manager.reset_environment(reset_kwargs)
+            elapsed_s = time.monotonic() - start
             state = web_manager.get_state()
             return (
-                _format_repl_response(data),
+                _format_repl_response(data, elapsed_s=elapsed_s),
                 json.dumps(data, indent=2, sort_keys=True),
                 json.dumps(state, indent=2, sort_keys=True),
-                "REPL reset complete.",
+                f"Reset complete in {elapsed_s:.1f}s.",
             )
         except Exception as exc:
             return ("", "", "", f"Error: {exc}")
 
-    async def run_code(code: str):
-        if not code.strip():
+    async def run_code(code: Optional[str]):
+        if not code or not code.strip():
             return ("", "", "", "Enter Python code to run.")
 
         try:
+            start = time.monotonic()
             data = await web_manager.step_environment({"code": code})
+            elapsed_s = time.monotonic() - start
             state = web_manager.get_state()
             return (
-                _format_repl_response(data),
+                _format_repl_response(data, elapsed_s=elapsed_s),
                 json.dumps(data, indent=2, sort_keys=True),
                 json.dumps(state, indent=2, sort_keys=True),
-                "Code executed.",
+                f"Code executed in {elapsed_s:.1f}s.",
             )
         except Exception as exc:
             return ("", "", "", f"Error: {exc}")
@@ -122,62 +289,132 @@ def build_repl_gradio_app(
         except Exception as exc:
             return f"Error: {exc}"
 
-    with gr.Blocks(title=f"{title} - REPL") as blocks:
+    def load_example(label: str):
+        if not label or label not in _EXAMPLES:
+            return (
+                gr.update(),
+                gr.update(),
+                gr.update(),
+                gr.update(),
+                gr.update(),
+            )
+        ctx, task, code, expected = _EXAMPLES[label]
+        # Auto-collapse the Context accordion once populated to keep the
+        # viewport focused on Task + Code + Output.
+        return ctx, task, code, expected, gr.update(open=False)
+
+    def load_uploaded_document(file_obj):
+        if file_obj is None:
+            return gr.update(), gr.update(), gr.update()
+        file_path = (
+            file_obj if isinstance(file_obj, str) else getattr(file_obj, "name", None)
+        )
+        if not file_path:
+            return gr.update(), "Upload error: could not read file path.", gr.update()
+        try:
+            text = _extract_text_from_upload(file_path)
+        except Exception as exc:
+            return gr.update(), f"Upload error: {exc}", gr.update()
+        return (
+            text,
+            f"Loaded {Path(file_path).name} ({len(text)} chars) into Context.",
+            gr.update(open=False),
+        )
+
+    with gr.Blocks(title=title) as blocks:
         gr.Markdown(
-            "# REPL Control Panel\n\n"
-            "Load a problem into the REPL, execute Python, and inspect state without "
-            "leaving the Space."
+            "# REPL Control Panel\n"
+            "*Recursive Language Model REPL — run agentic Python with recursive LM calls.*"
         )
+
         with gr.Row():
             with gr.Column(scale=2):
-                context = gr.Textbox(
-                    label="Context",
-                    placeholder="Problem context or source text...",
-                    lines=8,
-                )
+                with gr.Row():
+                    example_dropdown = gr.Dropdown(
+                        label="💡 Load example",
+                        choices=list(_EXAMPLES.keys()),
+                        value=None,
+                    )
+                    upload = gr.File(
+                        label="📎 Or drop a .txt / .pdf",
+                        file_types=[".txt", ".pdf"],
+                        type="filepath",
+                    )
                 task_prompt = gr.Textbox(
                     label="Task Prompt",
                     placeholder="What should the agent solve?",
-                    lines=3,
+                    lines=2,
                 )
-                with gr.Accordion("Optional Model Settings", open=False):
+                with gr.Accordion("Context", open=True) as context_accordion:
+                    context = gr.Textbox(
+                        label="",
+                        placeholder="Problem context or source text (auto-fills from Load example / Drop)...",
+                        lines=3,
+                        show_label=False,
+                    )
+                with gr.Accordion("Advanced options", open=False):
                     hf_token = gr.Textbox(
-                        label="Hugging Face Token",
-                        placeholder="Used only for this reset; not persisted in state",
+                        label="Hugging Face Token (required for llm_query / rlm_query)",
+                        placeholder="hf_...  — used only for this reset; not persisted",
                         type="password",
                     )
+                    expected_answer = gr.Textbox(
+                        label="Expected answer (activates reward scoring)",
+                        placeholder="e.g. BANANA-747  —  leave blank to skip",
+                    )
                     llm_model = gr.Textbox(
                         label="LLM Model",
-                        placeholder="Optional override for llm_query / rlm_query",
+                        placeholder="e.g. Qwen/Qwen2.5-7B-Instruct",
+                        info=(
+                            "Blank = server default (Qwen/Qwen3.5-9B, tuned for the "
+                            "`enable_thinking=False` chat template)."
+                        ),
                     )
+                with gr.Row():
+                    reset_btn = gr.Button("🔁 Reset episode", variant="secondary")
+                    status = gr.Textbox(
+                        label="Status",
+                        interactive=False,
+                        lines=1,
+                        scale=1,
+                    )
+            with gr.Column(scale=3):
                 code = gr.Textbox(
                     label="Python Code",
                     placeholder="count = len(context.split())",
-                    lines=10,
+                    lines=8,
                 )
                 with gr.Row():
-                    reset_btn = gr.Button("Reset", variant="secondary")
-                    run_btn = gr.Button("Run", variant="primary")
-                    state_btn = gr.Button("Get state", variant="secondary")
-                status = gr.Textbox(label="Status", interactive=False)
-            with gr.Column(scale=3):
-                session_view = gr.Markdown(
-                    value="# REPL Session\n\nReset the environment to start."
-                )
-                raw_json = gr.Code(
-                    label="Raw JSON response",
-                    language="json",
-                    interactive=False,
-                )
-                state_json = gr.Code(
-                    label="Session state",
-                    language="json",
-                    interactive=False,
-                )
+                    run_btn = gr.Button("▶ Run", variant="primary")
+                    state_btn = gr.Button("📋 Get state", variant="secondary")
+                session_view = gr.Markdown(value="_Reset and Run to see output._")
+                with gr.Accordion("Raw JSON response", open=False):
+                    raw_json = gr.Code(
+                        label="",
+                        language="json",
+                        interactive=False,
+                    )
+                with gr.Accordion("Session state", open=False):
+                    state_json = gr.Code(
+                        label="",
+                        language="json",
+                        interactive=False,
+                    )
+
+        with gr.Accordion("REPL helpers reference & typical flow", open=False):
+            gr.Markdown(_HELPERS_MD)
+
+        if quick_start_md:
+            with gr.Accordion("API quick start (Python client snippet)", open=False):
+                gr.Markdown(quick_start_md)
+
+        if readme_content:
+            with gr.Accordion("Full README", open=False):
+                gr.Markdown(readme_content)
 
         reset_btn.click(
             fn=reset_repl,
-            inputs=[context, task_prompt, hf_token, llm_model],
+            inputs=[context, task_prompt, expected_answer, hf_token, llm_model],
             outputs=[session_view, raw_json, state_json, status],
         )
         run_btn.click(
@@ -191,5 +428,21 @@ def build_repl_gradio_app(
             outputs=[session_view, raw_json, state_json, status],
         )
         state_btn.click(fn=get_state_sync, outputs=[state_json])
+        example_dropdown.change(
+            fn=load_example,
+            inputs=[example_dropdown],
+            outputs=[
+                context,
+                task_prompt,
+                code,
+                expected_answer,
+                context_accordion,
+            ],
+        )
+        upload.change(
+            fn=load_uploaded_document,
+            inputs=[upload],
+            outputs=[context, status, context_accordion],
+        )
 
     return blocks

server/repl_environment.py

@@ -33,7 +33,12 @@ try:
     from ..models import CodeBlockResult, REPLAction, REPLObservation, REPLState
 except ImportError:
     try:
-        from repl_env.models import CodeBlockResult, REPLAction, REPLObservation, REPLState
+        from repl_env.models import (
+            CodeBlockResult,
+            REPLAction,
+            REPLObservation,
+            REPLState,
+        )
     except ImportError:
         from models import CodeBlockResult, REPLAction, REPLObservation, REPLState
 
@@ -47,10 +52,11 @@ except ImportError:
         from repl_env.rubrics import REPLRubric
         from .python_executor import PythonExecutor
     except ImportError:
-        from .python_executor import PythonExecutor
         from recursive_controller import create_server_recursive_controller
         from rubrics import REPLRubric
 
+        from .python_executor import PythonExecutor
+
 
 class REPLEnvironment(Environment):
     """
@@ -138,6 +144,12 @@ class REPLEnvironment(Environment):
         self._executor: Optional[PythonExecutor] = None
         self._runtime_controller = None
         self._runtime_controller_chat_fn: Optional[Callable[..., str]] = None
+        self._current_llm_model: Optional[str] = None
+
+    @staticmethod
+    def _resolve_model(llm_model: Optional[str]) -> str:
+        """Resolve the effective model name: explicit arg > env var > hard default."""
+        return llm_model or os.environ.get("LLM_MODEL", "Qwen/Qwen3.5-9B")
 
     @staticmethod
     def _build_hf_chat_fn(
@@ -149,7 +161,7 @@ class REPLEnvironment(Environment):
         except ImportError:
             raise RuntimeError("huggingface_hub is required for HF-backed recursion")
 
-        default_model = llm_model or os.environ.get("LLM_MODEL", "Qwen/Qwen3.5-9B")
+        default_model = REPLEnvironment._resolve_model(llm_model)
         client = InferenceClient(model=default_model, token=hf_token, timeout=300)
 
         def chat_fn(messages: list[dict[str, str]], model: str | None = None) -> str:
@@ -255,13 +267,23 @@ class REPLEnvironment(Environment):
         self.rlm_max_depth = runtime_rlm_max_depth
         self.rlm_max_iterations = runtime_rlm_max_iterations
 
-        # Create or rebuild LLM functions when needed.
         # Token resolution: explicit hf_token > HF_TOKEN env var > cached HF login.
-        if not self.llm_query_fn:
+        # Compare resolved model names so reset(llm_model=<default>) and
+        # reset() are treated as equal and don't trigger a redundant rebuild.
+        resolved_model = self._resolve_model(llm_model)
+        has_runtime_llm = self._runtime_controller is not None
+        model_changed = (
+            has_runtime_llm and resolved_model != self._current_llm_model
+        )
+        token_provided = hf_token is not None
+        if not self.llm_query_fn or model_changed or token_provided:
             effective_token = (
                 hf_token if hf_token is not None else os.environ.get("HF_TOKEN")
             )
+            # Pass raw llm_model; _build_hf_chat_fn resolves it via
+            # _resolve_model so the client matches the cache key below.
             self._create_llm_functions(effective_token, llm_model)
+            self._current_llm_model = resolved_model
         elif depth_changed and self._runtime_controller is not None:
             # Rebuild controller with new depth/iteration config but reuse
             # the existing chat_fn — don't require re-providing credentials.
@@ -514,10 +536,18 @@ class REPLEnvironment(Environment):
         Returns:
             Final answer string or None if not found
         """
-        # Pattern 1: RLM-style FINAL(answer)
-        final_match = re.search(r"FINAL\((.*?)\)", stdout, re.DOTALL)
-        if final_match:
-            return final_match.group(1).strip()
+        # Pattern 1: RLM-style FINAL(answer). Paren-counting handles nested
+        # parens (e.g. FINAL(f(x))) and multi-line values without regex flag trade-offs.
+        idx = stdout.find("FINAL(")
+        if idx != -1:
+            depth, start = 0, idx + len("FINAL")
+            for i, ch in enumerate(stdout[idx + len("FINAL"):], start=idx + len("FINAL")):
+                if ch == "(":
+                    depth += 1
+                elif ch == ")":
+                    depth -= 1
+                    if depth == 0:
+                        return stdout[start + 1 : i].strip()
 
         # Pattern 2: RLM-style FINAL_VAR(variable_name)
         final_var_match = re.search(r"FINAL_VAR\((\w+)\)", stdout)

src/__init__.py

@@ -1,7 +1,3 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """EnvTorch: Standardized agentic execution environments."""

src/core/README.md

@@ -1,21 +1,15 @@
-# <img width="35" height="35" alt="image" src="https://github.com/user-attachments/assets/2700a971-e5d6-4036-b03f-2f89c9791609" /> OpenEnv: Agentic Execution Environments
+# OpenEnv: Agentic Execution Environments
 
 An e2e framework for creating, deploying and using isolated execution environments for agentic RL training, built using Gymnasium style simple APIs. OpenEnv provides a standard for interacting with agentic execution environments via simple Gymnasium style APIs - step(), reset(), state(). Users of agentic execution environments can interact with the environment during RL training loops using these simple APIs.
 
 In addition to making it easier for researchers and RL framework writers, we also provide tools for environment creators making it easier for them to create richer environments and make them available over familiar protocols like HTTP and packaged using canonical technologies like docker. Environment creators can use the OpenEnv framework to create environments that are isolated, secure, and easy to deploy and use.
 
+> [!NOTE]
+> OpenEnv is in early development, see the note in the [project README](../../../README.md) before contributing. The [README](../../../README.md#community-support--acknowledgments) lists the technical committee that governs the project, and the [charter](../../../GOVERNANCE.md) explains how it works.
 
 ## Overview
 `openenv.core` provides the foundational building blocks for creating and interacting with containerized environments over HTTP. It enables you to build agent environments that can be deployed as Docker containers and accessed via a simple HTTP API.
 
-> ⚠️ **Early Development Warning** OpenEnv is currently in an experimental
-> stage. You should expect bugs, incomplete features, and APIs that may change
-> in future versions. The project welcomes bugfixes, but to make sure things are
-> well coordinated you should discuss any significant change before starting the
-> work. It's recommended that you signal your intention to contribute in the
-> issue tracker, either by filing a new issue or by claiming an existing one.
-
-
 # OpenEnv Core
 
 Core components for OpenEnv - a framework for building HTTP-based agentic environments.
@@ -28,16 +22,42 @@ Core components for OpenEnv - a framework for building HTTP-based agentic enviro
 - **Container Providers**: Pluggable architecture for running containers (Docker, Kubernetes, etc.)
 - **Type System**: Strongly-typed Action/Observation/State interfaces
 - **Web Interface**: Optional web UI for interacting with environments
+- **Experimental Harness Helpers**: `openenv.core.harness` provides
+  `ResourceSessionFactory`, `StepEnvSessionAdapter`, and `HarnessAdapter` for
+  MCP-first training/evaluation harnesses
+
+## Experimental Harness Helpers
+
+OpenEnv now includes an additive harness-facing layer for cases where a rollout
+driver should interact with environment resources through tools instead of
+calling `reset()` / `step()` directly. These helpers are currently
+experimental while RFC 005 remains in review, so import them from
+`openenv.core.harness`.
+
+Core pieces:
+
+- `ResourceSessionFactory.create(task, seed, episode_id)` creates an isolated
+  per-rollout session
+- `StepEnvSessionAdapter` wraps an existing client and exposes session tools
+- `SessionMCPBridge` exposes a session through MCP-style JSON-RPC
+- `MCPHarnessAdapter` drives white-box rollouts where the trainer still owns
+  model sampling, token IDs, and logprobs
+- `CLIHarnessAdapter` supports opaque evaluation harnesses over the same session
+  surface
+- `build_harness_rollout_func(...)` builds a TRL-compatible `rollout_func`
+
+This layer is additive. Existing environment clients and MCP environments keep
+their current control-plane APIs unchanged.
 
 ## Installation
 
 ```bash
-pip install "openenv[core]"
+pip install "openenv"
 ```
 
 For development:
 ```bash
-pip install "openenv[core]"
+pip install "openenv"
 ```
 
 ## Quick Start
@@ -52,14 +72,17 @@ from openenv.core import EnvClient, StepResult
 from dataclasses import dataclass
 from typing import Any
 
+
 @dataclass
 class MyAction:
     text: str
 
+
 @dataclass
 class MyObservation:
     response: str
 
+
 class MyEnvClient(EnvClient[MyAction, MyObservation, Any]):
     def _step_payload(self, action: MyAction) -> dict:
         return {"text": action.text}
@@ -69,12 +92,13 @@ class MyEnvClient(EnvClient[MyAction, MyObservation, Any]):
         return StepResult(
             observation=MyObservation(**obs_data),
             reward=payload.get("reward"),
-            done=payload.get("done", False)
+            done=payload.get("done", False),
         )
 
     def _parse_state(self, payload: dict) -> Any:
         return payload
 
+
 # Async usage (recommended)
 async def main():
     client = await MyEnvClient.from_docker_image("my-env:latest")
@@ -82,6 +106,7 @@ async def main():
         result = await client.reset()
         step_result = await client.step(MyAction(text="hello"))
 
+
 asyncio.run(main())
 
 # Sync usage (via .sync() wrapper)
@@ -96,26 +121,26 @@ with MyEnvClient(base_url="http://localhost:8000").sync() as client:
 from openenv.core.env_server import Environment, HTTPEnvServer, create_app
 from dataclasses import dataclass
 
+
 @dataclass
 class MyAction:
     text: str
 
+
 @dataclass
 class MyObservation:
     response: str
     reward: float = 0.0
     done: bool = False
 
+
 class MyEnvironment(Environment):
     def reset(self) -> MyObservation:
         return MyObservation(response="Ready")
 
     def step(self, action: MyAction) -> MyObservation:
-        return MyObservation(
-            response=f"Echo: {action.text}",
-            reward=1.0,
-            done=False
-        )
+        return MyObservation(response=f"Echo: {action.text}", reward=1.0, done=False)
+
 
 # Create FastAPI app
 env = MyEnvironment()
@@ -199,7 +224,7 @@ Base interface for environment implementations:
 
 ## License
 
-This project is licensed under the BSD-3-Clause License - see the LICENSE file for details.
+This project is licensed under the BSD 3-Clause License - see the LICENSE file for details.
 
 ## Contributing
 
@@ -207,6 +232,6 @@ Contributions are welcome! Please see the main OpenEnv repository for contributi
 
 ## Links
 
-- **Homepage**: https://github.com/meta-pytorch/OpenEnv
-- **Documentation**: https://github.com/meta-pytorch/OpenEnv/blob/main/README.md
-- **Bug Tracker**: https://github.com/meta-pytorch/OpenEnv/issues
+- **Homepage**: https://github.com/huggingface/OpenEnv
+- **Documentation**: https://github.com/huggingface/OpenEnv/blob/main/README.md
+- **Bug Tracker**: https://github.com/huggingface/OpenEnv/issues

src/core/__init__.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """Core components for agentic environments."""
 

src/core/client_types.py

@@ -1,6 +1,6 @@
 # Type definitions for EnvTorch
 from dataclasses import dataclass
-from typing import Generic, Optional, TypeVar
+from typing import Any, Dict, Generic, Optional, TypeVar
 
 # Generic type for observations
 ObsT = TypeVar("ObsT")
@@ -13,11 +13,17 @@ class StepResult(Generic[ObsT]):
     Represents the result of one environment step.
 
     Attributes:
-        observation: The environment's observation after the action.
-        reward: Scalar reward for this step (optional).
-        done: Whether the episode is finished.
+        observation:
+            The environment's observation after the action.
+        reward (`float`, *optional*):
+            Scalar reward for this step.
+        done (`bool`, *optional*, defaults to `False`):
+            Whether the episode is finished.
+        metadata (`dict`, *optional*):
+            Additional metadata returned alongside the observation.
     """
 
     observation: ObsT
     reward: Optional[float] = None
     done: bool = False
+    metadata: Optional[Dict[str, Any]] = None

src/core/containers/__init__.py

@@ -1,7 +1,3 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """Container management for environment servers."""

src/core/containers/images/Dockerfile

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 #
 # OpenEnv Base Image

src/core/containers/runtime/__init__.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """Container runtime providers."""
 
@@ -15,6 +11,11 @@ from .providers import (
 )
 from .uv_provider import UVProvider
 
+# Note: optional cloud providers that require extra SDKs (e.g.
+# `ACASandboxProvider`, `DaytonaProvider`) are intentionally NOT re-exported
+# here. Import them from their module, e.g.
+# `from openenv.core.containers.runtime.aca_provider import ACASandboxProvider`.
+
 __all__ = [
     "ContainerProvider",
     "DockerSwarmProvider",

src/core/containers/runtime/_server_config.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Shared helpers for reading environment server configuration."""
+
+from __future__ import annotations
+
+import json
+import re
+
+import yaml
+
+
+def parse_openenv_app_field(yaml_content: str) -> str | None:
+    """Extract the top-level ``app`` value from raw ``openenv.yaml`` content."""
+    try:
+        data = yaml.safe_load(yaml_content) or {}
+    except Exception:
+        return None
+
+    if not isinstance(data, dict):
+        return None
+
+    value = data.get("app")
+    if isinstance(value, str):
+        value = value.strip()
+        return value if value else None
+    return None
+
+
+def parse_dockerfile_cmd(dockerfile_content: str) -> str | None:
+    """Extract the server command from the last Dockerfile ``CMD``."""
+    last_cmd: str | None = None
+    for line in dockerfile_content.splitlines():
+        stripped = line.strip()
+        if stripped.startswith("#"):
+            continue
+        match = re.match(r"CMD\s+(.+)", stripped, flags=re.IGNORECASE)
+        if match:
+            last_cmd = match.group(1).strip()
+
+    if last_cmd is None:
+        return None
+
+    if last_cmd.startswith("["):
+        try:
+            parts = json.loads(last_cmd)
+            if isinstance(parts, list) and all(isinstance(p, str) for p in parts):
+                return " ".join(parts)
+        except (json.JSONDecodeError, TypeError):
+            # Invalid exec-form JSON falls back to the raw CMD below.
+            pass
+
+    return last_cmd if last_cmd else None

src/core/containers/runtime/aca_provider.py

@@ -0,0 +1,690 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""Azure Container Apps Sandbox provider for OpenEnv environments.
+
+Note: the ``RFC 002 security invariant S<n>`` references in this module point to
+the **proposed** "Cloud Sandbox Providers" amendment to RFC 002, which is pending
+review/sign-off by the RFC authors. They describe the security properties this
+provider enforces; treat them as proposed (not yet ratified) until that sign-off.
+"""
+
+from __future__ import annotations
+
+import os
+import shlex
+import time
+import warnings
+from typing import Any, Dict, Literal, Optional
+
+from .providers import ContainerProvider
+
+
+_DEFAULT_ACA_PORT = 8000
+# 15 minutes matches the conservative auto-stop default used by other cloud
+# sandbox providers (e.g. Daytona). A live RL rollout holds one WebSocket for
+# the whole episode, so an aggressive auto-suspend can silently drop the
+# session between slow steps. See RFC 002 "Cloud Sandbox Providers".
+_DEFAULT_AUTO_SUSPEND_SECONDS = 900
+_VALID_SUSPEND_MODES = {"Memory", "Disk"}
+
+
+def _require_secure_url(url: str) -> str:
+    """Enforce https/wss transport (RFC 002 security invariant S1).
+
+    `EnvClient` derives its WebSocket URL from this base URL, so a plaintext
+    URL would become a cleartext `ws://` connection. The offending URL is
+    deliberately omitted from the error because an anonymous ACA port URL is a
+    bearer capability (S2) that must not leak into logs.
+    """
+    if not isinstance(url, str) or not url.lower().startswith("https://"):
+        raise RuntimeError(
+            "ACA sandbox returned a non-HTTPS port URL. OpenEnv requires an "
+            "https/wss base_url so EnvClient traffic is encrypted. Refusing to "
+            "connect over plaintext."
+        )
+    return url
+
+
+def _validate_anonymous_port(anonymous_port: Optional[bool]) -> bool:
+    """Require explicit, public-only ingress opt-in (RFC 002 invariant S2).
+
+    Public exposure is never implicit (`None` is rejected), and authenticated
+    ACA ports are rejected because `EnvClient` has no WebSocket auth mechanism
+    yet.
+    """
+    if anonymous_port is None:
+        raise ValueError(
+            "ACASandboxProvider requires explicit anonymous_port=True. "
+            "ACA anonymous ports are public ingress; authenticated ACA "
+            "ports are not supported by EnvClient until a client-auth RFC "
+            "defines WebSocket credentials."
+        )
+    if not anonymous_port:
+        raise ValueError(
+            "authenticated ACA sandbox ports are not supported by EnvClient "
+            "yet. Use anonymous_port=True only for intentionally exposed "
+            "sandbox URLs, or wait for a client-auth RFC."
+        )
+    return True
+
+
+def _resolve_required(value: str | None, env_name: str, parameter: str) -> str:
+    resolved = value or os.environ.get(env_name)
+    if not resolved:
+        raise ValueError(
+            f"{parameter} is required. Pass it to ACASandboxProvider or set "
+            f"the {env_name} environment variable."
+        )
+    return resolved
+
+
+def _raise_install_error(exc: ImportError) -> None:
+    raise RuntimeError(
+        "Azure Container Apps sandbox support requires optional dependencies. "
+        "Install them with `pip install openenv[aca]`."
+    ) from exc
+
+
+class _DefaultACASandboxAdapter:
+    """Thin adapter over the preview `azure-containerapps-sandbox` SDK.
+
+    The provider talks to this private adapter instead of spreading SDK details
+    through its own logic; tests inject a duck-typed fake in its place.
+    """
+
+    def __init__(
+        self,
+        *,
+        subscription_id: str,
+        resource_group: str,
+        sandbox_group: str,
+        region: str | None,
+        endpoint: str | None,
+        credential: Any,
+        sdk_kwargs: dict[str, Any],
+    ):
+        try:
+            from azure.containerapps.sandbox import (
+                endpoint_for_region,
+                SandboxGroupClient,
+            )
+        except ImportError as exc:  # pragma: no cover - exercised via provider tests
+            _raise_install_error(exc)
+
+        if credential is None:
+            try:
+                from azure.identity import DefaultAzureCredential
+            except (
+                ImportError
+            ) as exc:  # pragma: no cover - exercised via provider tests
+                _raise_install_error(exc)
+
+            credential = DefaultAzureCredential()
+
+        if endpoint is None:
+            if region is None:
+                raise ValueError(
+                    "region is required when endpoint is not provided. Pass it to "
+                    "ACASandboxProvider or set AZURE_REGION."
+                )
+            endpoint = endpoint_for_region(region)
+
+        self._client = SandboxGroupClient(
+            endpoint,
+            credential,
+            subscription_id=subscription_id,
+            resource_group=resource_group,
+            sandbox_group=sandbox_group,
+            **sdk_kwargs,
+        )
+
+    def create_sandbox(
+        self,
+        *,
+        disk: str | None,
+        disk_id: str | None,
+        env_vars: dict[str, str] | None,
+        labels: dict[str, str] | None,
+        egress_policy: Any,
+        cpu: str,
+        memory: str,
+        disk_size: str | None,
+        auto_suspend_seconds: int,
+        auto_suspend_mode: str,
+        create_timeout: int,
+        polling_interval: int,
+    ) -> Any:
+        kwargs: dict[str, Any] = {
+            "disk": disk,
+            "disk_id": disk_id,
+            "cpu": cpu,
+            "memory": memory,
+            "disk_size": disk_size,
+            "auto_suspend_seconds": auto_suspend_seconds,
+            "auto_suspend_mode": auto_suspend_mode,
+            "labels": labels,
+            "environment": env_vars,
+            "egress_policy": egress_policy,
+            "polling_timeout": create_timeout,
+            "polling_interval": polling_interval,
+        }
+        kwargs = {key: value for key, value in kwargs.items() if value is not None}
+        # `begin_create_sandbox(...).result()` blocks until the sandbox is
+        # *Running* and returns a sandbox-scoped `SandboxClient` (not a plain
+        # data model). That client exposes exec / add_port / begin_delete
+        # directly, which is the handle the rest of this adapter operates on.
+        return self._client.begin_create_sandbox(**kwargs).result()
+
+    def add_port(self, sandbox: Any, *, port: int, anonymous: bool) -> str:
+        sandbox_port = sandbox.add_port(port, anonymous=anonymous)
+        url = getattr(sandbox_port, "url", None)
+        if not url and isinstance(sandbox_port, dict):
+            url = sandbox_port.get("url")
+        if not url:
+            raise RuntimeError(
+                "ACA sandbox did not return a URL for the exposed port. "
+                "OpenEnv requires a base_url that EnvClient can connect to directly."
+            )
+        return str(url)
+
+    def exec(
+        self,
+        sandbox: Any,
+        command: str,
+        *,
+        working_directory: str | None = None,
+    ) -> Any:
+        return sandbox.exec(command, working_directory=working_directory)
+
+    def delete_sandbox(self, sandbox: Any) -> None:
+        if hasattr(sandbox, "begin_delete"):
+            sandbox.begin_delete().result()
+        else:
+            sandbox.delete()
+
+    def close(self) -> None:
+        self._client.close()
+
+
+def _source_from_image(image: str) -> tuple[str | None, str | None]:
+    if not image or not image.strip():
+        raise ValueError("image must be a non-empty ACA disk or disk-id reference")
+    image = image.strip()
+
+    if image.startswith("disk:"):
+        disk = image[len("disk:") :].strip()
+        if not disk:
+            raise ValueError("disk: source must include a disk image name")
+        return disk, None
+    if image.startswith("disk-id:"):
+        disk_id = image[len("disk-id:") :].strip()
+        if not disk_id:
+            raise ValueError("disk-id: source must include a disk image id")
+        return None, disk_id
+
+    # A bare string is treated as a public disk-image name. Guard the common
+    # migration foot-gun: an ACA sandbox source is NOT a Docker/OCI registry
+    # image, so a value that looks like one (a registry path with "/", or an
+    # "image:tag") is almost certainly a mistake and would otherwise surface as
+    # an opaque ACA SDK error. Fail fast with guidance instead.
+    if "/" in image or ":" in image:
+        raise ValueError(
+            f"{image!r} looks like a container/OCI image reference, but "
+            "ACASandboxProvider sources are Azure Container Apps sandbox disks, "
+            "not Docker registry images. Use a bare public disk name, "
+            "'disk:<name>' for a public disk image, or 'disk-id:<id>' for a "
+            "private one (building a disk image from a container image is a "
+            "separate, out-of-band ACA step)."
+        )
+
+    return image, None
+
+
+def _exec_stdout(result: Any) -> str:
+    for attr in ("stdout", "result", "output"):
+        value = getattr(result, attr, None)
+        if value is not None:
+            return str(value)
+    if isinstance(result, dict):
+        for key in ("stdout", "result", "output"):
+            value = result.get(key)
+            if value is not None:
+                return str(value)
+    return str(result)
+
+
+class ACASandboxProvider(ContainerProvider):
+    """Container provider backed by Azure Container Apps Sandboxes.
+
+    `start_container`'s `image` is an ACA sandbox source (`disk:<name>` for a
+    public disk image or `disk-id:<id>` for a private one), not a Docker
+    registry image. The provider boots an OpenEnv server inside the sandbox,
+    exposes it on an anonymous ACA port, and returns an `https://` URL that
+    `EnvClient` connects to over `wss://`.
+
+    The environment runs untrusted code, so the provider is secure by default
+    (proposed RFC 002 security invariants): it requires explicit `anonymous_port=True`
+    ingress (S2), enforces https/wss transport (S1), offers
+    `deny_all_egress()` to block the cloud metadata/IMDS endpoint (S3), and
+    never surfaces raw sandbox output unless `surface_server_logs=True` (S4).
+
+    Validated end to end against a live ACA sandbox group; the underlying
+    `azure-containerapps-sandbox` SDK is preview, so pin it and re-validate
+    after upgrades (see `tests/test_core/test_aca_provider_integration.py`).
+
+    `close()` stops the active sandbox *and* releases the underlying SDK client,
+    so prefer using the provider as a context manager (or call `close()`
+    explicitly) for deterministic cleanup. `close()`/`__enter__`/`__exit__` are
+    defined on `ContainerProvider` (default no-op), so a caller holding a bare
+    `ContainerProvider` reference can release the client polymorphically.
+
+    Only one sandbox is active per provider: calling `start_container` again
+    before `stop_container()`/`close()` raises `RuntimeError` rather than
+    orphaning the running sandbox.
+
+    ```python
+    with ACASandboxProvider(
+        image="disk:my-env",
+        anonymous_port=True,
+        ...,
+    ) as provider:
+        base_url = provider.start_container(cmd=...)
+        ...
+    # sandbox deleted and SDK client closed on exit
+    ```
+    """
+
+    def __init__(
+        self,
+        *,
+        subscription_id: Optional[str] = None,
+        resource_group: Optional[str] = None,
+        sandbox_group: Optional[str] = None,
+        region: Optional[str] = None,
+        endpoint: Optional[str] = None,
+        credential: Any = None,
+        cpu: str = "1000m",
+        memory: str = "2048Mi",
+        disk_size: Optional[str] = None,
+        auto_suspend_seconds: int = _DEFAULT_AUTO_SUSPEND_SECONDS,
+        auto_suspend_mode: Literal["Memory", "Disk"] = "Memory",
+        labels: Optional[Dict[str, str]] = None,
+        egress_policy: Any = None,
+        anonymous_port: Literal[True] | None = None,
+        image: Optional[str] = None,
+        env_vars: Optional[Dict[str, str]] = None,
+        cmd: Optional[str] = None,
+        working_directory: Optional[str] = None,
+        surface_server_logs: bool = False,
+        create_timeout: int = 300,
+        polling_interval: int = 3,
+        sdk_kwargs: Optional[dict[str, Any]] = None,
+        _adapter: Any = None,
+    ):
+        if auto_suspend_mode not in _VALID_SUSPEND_MODES:
+            raise ValueError("auto_suspend_mode must be 'Memory' or 'Disk'")
+        _validate_anonymous_port(anonymous_port)
+
+        self.subscription_id = subscription_id or os.environ.get(
+            "AZURE_SUBSCRIPTION_ID"
+        )
+        self.resource_group = resource_group or os.environ.get("AZURE_RESOURCE_GROUP")
+        self.sandbox_group = sandbox_group or os.environ.get("AZURE_SANDBOX_GROUP")
+        self.region = region or os.environ.get("AZURE_REGION")
+        self.endpoint = endpoint
+
+        self.cpu = cpu
+        self.memory = memory
+        self.disk_size = disk_size
+        self.auto_suspend_seconds = auto_suspend_seconds
+        self.auto_suspend_mode = auto_suspend_mode
+        self.labels = dict(labels or {})
+        self.egress_policy = egress_policy
+        self.anonymous_port = anonymous_port
+        self._image = image
+        self._env_vars = dict(env_vars or {}) if env_vars is not None else None
+        self.cmd = cmd
+        self.working_directory = working_directory
+        self.surface_server_logs = surface_server_logs
+        self.create_timeout = create_timeout
+        self.polling_interval = polling_interval
+
+        if _adapter is None:
+            self._adapter = _DefaultACASandboxAdapter(
+                subscription_id=_resolve_required(
+                    self.subscription_id, "AZURE_SUBSCRIPTION_ID", "subscription_id"
+                ),
+                resource_group=_resolve_required(
+                    self.resource_group, "AZURE_RESOURCE_GROUP", "resource_group"
+                ),
+                sandbox_group=_resolve_required(
+                    self.sandbox_group, "AZURE_SANDBOX_GROUP", "sandbox_group"
+                ),
+                region=self.region,
+                endpoint=self.endpoint,
+                credential=credential,
+                sdk_kwargs=dict(sdk_kwargs or {}),
+            )
+            self._owns_adapter = True
+        else:
+            self._adapter = _adapter
+            self._owns_adapter = False
+
+        self._sandbox: Any = None
+        self._base_url: Optional[str] = None
+        self._port = _DEFAULT_ACA_PORT
+        self._started_server = False
+        # Injected env-var values, used to scrub captured server output before
+        # it is ever surfaced in an error (RFC 002 security invariant S4).
+        self._redact_values: set[str] = set()
+
+    @staticmethod
+    def deny_all_egress(allow: Optional[list[str]] = None) -> Any:
+        """Build a default-deny ACA `EgressPolicy` (RFC 002 security invariant S3).
+
+        Untrusted RL/agent code should not be able to exfiltrate data or reach
+        the cloud metadata/IMDS endpoint. This returns an `EgressPolicy` whose
+        default action is `Deny`, with an optional host allowlist (for example a
+        model endpoint or package registry the environment legitimately needs).
+
+        ```python
+        provider = ACASandboxProvider(
+            anonymous_port=True,
+            egress_policy=ACASandboxProvider.deny_all_egress(
+                allow=["my-model.openai.azure.com"]
+            ),
+        )
+        ```
+        """
+        try:
+            from azure.containerapps.sandbox import EgressHostRule, EgressPolicy
+        except ImportError as exc:  # pragma: no cover - exercised via provider tests
+            _raise_install_error(exc)
+
+        host_rules = [
+            EgressHostRule(pattern=pattern, action="Allow") for pattern in (allow or [])
+        ]
+        return EgressPolicy(default_action="Deny", host_rules=host_rules)
+
+    def start_container(
+        self,
+        image: Optional[str] = None,
+        port: Optional[int] = None,
+        env_vars: Optional[Dict[str, str]] = None,
+        **kwargs: Any,
+    ) -> str:
+        """Start an OpenEnv server in an ACA sandbox and return its base URL.
+
+        `image` is an ACA sandbox *source*, not a Docker/OCI registry image: a
+        bare string or `disk:<name>` for a public disk image, or `disk-id:<id>`
+        for a private one. It may be omitted when supplied to the constructor.
+        A value that looks like a container image (a registry path or
+        `name:tag`) is rejected with guidance. Only port 8000 is supported.
+        `**kwargs` accepts per-start overrides (`cmd`, `labels`,
+        `egress_policy`, `anonymous_port`); unknown options raise `ValueError`
+        so typos cannot silently change sandbox behavior.
+        """
+        if self._sandbox is not None:
+            raise RuntimeError(
+                "ACASandboxProvider already has an active sandbox. Call "
+                "stop_container() (or close()) before starting another — a "
+                "second start would orphan the running sandbox."
+            )
+
+        bind_port = port or _DEFAULT_ACA_PORT
+        if bind_port != _DEFAULT_ACA_PORT:
+            raise ValueError(
+                f"ACASandboxProvider only supports port {_DEFAULT_ACA_PORT} "
+                f"(got {bind_port})."
+            )
+
+        effective_image = image if image is not None else self._image
+        if effective_image is None:
+            raise ValueError(
+                "ACASandboxProvider requires an image. Pass it to the constructor "
+                "or start_container()."
+            )
+        effective_env_vars = self._env_vars if env_vars is None else env_vars
+
+        disk, disk_id = _source_from_image(effective_image)
+        labels = dict(self.labels)
+        labels.update(kwargs.pop("labels", {}) or {})
+        cmd = kwargs.pop("cmd", None) or self.cmd
+        egress_policy = kwargs.pop("egress_policy", self.egress_policy)
+        anonymous_port = kwargs.pop("anonymous_port", self.anonymous_port)
+        # Re-validate in case anonymous_port was overridden at start time: public
+        # exposure must stay an explicit, opt-in decision (S2).
+        _validate_anonymous_port(anonymous_port)
+
+        if kwargs:
+            unknown = ", ".join(sorted(kwargs))
+            raise ValueError(f"Unsupported ACASandboxProvider start options: {unknown}")
+
+        # Security (RFC 002 S3): the env server runs untrusted code and ACA's
+        # EgressPolicy defaults to default_action="Allow", so an unset policy
+        # means unrestricted egress — the workload could exfiltrate data or
+        # steal the sandbox managed-identity token from the cloud metadata
+        # endpoint. Warn loudly; do not silently allow it.
+        if egress_policy is None:
+            warnings.warn(
+                "ACASandboxProvider is starting a sandbox with unrestricted "
+                "egress (no egress_policy). The environment runs untrusted code; "
+                "prefer a default-deny policy via "
+                "ACASandboxProvider.deny_all_egress(allow=[...]) and ensure the "
+                "cloud metadata/IMDS endpoint is blocked.",
+                stacklevel=2,
+            )
+
+        # Record injected secret values so captured server output can be scrubbed
+        # before it is ever surfaced in an error (security invariant S4).
+        self._redact_values = {
+            value for value in (effective_env_vars or {}).values() if value
+        }
+
+        # Two-stage cleanup. A create failure created nothing, so just drop the
+        # recorded secrets and re-raise (notably it must NOT delete a sandbox
+        # this call did not create — e.g. one left over on the provider). Only a
+        # failure AFTER the sandbox exists runs stop_container() to delete it.
+        try:
+            self._sandbox = self._adapter.create_sandbox(
+                disk=disk,
+                disk_id=disk_id,
+                env_vars=dict(effective_env_vars or {}) if effective_env_vars else None,
+                labels=labels or None,
+                egress_policy=egress_policy,
+                cpu=self.cpu,
+                memory=self.memory,
+                disk_size=self.disk_size,
+                auto_suspend_seconds=self.auto_suspend_seconds,
+                auto_suspend_mode=self.auto_suspend_mode,
+                create_timeout=self.create_timeout,
+                polling_interval=self.polling_interval,
+            )
+        except Exception:
+            self._redact_values = set()
+            raise
+
+        try:
+            if cmd:
+                self._start_server(cmd)
+            self._base_url = _require_secure_url(
+                self._adapter.add_port(
+                    self._sandbox,
+                    port=bind_port,
+                    anonymous=anonymous_port,
+                )
+            )
+            self._port = bind_port
+        except Exception:
+            # A cleanup failure here must not mask the original error: swallow
+            # any exception from stop_container() so the root cause propagates.
+            try:
+                self.stop_container()
+            except Exception:
+                pass
+            raise
+
+        return self._base_url
+
+    def _start_server(self, cmd: str) -> None:
+        # `cmd` is trusted orchestrator configuration, not agent/environment
+        # input — callers must not pass agent-controlled text as `cmd`, and any
+        # dynamic value interpolated into it must be quoted by the caller. Given
+        # that, running it through the subshell below is safe (RFC 002 amendment
+        # S5 — no command injection from untrusted sources). `working_directory`
+        # is still shlex-quoted defensively.
+        if self.working_directory:
+            command = f"cd {shlex.quote(self.working_directory)} && {cmd}"
+        else:
+            command = cmd
+        escaped = shlex.quote(command)
+        # The adapter's `exec` maps to the ACA SDK's `executeShellCommand`
+        # endpoint, which runs its argument through a POSIX shell in the sandbox
+        # (the SDK docstring says "Execute a shell command" and tells callers to
+        # `shlex.quote()` interpolated input). That shell is what makes the
+        # backgrounding (`&`), output redirection, and `$!` PID capture below
+        # work; if a future preview SDK changes `exec` to bypass the shell, this
+        # start sequence must be revisited (re-validate via the integration test).
+        self._adapter.exec(
+            self._sandbox,
+            f"nohup bash -c {escaped} > /tmp/openenv-server.log 2>&1 & "
+            "echo $! > /tmp/openenv-server.pid",
+        )
+        self._started_server = True
+
+    def _redact(self, text: str, *, max_chars: int = 2000) -> str:
+        """Scrub injected secret values and bound length before surfacing output.
+
+        Replaces any injected env-var value with `***` and keeps only the tail.
+        This is best-effort (exact-match only), which is why server output is
+        withheld entirely unless `surface_server_logs=True` (RFC 002 S4).
+        """
+        redacted = text or ""
+        for value in self._redact_values:
+            redacted = redacted.replace(value, "***")
+        if len(redacted) > max_chars:
+            redacted = "...(truncated)...\n" + redacted[-max_chars:]
+        return redacted
+
+    def _server_died_message(self) -> str:
+        """Build the startup-crash error, secure by default (RFC 002 S4).
+
+        Untrusted code can print secrets then force a crash to exfiltrate them
+        through the exception (which lands in orchestrator/CI logs), so sandbox
+        output is excluded unless `surface_server_logs=True`, in which case a
+        best-effort redacted, bounded excerpt is included.
+        """
+        base = (
+            "ACA sandbox server process died during startup. Server output is "
+            "not surfaced to avoid leaking secrets injected into the sandbox; "
+            "retrieve /tmp/openenv-server.log from the sandbox out of band, or "
+            "construct the provider with surface_server_logs=True to include a "
+            "redacted excerpt."
+        )
+        if not self.surface_server_logs:
+            return base
+        if self._sandbox is None:
+            return base
+        log = self._redact(
+            _exec_stdout(
+                self._adapter.exec(
+                    self._sandbox,
+                    "cat /tmp/openenv-server.log 2>/dev/null",
+                )
+            )
+        )
+        return (
+            "ACA sandbox server process died during startup. The excerpt below "
+            "is the sandbox server output with injected secret values redacted "
+            "(best-effort); it may still contain secrets the workload printed "
+            f"by other means.\nLog (redacted):\n{log}"
+        )
+
+    def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
+        """Wait for the ACA-hosted OpenEnv server to answer on `/health`.
+
+        A `200` on `/health` proves HTTP reachability but not that the exposed
+        ACA port proxies the `/ws` WebSocket upgrade `EnvClient` needs; the
+        integration test covers the full `wss://` round-trip. If the captured
+        server process dies during startup this raises `RuntimeError` (sandbox
+        output is withheld unless `surface_server_logs=True`; see RFC 002 S4).
+        """
+        # Imported lazily inside the method that needs it, matching the other
+        # providers in this package (e.g. DaytonaProvider).
+        import requests
+
+        deadline = time.time() + timeout_s
+        health_url = f"{base_url}/health"
+
+        while time.time() < deadline:
+            try:
+                response = requests.get(health_url, timeout=5.0)
+                if response.status_code == 200:
+                    return
+            except requests.RequestException:
+                pass
+
+            if self._started_server and self._sandbox is not None:
+                status = _exec_stdout(
+                    self._adapter.exec(
+                        self._sandbox,
+                        "kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null "
+                        "&& echo RUNNING || echo DEAD",
+                    )
+                )
+                if "DEAD" in status:
+                    raise RuntimeError(self._server_died_message())
+
+            time.sleep(1.0)
+
+        raise TimeoutError(
+            f"ACA sandbox did not become ready within {timeout_s}s. "
+            "If the sandbox image starts the server itself, pass cmd= to let "
+            "ACASandboxProvider capture /tmp/openenv-server.log for early crash "
+            "diagnostics."
+        )
+
+    def stop_container(self) -> None:
+        """Delete the active ACA sandbox."""
+        if self._sandbox is None:
+            # Still drop any injected secret values recorded by a failed start
+            # (e.g. create_sandbox raised before a sandbox handle existed).
+            self._redact_values = set()
+            return
+
+        try:
+            self._adapter.delete_sandbox(self._sandbox)
+        finally:
+            self._sandbox = None
+            self._base_url = None
+            self._started_server = False
+            # Drop the previous episode's injected secret values; a new start
+            # repopulates them (avoids carrying stale values across sandboxes).
+            self._redact_values = set()
+
+    @property
+    def base_url(self) -> str:
+        """URL returned by the last `start_container`."""
+        if self._base_url is None:
+            raise RuntimeError(
+                "ACASandboxProvider has no active base_url. Start the provider "
+                "before reading base_url."
+            )
+        return self._base_url
+
+    def close(self) -> None:
+        """Stop the active sandbox and close the underlying SDK client.
+
+        Overrides the base no-op so a caller can release the preview SDK client
+        deterministically (also invoked on context-manager exit). The base
+        `ContainerProvider` defines `close()`/`__enter__`/`__exit__`, so callers
+        holding a bare `ContainerProvider` can release resources polymorphically.
+        """
+        self.stop_container()
+        if self._owns_adapter:
+            self._adapter.close()
+            # don't close a second time on a subsequent close()/context exit
+            self._owns_adapter = False
+
+
+__all__ = ["ACASandboxProvider"]

src/core/containers/runtime/daytona_provider.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 Daytona container provider for running OpenEnv environments in Daytona cloud sandboxes.
@@ -12,14 +8,12 @@ Requires the ``daytona`` SDK: ``pip install daytona>=0.10``
 
 from __future__ import annotations
 
-import json
 import os
 import shlex
 import time
 from typing import Any, Callable, Dict, Optional
 
-import yaml
-
+from ._server_config import parse_dockerfile_cmd, parse_openenv_app_field
 from .providers import ContainerProvider
 
 
@@ -40,6 +34,8 @@ class DaytonaProvider(ContainerProvider):
     def __init__(
         self,
         *,
+        image: Optional[str] = None,
+        env_vars: Optional[Dict[str, str]] = None,
         api_key: Optional[str] = None,
         public: bool = False,
         resources: Optional[Any] = None,
@@ -51,6 +47,11 @@ class DaytonaProvider(ContainerProvider):
     ):
         """
         Args:
+            image: Docker image, snapshot reference, or ``"dockerfile:<path>"``
+                source to use when ``start_container()`` is called without an
+                image.
+            env_vars: Environment variables to use when ``start_container()`` is
+                called without explicit ``env_vars``.
             api_key: Daytona API key. Falls back to ``DAYTONA_API_KEY`` env var.
             public: If True, the sandbox preview is publicly accessible.
             resources: Optional ``daytona.Resources`` instance for CPU/memory.
@@ -71,6 +72,8 @@ class DaytonaProvider(ContainerProvider):
             config_kwargs["target"] = target
 
         self._daytona = Daytona(DaytonaConfig(**config_kwargs))
+        self._image = image
+        self._env_vars = env_vars
         self._public = public
         self._resources = resources
         self._auto_stop_interval = auto_stop_interval
@@ -144,19 +147,7 @@ class DaytonaProvider(ContainerProvider):
 
         Uses PyYAML to handle comments, quotes, and nested keys correctly.
         """
-        try:
-            data = yaml.safe_load(yaml_content) or {}
-        except Exception:
-            return None
-
-        if not isinstance(data, dict):
-            return None
-
-        value = data.get("app")
-        if isinstance(value, str):
-            value = value.strip()
-            return value if value else None
-        return None
+        return parse_openenv_app_field(yaml_content)
 
     @staticmethod
     def _parse_dockerfile_cmd(dockerfile_content: str) -> Optional[str]:
@@ -171,32 +162,7 @@ class DaytonaProvider(ContainerProvider):
         Returns:
             The command as a single string, or ``None`` if no ``CMD`` found.
         """
-        import re
-
-        last_cmd: Optional[str] = None
-        for line in dockerfile_content.splitlines():
-            stripped = line.strip()
-            # Skip comments
-            if stripped.startswith("#"):
-                continue
-            match = re.match(r"CMD\s+(.+)", stripped, flags=re.IGNORECASE)
-            if match:
-                last_cmd = match.group(1).strip()
-
-        if last_cmd is None:
-            return None
-
-        # Exec form: CMD ["executable", "param1", ...]
-        if last_cmd.startswith("["):
-            try:
-                parts = json.loads(last_cmd)
-                if isinstance(parts, list) and all(isinstance(p, str) for p in parts):
-                    return " ".join(parts)
-            except (json.JSONDecodeError, TypeError):
-                pass
-
-        # Shell form: CMD executable param1 ...
-        return last_cmd if last_cmd else None
+        return parse_dockerfile_cmd(dockerfile_content)
 
     @staticmethod
     def strip_buildkit_syntax(dockerfile_content: str) -> str:
@@ -349,7 +315,7 @@ class DaytonaProvider(ContainerProvider):
 
     def start_container(
         self,
-        image: str,
+        image: Optional[str] = None,
         port: Optional[int] = None,
         env_vars: Optional[Dict[str, str]] = None,
         **kwargs: Any,
@@ -370,7 +336,8 @@ class DaytonaProvider(ContainerProvider):
             image: Docker image name (e.g. ``"echo-env:latest"``),
                    ``"snapshot:<name>"`` to create from a pre-built snapshot,
                    or ``"dockerfile:<path>"`` returned by
-                   :meth:`image_from_dockerfile`.
+                   :meth:`image_from_dockerfile`. May be omitted when supplied
+                   to the constructor.
             port: Must be ``None`` or ``8000``. Daytona exposes port 8000
                   via its preview proxy; other ports raise ``ValueError``.
             env_vars: Environment variables forwarded to the sandbox.
@@ -386,6 +353,14 @@ class DaytonaProvider(ContainerProvider):
                 "The Daytona preview proxy routes to port 8000 inside the sandbox."
             )
 
+        effective_image = image if image is not None else self._image
+        if effective_image is None:
+            raise ValueError(
+                "DaytonaProvider requires an image. Pass it to the constructor "
+                "or start_container()."
+            )
+        effective_env_vars = self._env_vars if env_vars is None else env_vars
+
         # Resolve the server command (may be None; discovery happens after
         # sandbox creation when we can inspect the filesystem).
         cmd = kwargs.pop("cmd", None) or self._cmd
@@ -395,24 +370,24 @@ class DaytonaProvider(ContainerProvider):
 
         # Build creation params
         create_kwargs: Dict[str, Any] = {}
-        if env_vars:
-            create_kwargs["env_vars"] = env_vars
+        if effective_env_vars:
+            create_kwargs["env_vars"] = effective_env_vars
         if self._public:
             create_kwargs["public"] = True
         if self._auto_stop_interval != 15:
             create_kwargs["auto_stop_interval"] = self._auto_stop_interval
 
-        if image.startswith("snapshot:"):
+        if effective_image.startswith("snapshot:"):
             from daytona import CreateSandboxFromSnapshotParams
 
-            snapshot_name = image[len("snapshot:") :]
+            snapshot_name = effective_image[len("snapshot:") :]
             params = CreateSandboxFromSnapshotParams(
                 snapshot=snapshot_name, **create_kwargs
             )
-        elif image.startswith("dockerfile:"):
+        elif effective_image.startswith("dockerfile:"):
             from daytona import CreateSandboxFromImageParams, Image
 
-            dockerfile_path = image[len("dockerfile:") :]
+            dockerfile_path = effective_image[len("dockerfile:") :]
             meta = self._dockerfile_registry.get(dockerfile_path)
             if meta is None:
                 raise ValueError(
@@ -445,7 +420,7 @@ class DaytonaProvider(ContainerProvider):
         else:
             from daytona import CreateSandboxFromImageParams
 
-            img_kwargs = {"image": image, **create_kwargs}
+            img_kwargs = {"image": effective_image, **create_kwargs}
             if self._resources is not None:
                 img_kwargs["resources"] = self._resources
             params = CreateSandboxFromImageParams(**img_kwargs)

src/core/containers/runtime/hf_sandbox_provider.py

@@ -0,0 +1,342 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+"""Hugging Face-backed provider for OpenEnv environment servers."""
+
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import socket
+import threading
+import time
+from contextlib import suppress
+from typing import Any
+
+import requests
+from fastapi import FastAPI, Request, Response, WebSocket, WebSocketDisconnect
+from websockets.asyncio.client import connect as ws_connect
+from websockets.exceptions import ConnectionClosed
+
+from .providers import ContainerProvider
+
+
+_DEFAULT_PORT = 8000
+_SERVER_COMMAND = (
+    'export SBX_PROXY_DIR="${SBX_PROXY_DIR:-$HOME/.sbx/proxy}"; '
+    'mkdir -p "$SBX_PROXY_DIR"; '
+    'nohup server >"$HOME/openenv-server.log" 2>&1 &'
+)
+_MAX_WS_MESSAGE_SIZE = 100 * 1024 * 1024
+_HOP_BY_HOP_HEADERS = {
+    "connection",
+    "content-encoding",
+    "content-length",
+    "host",
+    "keep-alive",
+    "proxy-authenticate",
+    "proxy-authorization",
+    "te",
+    "trailer",
+    "transfer-encoding",
+    "upgrade",
+}
+_POOLS: dict[tuple[str, str], Any] = {}
+_POOL_LOCK = threading.Lock()
+
+
+def _find_available_port() -> int:
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.bind(("127.0.0.1", 0))
+        sock.listen(1)
+        return sock.getsockname()[1]
+
+
+def _to_ws_url(url: str) -> str:
+    if url.startswith("https://"):
+        return "wss://" + url[len("https://") :]
+    if url.startswith("http://"):
+        return "ws://" + url[len("http://") :]
+    return url
+
+
+def _pool_name(image: str, flavor: str) -> str:
+    digest = hashlib.sha1(f"{image}\0{flavor}".encode("utf-8")).hexdigest()[:12]
+    return f"openenv-{digest}"
+
+
+def _get_sandbox_pool_cls() -> Any:
+    try:
+        from huggingface_hub import SandboxPool
+    except ImportError as exc:
+        raise RuntimeError(
+            "HFSandboxProvider requires a huggingface_hub version with "
+            "SandboxPool.create and Sandbox.proxy_url_for support."
+        ) from exc
+    return SandboxPool
+
+
+def _get_pool(image: str, flavor: str) -> Any:
+    key = (image, flavor)
+    with _POOL_LOCK:
+        pool = _POOLS.get(key)
+        if pool is None:
+            SandboxPool = _get_sandbox_pool_cls()
+            pool = SandboxPool(
+                image=image,
+                flavor=flavor,
+                name=_pool_name(image, flavor),
+            )
+            _POOLS[key] = pool
+        return pool
+
+
+class _LocalAuthProxy:
+    def __init__(self, *, target_url: str, headers: dict[str, str]):
+        self.target_url = target_url.rstrip("/")
+        self.headers = headers
+        self.port = _find_available_port()
+        self._server: Any = None
+        self._thread: threading.Thread | None = None
+
+    @property
+    def base_url(self) -> str:
+        return f"http://127.0.0.1:{self.port}"
+
+    def start(self) -> str:
+        app = FastAPI()
+
+        @app.api_route(
+            "/{path:path}",
+            methods=["GET", "POST", "PUT", "PATCH", "DELETE", "OPTIONS"],
+        )
+        async def proxy_http(path: str, request: Request) -> Response:
+            query = request.url.query
+            target = f"{self.target_url}/{path}"
+            if query:
+                target = f"{target}?{query}"
+            headers = dict(self.headers)
+            content_type = request.headers.get("content-type")
+            if content_type is not None:
+                headers["Content-Type"] = content_type
+            body = await request.body()
+            try:
+                upstream = await asyncio.to_thread(
+                    requests.request,
+                    request.method,
+                    target,
+                    data=body,
+                    headers=headers,
+                    timeout=60.0,
+                    allow_redirects=True,
+                )
+            except requests.RequestException:
+                return Response(
+                    content=b"upstream HF job unreachable",
+                    status_code=502,
+                )
+            response_headers = {
+                key: value
+                for key, value in upstream.headers.items()
+                if key.lower() not in _HOP_BY_HOP_HEADERS
+            }
+            return Response(
+                content=upstream.content,
+                status_code=upstream.status_code,
+                headers=response_headers,
+            )
+
+        @app.websocket("/{path:path}")
+        async def proxy_websocket(path: str, websocket: WebSocket) -> None:
+            query = websocket.url.query
+            target = f"{_to_ws_url(self.target_url)}/{path}"
+            if query:
+                target = f"{target}?{query}"
+            headers = {**self.headers, "accept": "*/*"}
+            upstream = await self._connect_upstream_websocket(target, headers)
+            await websocket.accept()
+            try:
+                to_upstream = asyncio.create_task(
+                    self._client_to_upstream(websocket, upstream)
+                )
+                to_client = asyncio.create_task(
+                    self._upstream_to_client(websocket, upstream)
+                )
+                done, pending = await asyncio.wait(
+                    {to_upstream, to_client},
+                    return_when=asyncio.FIRST_COMPLETED,
+                )
+                for task in pending:
+                    task.cancel()
+                for task in done:
+                    with suppress(ConnectionClosed, WebSocketDisconnect):
+                        task.result()
+            finally:
+                with suppress(Exception):
+                    await upstream.close()
+
+        import uvicorn
+
+        config = uvicorn.Config(
+            app,
+            host="127.0.0.1",
+            port=self.port,
+            log_level="warning",
+            access_log=False,
+        )
+        self._server = uvicorn.Server(config)
+        self._thread = threading.Thread(target=self._server.run, daemon=True)
+        self._thread.start()
+        while not self._server.started:
+            if not self._thread.is_alive():
+                raise RuntimeError("HF sandbox auth proxy failed to start")
+            time.sleep(0.05)
+        return self.base_url
+
+    async def _connect_upstream_websocket(
+        self, target: str, headers: dict[str, str]
+    ) -> Any:
+        last_error: Exception | None = None
+        for _ in range(5):
+            try:
+                return await ws_connect(
+                    target,
+                    additional_headers=headers,
+                    max_size=_MAX_WS_MESSAGE_SIZE,
+                    compression=None,
+                )
+            except Exception as exc:
+                last_error = exc
+                await asyncio.sleep(0.5)
+        assert last_error is not None
+        raise last_error
+
+    async def _client_to_upstream(self, websocket: WebSocket, upstream: Any) -> None:
+        # EnvClient sends JSON text frames; binary frames are only relayed downstream.
+        async for message in websocket.iter_text():
+            await upstream.send(message)
+
+    async def _upstream_to_client(self, websocket: WebSocket, upstream: Any) -> None:
+        async for message in upstream:
+            if isinstance(message, bytes):
+                await websocket.send_bytes(message)
+            else:
+                await websocket.send_text(message)
+
+    def stop(self) -> None:
+        if self._server is None or self._thread is None:
+            return
+        self._server.should_exit = True
+        self._thread.join(timeout=5.0)
+        self._server = None
+        self._thread = None
+
+
+class HFSandboxProvider(ContainerProvider):
+    """Run an OpenEnv server on Hugging Face infrastructure.
+
+    The server runs in a `huggingface_hub` sandbox pool, which bills its hosts as
+    Hugging Face Jobs. `start_container()` returns a local proxy URL that adds the
+    Hugging Face auth headers, so any `EnvClient` can connect to it unchanged.
+
+    Args:
+        image (`str`):
+            Image to run, for example `hf.co/spaces/openenv/coding_env`. It must
+            provide a `server` command that serves OpenEnv on port 8000.
+        env_vars (`dict[str, str]`, *optional*):
+            Environment variables for the server.
+        flavor (`str`, *optional*, defaults to `"cpu-basic"`):
+            Hugging Face hardware flavor of the sandbox hosts.
+    """
+
+    def __init__(
+        self,
+        *,
+        image: str,
+        env_vars: dict[str, str] | None = None,
+        flavor: str = "cpu-basic",
+    ):
+        self.image = image
+        self.env_vars = env_vars
+        self.flavor = flavor
+        self._sandbox: Any = None
+        self._proxy: _LocalAuthProxy | None = None
+
+    def start_container(
+        self,
+        image: str | None = None,
+        port: int | None = None,
+        env_vars: dict[str, str] | None = None,
+        **kwargs: Any,
+    ) -> str:
+        if self._sandbox is not None:
+            raise RuntimeError("HFSandboxProvider already has an active service")
+        if kwargs:
+            unsupported = ", ".join(sorted(kwargs))
+            raise ValueError(
+                f"HFSandboxProvider does not support start_container kwargs: {unsupported}"
+            )
+
+        if port not in (None, _DEFAULT_PORT):
+            raise ValueError(
+                f"HFSandboxProvider only supports port {_DEFAULT_PORT} (got {port})."
+            )
+
+        effective_image = self.image if image is None else image
+        effective_env = self.env_vars if env_vars is None else env_vars
+        pool = _get_pool(effective_image, self.flavor)
+        self._sandbox = pool.create(env=effective_env)
+        try:
+            if not hasattr(self._sandbox, "proxy_url_for") or not hasattr(
+                self._sandbox, "proxy_headers"
+            ):
+                raise RuntimeError(
+                    "HFSandboxProvider requires a huggingface_hub version with "
+                    "Sandbox.proxy_url_for and Sandbox.proxy_headers support."
+                )
+            self._sandbox.run(
+                _SERVER_COMMAND,
+                shell=True,
+            )
+            self._proxy = _LocalAuthProxy(
+                target_url=self._sandbox.proxy_url_for(_DEFAULT_PORT, "/"),
+                headers=self._sandbox.proxy_headers,
+            )
+            return self._proxy.start()
+        except Exception:
+            self.stop_container()
+            raise
+
+    def stop_container(self) -> None:
+        if self._proxy is not None:
+            self._proxy.stop()
+            self._proxy = None
+        if self._sandbox is not None:
+            sandbox = self._sandbox
+            self._sandbox = None
+            with suppress(Exception):
+                sandbox.kill()
+
+    def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
+        deadline = time.time() + timeout_s
+        health_url = f"{base_url}/health"
+        while time.time() < deadline:
+            try:
+                response = requests.get(health_url, timeout=5.0)
+                if response.status_code == 200:
+                    return
+            except requests.exceptions.RequestException:
+                pass
+            time.sleep(1.0)
+        raise TimeoutError(
+            f"HF sandbox job at {base_url} did not become ready within {timeout_s}s"
+        )
+
+    def close(self) -> None:
+        self.stop_container()
+
+
+__all__ = ["HFSandboxProvider"]

src/core/containers/runtime/modal_provider.py

@@ -0,0 +1,697 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Modal container provider for running OpenEnv environments in Modal sandboxes.
+
+Requires the ``modal`` SDK: ``pip install modal>=1.4``
+
+The provider boots an OpenEnv server inside a Modal sandbox, exposes it on an
+encrypted tunnel, and returns an ``https://`` URL that ``EnvClient`` connects to
+over ``wss://``.
+
+Supports both the stable Sandbox API (``modal.Sandbox.create``) and the beta
+Sandbox v2 API (``modal.Sandbox._experimental_create``). Sandbox v2 is opt-in
+via ``use_sandbox_v2=True`` because it is an experimental, private SDK feature;
+see https://modal.com/docs/guide/sandbox-v2 .
+"""
+
+from __future__ import annotations
+
+import logging
+import shlex
+import time
+from typing import Any, Dict, Optional
+
+from ._server_config import parse_dockerfile_cmd, parse_openenv_app_field
+from .providers import ContainerProvider
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_MODAL_PORT = 8000
+
+# Modal's default per-container resource requests. We pass configured cpu/memory
+# as the *limit* half of Modal's ``(request, limit)`` tuple so they cap usage
+# rather than reserving (and billing for) the full amount. See
+# https://modal.com/docs/guide/resources .
+_DEFAULT_CPU_REQUEST = 0.125
+_DEFAULT_MEMORY_REQUEST = 128
+
+
+def _require_secure_url(url: str) -> str:
+    """Enforce https/wss transport (RFC 002 security invariant S1).
+
+    ``EnvClient`` derives its WebSocket URL from this base URL, so a plaintext
+    URL would become a cleartext ``ws://`` connection. The offending URL is
+    deliberately omitted from the error because a Modal tunnel URL is a bearer
+    capability that must not leak into logs.
+    """
+    if not isinstance(url, str) or not url.lower().startswith("https://"):
+        raise RuntimeError(
+            "Modal sandbox returned a non-HTTPS tunnel URL. OpenEnv requires an "
+            "https/wss base_url so EnvClient traffic is encrypted. Refusing to "
+            "connect over plaintext."
+        )
+    return url
+
+
+class _DefaultModalAdapter:
+    """Thin adapter over the ``modal`` SDK.
+
+    The provider talks to this private adapter instead of spreading SDK details
+    through its own logic; tests inject a duck-typed fake in its place. Keeping
+    the SDK surface here also localizes the feature check for the private
+    Sandbox v2 (``_experimental_create``) API.
+    """
+
+    def __init__(self, *, app_name: str, use_sandbox_v2: bool):
+        import modal
+
+        self._modal = modal
+        self._app_name = app_name
+        self._use_sandbox_v2 = use_sandbox_v2
+
+        # Sandbox v2 rides on a private SDK entry point that is not present in
+        # every modal release. Fail fast at construction with clear guidance
+        # rather than at start_container time with an AttributeError.
+        if use_sandbox_v2 and not hasattr(modal.Sandbox, "_experimental_create"):
+            raise RuntimeError(
+                "use_sandbox_v2=True requires modal.Sandbox._experimental_create, "
+                "which is not available in the installed modal SDK. Upgrade modal, "
+                "or use the stable Sandbox API (use_sandbox_v2=False). Sandbox v2 "
+                "is an experimental feature; contact support@modal.com for access."
+            )
+
+    def image_from_registry(self, tag: str) -> Any:
+        return self._modal.Image.from_registry(tag)
+
+    def image_from_dockerfile(self, dockerfile_path: str, context_dir: str) -> Any:
+        return self._modal.Image.from_dockerfile(
+            dockerfile_path, context_dir=context_dir
+        )
+
+    def create_sandbox(
+        self,
+        *,
+        image: Any,
+        encrypted_ports: list[int],
+        timeout: int,
+        env: Optional[dict[str, str]],
+        extra: dict[str, Any],
+    ) -> Any:
+        app = self._modal.App.lookup(self._app_name, create_if_missing=True)
+        kwargs: Dict[str, Any] = {
+            "app": app,
+            "image": image,
+            "encrypted_ports": encrypted_ports,
+            "timeout": timeout,
+            **extra,
+        }
+        if env:
+            kwargs["env"] = dict(env)
+
+        # The sandbox is started with a keep-alive entrypoint; the server
+        # command is launched via exec afterwards (see ModalProvider).
+        if self._use_sandbox_v2:
+            return self._modal.Sandbox._experimental_create(
+                "sleep", "infinity", **kwargs
+            )
+        return self._modal.Sandbox.create("sleep", "infinity", **kwargs)
+
+    def exec(self, sandbox: Any, command: str, *, timeout: int = 10) -> str:
+        """Run *command* through a shell inside *sandbox* and return its stdout."""
+        proc = sandbox.exec("bash", "-c", command, timeout=timeout)
+        try:
+            out = proc.stdout.read()
+        except Exception:
+            out = ""
+        proc.wait()
+        return out or ""
+
+    def tunnel_url(self, sandbox: Any, port: int) -> str:
+        return sandbox.tunnels()[port].url
+
+    def terminate(self, sandbox: Any) -> None:
+        sandbox.terminate()
+
+
+class ModalProvider(ContainerProvider):
+    """
+    Container provider that runs environments in Modal sandboxes.
+
+    ``start_container``'s ``image`` is either a registry tag
+    (``"echo-env:latest"``) or a ``"dockerfile:<path>"`` reference returned by
+    :meth:`image_from_dockerfile`. The server is exposed on an encrypted Modal
+    tunnel and the returned ``https://`` URL is what ``EnvClient`` connects to
+    over ``wss://``.
+
+    The environment runs untrusted code, so the provider is secure by default:
+    it enforces https/wss transport, treats the tunnel URL as a bearer secret
+    (never interpolated into errors), and never surfaces raw sandbox output
+    unless ``surface_server_logs=True``.
+
+    Only one sandbox is active per provider: calling ``start_container`` again
+    before ``stop_container()``/``close()`` raises ``RuntimeError`` rather than
+    orphaning the running sandbox. ``close()`` (and context-manager exit) stops
+    the active sandbox.
+
+    Example:
+        ```python
+        with ModalProvider(app_name="openenv", cpu=2.0, memory=4096) as provider:
+            image = ModalProvider.image_from_dockerfile(
+                "envs/echo_env/server/Dockerfile"
+            )
+            base_url = provider.start_container(image)
+            provider.wait_for_ready(base_url)
+        # sandbox terminated on exit
+        ```
+
+    Sandbox v2 (beta) is opt-in:
+        ```python
+        provider = ModalProvider(app_name="openenv", use_sandbox_v2=True)
+        ```
+    """
+
+    _dockerfile_registry: Dict[str, Dict[str, Any]] = {}
+
+    def __init__(
+        self,
+        *,
+        image: str | None = None,
+        env_vars: dict[str, str] | None = None,
+        app_name: str = "openenv",
+        use_sandbox_v2: bool = False,
+        timeout: int = 300,
+        cmd: str | None = None,
+        cpu: float | None = None,
+        memory: int | None = None,
+        surface_server_logs: bool = False,
+        _adapter: Any = None,
+    ):
+        """
+        Args:
+            image (`str`, *optional*):
+                Registry image tag or ``"dockerfile:<path>"`` source to use
+                when ``start_container()`` is called without an image.
+            env_vars (`dict`, *optional*):
+                Environment variables to use when ``start_container()`` is
+                called without explicit ``env_vars``.
+            app_name (`str`, *optional*, defaults to `"openenv"`):
+                Modal app name the sandbox is created under. Looked up (and
+                created if missing) via ``modal.App.lookup``.
+            use_sandbox_v2 (`bool`, *optional*, defaults to `False`):
+                When `True`, sandboxes are created via the beta
+                ``modal.Sandbox._experimental_create`` API (Sandbox v2). This is
+                an experimental, private SDK feature and is therefore off by
+                default; a feature check fails fast at construction if the
+                installed SDK lacks it. See https://modal.com/docs/guide/sandbox-v2 .
+            timeout (`int`, *optional*, defaults to `300`):
+                Maximum sandbox lifetime in seconds.
+            cmd (`str`, *optional*):
+                Shell command to start the server inside the sandbox. When
+                omitted, the command is auto-discovered from ``openenv.yaml``
+                (falling back to the Dockerfile ``CMD``).
+            cpu (`float`, *optional*):
+                Hard CPU-core limit for the sandbox. Passed as the limit half of
+                Modal's ``(request, limit)`` tuple (request stays at Modal's
+                default), so it caps usage rather than reserving cores. When
+                `None`, Modal's default applies.
+            memory (`int`, *optional*):
+                Hard memory limit in MiB for the sandbox (containers exceeding it
+                are OOM-killed). Passed as the limit half of Modal's
+                ``(request, limit)`` tuple. When `None`, Modal's default applies.
+            surface_server_logs (`bool`, *optional*, defaults to `False`):
+                When `False` (default), captured sandbox output is withheld from
+                raised errors so secrets the workload printed cannot leak into
+                orchestrator/CI logs. When `True`, a best-effort redacted,
+                length-bounded excerpt is included in startup-crash errors.
+        """
+        self._image = image
+        self._env_vars = env_vars
+        self._app_name = app_name
+        self._use_sandbox_v2 = use_sandbox_v2
+        self._timeout = timeout
+        self._cmd = cmd
+        self._cpu = cpu
+        self._memory = memory
+        self._surface_server_logs = surface_server_logs
+        self._sandbox: Any = None
+        self._base_url: str | None = None
+        # Injected env-var values, used to scrub captured server output before
+        # it is ever surfaced in an error.
+        self._redact_values: set[str] = set()
+
+        if _adapter is None:
+            # Import eagerly (inside the adapter) so SDK/configuration errors —
+            # including the Sandbox v2 feature check — surface at construction.
+            self._adapter: Any = _DefaultModalAdapter(
+                app_name=app_name, use_sandbox_v2=use_sandbox_v2
+            )
+        else:
+            self._adapter = _adapter
+
+        if use_sandbox_v2:
+            logger.info(
+                "Using Modal Sandbox v2 (experimental). This feature must be "
+                "explicitly enabled. Contact support@modal.com to get access."
+            )
+
+    def _discover_server_cmd(self, port: int = _DEFAULT_MODAL_PORT) -> str:
+        """Discover the server command from ``openenv.yaml`` inside the sandbox.
+
+        Finds the file, reads the ``app`` field, and constructs a command
+        of the form ``cd <env_root> && python -m uvicorn <app> --host 0.0.0.0 --port <port>``.
+
+        Raises:
+            ValueError: If ``openenv.yaml`` is not found or lacks an ``app`` field.
+        """
+        yaml_path = self._find_openenv_yaml()
+        if yaml_path is None:
+            raise ValueError(
+                "Could not find openenv.yaml inside the sandbox. "
+                "Pass an explicit cmd= to ModalProvider or start_container()."
+            )
+
+        content = self._adapter.exec(self._sandbox, f"cat {shlex.quote(yaml_path)}")
+        app = self._parse_app_field(content)
+        if app is None:
+            raise ValueError(
+                f"openenv.yaml at {yaml_path} does not contain an 'app' field. "
+                "Pass an explicit cmd= to ModalProvider or start_container()."
+            )
+
+        # The directory containing openenv.yaml is the env root
+        env_root = yaml_path.rsplit("/", 1)[0]
+        return (
+            f"cd {shlex.quote(env_root)} && "
+            f"python -m uvicorn {shlex.quote(app)} --host 0.0.0.0 --port {port}"
+        )
+
+    def _find_openenv_yaml(self) -> str | None:
+        """Locate ``openenv.yaml`` inside the sandbox.
+
+        Tries the modern layout path ``/app/env/openenv.yaml`` first,
+        then falls back to a ``find`` command for the old layout.
+        """
+        # Fast path: modern Dockerfile layout
+        out = self._adapter.exec(
+            self._sandbox, "test -f /app/env/openenv.yaml && echo found"
+        )
+        if "found" in (out or ""):
+            return "/app/env/openenv.yaml"
+
+        # Fallback: search for it (redirect stderr so error messages
+        # like "No such file or directory" don't get mistaken for paths).
+        path = self._adapter.exec(
+            self._sandbox,
+            "find /app -maxdepth 4 -name openenv.yaml -print -quit 2>/dev/null",
+        ).strip()
+        if path and path.startswith("/"):
+            return path
+
+        return None
+
+    @staticmethod
+    def _parse_app_field(yaml_content: str) -> str | None:
+        """Extract the ``app`` value from raw openenv.yaml content.
+
+        Uses PyYAML to handle comments, quotes, and nested keys correctly.
+        """
+        return parse_openenv_app_field(yaml_content)
+
+    @staticmethod
+    def _parse_dockerfile_cmd(dockerfile_content: str) -> str | None:
+        """Extract the server command from the last ``CMD`` in a Dockerfile.
+
+        Handles exec form (``CMD ["prog", "arg"]``) and shell form
+        (``CMD prog arg``).  When a Dockerfile has multiple ``CMD``
+        instructions (e.g. multi-stage builds), the last one wins - same
+        semantics as Docker itself.  Lines where ``CMD`` appears inside a
+        comment are ignored.
+
+        Returns:
+            The command as a single string, or ``None`` if no ``CMD`` found.
+        """
+        return parse_dockerfile_cmd(dockerfile_content)
+
+    @classmethod
+    def image_from_dockerfile(
+        cls,
+        dockerfile_path: str,
+        context_dir: str | None = None,
+    ) -> str:
+        """Validate a Dockerfile and return a ``dockerfile:`` URI for
+        :meth:`start_container`.
+
+        Eagerly validates the Dockerfile (existence, COPY sources) and stores
+        its content in an internal registry.  The actual ``modal.Image`` is
+        created later inside ``start_container``.
+
+        Args:
+            dockerfile_path (`str`):
+                Path to the Dockerfile on disk.
+            context_dir (`str`, *optional*):
+                Build context directory.  Defaults to the Dockerfile's
+                grandparent directory, matching the ``openenv init``
+                convention where Dockerfiles live in
+                ``<env>/server/Dockerfile`` and the build context is
+                ``<env>/``.  Pass explicitly for non-standard layouts
+                (e.g. ``context_dir="."`` for repo-root contexts).
+
+        Returns:
+            `str`: A ``"dockerfile:<abs_path>"`` string to pass to
+            ``start_container``.
+
+        Raises:
+            FileNotFoundError: If *dockerfile_path* does not exist.
+            ValueError: If *context_dir* is given but does not exist,
+                or if COPY sources in the Dockerfile cannot be found
+                under the resolved context directory.
+        """
+        import pathlib
+        import re
+
+        src = pathlib.Path(dockerfile_path).resolve()
+        if not src.is_file():
+            raise FileNotFoundError(f"Dockerfile not found: {dockerfile_path}")
+
+        if context_dir is not None:
+            ctx = pathlib.Path(context_dir)
+            if not ctx.is_dir():
+                raise ValueError(f"context_dir does not exist: {context_dir}")
+        else:
+            # Default: grandparent of the Dockerfile, matching the
+            # openenv init layout (<env>/server/Dockerfile -> <env>/).
+            ctx = src.parent.parent
+
+        content = src.read_text()
+
+        # Validate that COPY sources exist under the context directory.
+        # This catches mismatches early (e.g. a Dockerfile expecting repo
+        # root as context when we defaulted to the env directory).
+        for line in content.splitlines():
+            m = re.match(r"^\s*COPY\s+(?!--from=)(\S+)\s+", line, re.IGNORECASE)
+            if not m:
+                continue
+            copy_src = m.group(1)
+            if copy_src.startswith("/"):
+                continue
+            resolved = ctx / copy_src
+            if not resolved.exists() and not any(ctx.glob(copy_src)):
+                raise ValueError(
+                    f"Dockerfile COPY source '{copy_src}' not found "
+                    f"under context_dir '{ctx}'. This Dockerfile may "
+                    f"expect a different build context (e.g. the repo "
+                    f"root). Pass context_dir explicitly."
+                )
+
+        # Parse CMD from the Dockerfile so start_container can use it as a
+        # fallback when openenv.yaml is unavailable.
+        parsed_cmd = cls._parse_dockerfile_cmd(content)
+
+        cls._dockerfile_registry[str(src)] = {
+            "dockerfile_path": str(src),
+            "context_dir": str(ctx),
+            "server_cmd": parsed_cmd,
+        }
+
+        return f"dockerfile:{src}"
+
+    def _build_image(self, image: str) -> Any:
+        """Build the ``modal.Image`` for *image* (registry tag or dockerfile:)."""
+        if image.startswith("dockerfile:"):
+            dockerfile_path = image[len("dockerfile:") :]
+            meta = self._dockerfile_registry.get(dockerfile_path)
+            if meta is None:
+                raise ValueError(
+                    f"No registered Dockerfile metadata for {dockerfile_path}. "
+                    "Call ModalProvider.image_from_dockerfile() first."
+                )
+            return self._adapter.image_from_dockerfile(
+                meta["dockerfile_path"], meta["context_dir"]
+            )
+
+        # Plain registry tag (e.g. "echo-env:latest").
+        return self._adapter.image_from_registry(image)
+
+    def start_container(
+        self,
+        image: str | None = None,
+        port: int | None = None,
+        env_vars: dict[str, str] | None = None,
+        **kwargs: Any,
+    ) -> str:
+        """
+        Create a Modal sandbox from a Docker image or Dockerfile.
+
+        The sandbox is started with a keep-alive process and the server
+        command is launched via ``exec`` afterwards, mirroring the discovery
+        flow used by other cloud providers. The server command is resolved in
+        order:
+
+        1. Explicit ``cmd`` passed to the constructor.
+        2. ``cmd`` key in ``**kwargs`` (popped before forwarding).
+        3. Auto-discovered from ``openenv.yaml`` inside the sandbox.
+        4. ``CMD`` parsed from the Dockerfile (when *image* came from
+           ``image_from_dockerfile``).
+
+        Args:
+            image (`str`, *optional*):
+                Registry image tag (e.g. ``"echo-env:latest"``) or
+                ``"dockerfile:<path>"`` returned by
+                :meth:`image_from_dockerfile`. May be omitted when supplied to
+                the constructor.
+            port (`int`, *optional*):
+                Must be ``None`` or ``8000``. Modal exposes port 8000 via an
+                encrypted tunnel; other ports raise ``ValueError``.
+            env_vars (`dict`, *optional*):
+                Environment variables forwarded to the sandbox.
+            **kwargs:
+                ``cmd`` (`str`) to override the server command; any remaining
+                keyword arguments are forwarded to ``modal.Sandbox.create``.
+
+        Returns:
+            `str`: HTTPS tunnel URL for the sandbox (base_url).
+        """
+        if self._sandbox is not None:
+            raise RuntimeError(
+                "ModalProvider already has an active sandbox. Call "
+                "stop_container() (or close()) before starting another — a "
+                "second start would orphan the running sandbox."
+            )
+
+        if port is not None and port != _DEFAULT_MODAL_PORT:
+            raise ValueError(
+                f"ModalProvider only supports port {_DEFAULT_MODAL_PORT} "
+                f"(got {port}). The Modal tunnel routes to port "
+                f"{_DEFAULT_MODAL_PORT} inside the sandbox."
+            )
+
+        effective_image = image if image is not None else self._image
+        if effective_image is None:
+            raise ValueError(
+                "ModalProvider requires an image. Pass it to the constructor "
+                "or start_container()."
+            )
+        effective_env_vars = self._env_vars if env_vars is None else env_vars
+
+        # Resolve the server command (may be None; discovery happens after
+        # sandbox creation when we can inspect the filesystem).
+        cmd = kwargs.pop("cmd", None) or self._cmd
+
+        # CMD parsed from Dockerfile (populated for "dockerfile:" images).
+        parsed_cmd: str | None = None
+        if effective_image.startswith("dockerfile:"):
+            meta = self._dockerfile_registry.get(effective_image[len("dockerfile:") :])
+            if meta is not None:
+                parsed_cmd = meta.get("server_cmd")
+
+        modal_image = self._build_image(effective_image)
+
+        extra: Dict[str, Any] = dict(kwargs)
+        if self._cpu is not None:
+            extra["cpu"] = (_DEFAULT_CPU_REQUEST, self._cpu)
+        if self._memory is not None:
+            extra["memory"] = (_DEFAULT_MEMORY_REQUEST, self._memory)
+
+        # Record injected secret values so captured server output can be
+        # scrubbed before it is ever surfaced in an error.
+        self._redact_values = {
+            value for value in (effective_env_vars or {}).values() if value
+        }
+
+        # A create failure created nothing, so just drop the recorded secrets
+        # and re-raise (the double-start guard above guarantees there is no
+        # pre-existing sandbox to delete).
+        try:
+            self._sandbox = self._adapter.create_sandbox(
+                image=modal_image,
+                encrypted_ports=[_DEFAULT_MODAL_PORT],
+                timeout=self._timeout,
+                env=effective_env_vars,
+                extra=extra,
+            )
+        except Exception:
+            self._redact_values = set()
+            raise
+
+        try:
+            # Discover server command from openenv.yaml if not explicitly set.
+            if cmd is None:
+                try:
+                    cmd = self._discover_server_cmd()
+                except ValueError:
+                    # Fall back to CMD parsed from Dockerfile (if available).
+                    if parsed_cmd:
+                        cmd = parsed_cmd
+                    else:
+                        raise
+
+            # Launch the server in the background. Write the PID so we can
+            # check whether the process crashed in wait_for_ready().
+            escaped_cmd = shlex.quote(cmd)
+            self._adapter.exec(
+                self._sandbox,
+                f"nohup bash -c {escaped_cmd} > /tmp/openenv-server.log 2>&1 &"
+                " echo $! > /tmp/openenv-server.pid",
+            )
+
+            # Resolve the public tunnel URL for port 8000.
+            self._base_url = _require_secure_url(
+                self._adapter.tunnel_url(self._sandbox, _DEFAULT_MODAL_PORT)
+            )
+        except Exception:
+            # A cleanup failure here must not mask the original error: swallow
+            # any exception from stop_container() so the root cause propagates.
+            try:
+                self.stop_container()
+            except Exception:
+                pass
+            raise
+
+        return self._base_url
+
+    def stop_container(self) -> None:
+        """Terminate the Modal sandbox."""
+        if self._sandbox is None:
+            # Still drop any injected secret values recorded by a failed start.
+            self._redact_values = set()
+            return
+
+        try:
+            self._adapter.terminate(self._sandbox)
+        finally:
+            self._sandbox = None
+            self._base_url = None
+            self._redact_values = set()
+
+    def close(self) -> None:
+        """Stop the active sandbox.
+
+        Overrides the base no-op so a caller holding a bare ``ContainerProvider``
+        reference can release the sandbox polymorphically (also invoked on
+        context-manager exit). ``ModalProvider`` holds no separate SDK client, so
+        this is equivalent to ``stop_container()``.
+        """
+        self.stop_container()
+
+    @property
+    def base_url(self) -> str:
+        """URL returned by the last ``start_container``."""
+        if self._base_url is None:
+            raise RuntimeError(
+                "ModalProvider has no active base_url. Start the provider "
+                "before reading base_url."
+            )
+        return self._base_url
+
+    def _redact(self, text: str, *, max_chars: int = 2000) -> str:
+        """Scrub injected secret values and bound length before surfacing output.
+
+        Replaces any injected env-var value with `***` and keeps only the tail.
+        This is best-effort (exact-match only), which is why server output is
+        withheld entirely unless ``surface_server_logs=True``.
+        """
+        redacted = text or ""
+        for value in self._redact_values:
+            redacted = redacted.replace(value, "***")
+        if len(redacted) > max_chars:
+            redacted = "...(truncated)...\n" + redacted[-max_chars:]
+        return redacted
+
+    def _server_died_message(self) -> str:
+        """Build the startup-crash error, secure by default.
+
+        Untrusted code can print secrets then force a crash to exfiltrate them
+        through the exception (which lands in orchestrator/CI logs), so sandbox
+        output is excluded unless ``surface_server_logs=True``, in which case a
+        best-effort redacted, bounded excerpt is included.
+        """
+        base = (
+            "Modal sandbox server process died during startup. Server output is "
+            "not surfaced to avoid leaking secrets injected into the sandbox; "
+            "retrieve /tmp/openenv-server.log from the sandbox out of band, or "
+            "construct the provider with surface_server_logs=True to include a "
+            "redacted excerpt."
+        )
+        if not self._surface_server_logs or self._sandbox is None:
+            return base
+
+        log = self._redact(
+            self._adapter.exec(self._sandbox, "cat /tmp/openenv-server.log 2>/dev/null")
+        )
+        return (
+            "Modal sandbox server process died during startup. The excerpt below "
+            "is the sandbox server output with injected secret values redacted "
+            "(best-effort); it may still contain secrets the workload printed "
+            f"by other means.\nLog (redacted):\n{log}"
+        )
+
+    def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
+        """
+        Poll the /health endpoint until the sandbox is ready.
+
+        Uses a longer default timeout (120s) than local Docker providers
+        because Modal sandboxes may have cold-start latency.
+
+        Args:
+            base_url (`str`):
+                Tunnel URL returned by ``start_container()``.
+            timeout_s (`float`, *optional*, defaults to `120.0`):
+                Maximum seconds to wait.
+
+        Raises:
+            TimeoutError: If the sandbox doesn't become ready in time.
+            RuntimeError: If the server process died (detected via PID check).
+        """
+        import requests
+
+        health_url = f"{base_url}/health"
+
+        deadline = time.time() + timeout_s
+        while time.time() < deadline:
+            try:
+                response = requests.get(health_url, timeout=5.0)
+                if response.status_code == 200:
+                    return
+            except requests.RequestException:
+                pass
+
+            # Early exit: if the server process died, raise immediately
+            # instead of waiting for the full health-check timeout.
+            if self._sandbox is not None:
+                out = self._adapter.exec(
+                    self._sandbox,
+                    "kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null"
+                    " && echo RUNNING || echo DEAD",
+                )
+                if "DEAD" in (out or ""):
+                    raise RuntimeError(self._server_died_message())
+
+            time.sleep(1.0)
+
+        # The tunnel URL is a bearer capability, so it is deliberately omitted
+        # from the timeout error.
+        raise TimeoutError(f"Modal sandbox did not become ready within {timeout_s}s.")
+
+
+__all__ = ["ModalProvider"]

src/core/containers/runtime/novita_provider.py

@@ -0,0 +1,1216 @@
+# SPDX-License-Identifier: BSD-3-Clause
+
+"""
+Novita AI Sandbox container provider for running OpenEnv environments.
+
+Requires the ``novita-sandbox`` SDK: ``pip install openenv[novita]``
+
+The provider boots an OpenEnv server inside a Novita sandbox, exposes port 8000
+on the sandbox's public host, and returns an ``https://`` URL that ``EnvClient``
+connects to over ``wss://``.
+
+Note: the ``RFC 002 security invariant S<n>`` references in this module point to
+the **proposed** "Cloud Sandbox Providers" amendment to RFC 002, which is pending
+review/sign-off by the RFC authors. They describe the security properties this
+provider enforces; treat them as proposed (not yet ratified) until that sign-off.
+"""
+
+from __future__ import annotations
+
+import re
+import shlex
+import time
+from typing import Any, Callable, Dict, List, Optional
+
+from ._server_config import parse_openenv_app_field
+from .providers import ContainerProvider
+
+_DEFAULT_NOVITA_PORT = 8000
+
+# Novita defaults a sandbox to a 300s hard lifetime, which would kill an RL
+# rollout mid-episode. One hour is the documented ceiling for Hobby accounts, so
+# it is the largest default that cannot make ``create`` fail on plan limits;
+# longer-running work passes ``timeout=`` explicitly.
+_DEFAULT_SANDBOX_TIMEOUT_S = 3600
+
+# The sandbox is created with this as its start command so the image's own CMD
+# is not launched (see `_launch_server`). `sleep infinity` keeps the sandbox
+# alive with nothing bound to port 8000, leaving the port free for the server
+# this provider starts itself.
+_KEEPALIVE_CMD = "sleep infinity"
+_KEEPALIVE_READY_MS = 5_000
+
+# Resource defaults for a template built from a Dockerfile. The SDK's own
+# defaults, restated here so they are visible at the call site.
+_DEFAULT_TEMPLATE_CPU = 2
+_DEFAULT_TEMPLATE_MEMORY_MB = 1024
+
+# User the sandbox runs as. Root because the environment installs under /app,
+# which the image creates as root and the Novita parser's default user ("user")
+# cannot write to -- see `build_template`. Matches what a plain `docker run` of
+# the same image gives, which is what the other providers' users observe.
+_SANDBOX_USER = "root"
+
+_FROM_RE = re.compile(r"^\s*FROM\s+(?P<rest>.+?)\s*$", re.IGNORECASE)
+_COPY_FROM_RE = re.compile(
+    r"^\s*COPY\s+--from=(?P<stage>\S+)\s+(?P<src>\S+)\s+(?P<dst>\S+)\s*$",
+    re.IGNORECASE,
+)
+_COPY_FROM_REMAINDER_RE = re.compile(r"^\s*COPY\b.*\s--from(?:=|\s|$)", re.IGNORECASE)
+_USER_RE = re.compile(r"^\s*USER(?:\s+|$)", re.IGNORECASE)
+_RUN_RE = re.compile(r"^\s*RUN\s+(?P<rest>.*)$", re.IGNORECASE)
+_ARG_DEFAULT_RE = re.compile(
+    r"^\s*ARG\s+(?P<name>\w+)=(?P<value>\S+)\s*$", re.IGNORECASE
+)
+
+
+def _strip_mount_flags(content: str) -> str:
+    """Remove supported BuildKit cache mounts from ``RUN`` instructions.
+
+    Novita's template parser copies a ``RUN`` body verbatim into the build
+    command, so a leading ``--mount=type=cache,...`` would be handed to the
+    shell as an argument rather than honored as BuildKit syntax. Every in-repo
+    OpenEnv Dockerfile uses one to cache ``uv`` downloads, so explicit cache
+    mounts are stripped rather than left to fail the build. Other mount types
+    affect build semantics and cannot be represented by Novita's parser, so
+    they are rejected instead of silently removed.
+    """
+    lines = content.split("\n")
+    out: List[str] = []
+    in_run = False
+    for line in lines:
+        if _RUN_RE.match(line):
+            in_run = True
+        if in_run:
+            # Strip leading --mount flags from the RUN line itself and from any
+            # continuation lines that follow it.
+            stripped = line
+            prefix = ""
+            run_match = _RUN_RE.match(line)
+            if run_match:
+                prefix = line[: run_match.start("rest")]
+                stripped = run_match.group("rest")
+            while True:
+                mount = re.match(r"\s*--mount=(?P<spec>\S+)(?P<space>\s*)", stripped)
+                if not mount:
+                    break
+                spec = mount.group("spec")
+                options = spec.split(",")
+                types = [
+                    option.split("=", 1)[1].lower()
+                    for option in options
+                    if option.lower().startswith("type=")
+                ]
+                if types != ["cache"]:
+                    raise ValueError(
+                        "Novita Dockerfile rewriting only supports explicit "
+                        "RUN --mount=type=cache options; unsupported mount "
+                        "types must be handled by `openenv build` and a "
+                        "registry image."
+                    )
+                stripped = stripped[mount.end() :]
+            line = prefix + stripped
+            if not line.rstrip().endswith("\\"):
+                in_run = False
+        out.append(line)
+    return "\n".join(out)
+
+
+def _resolve_from_references(content: str) -> str:
+    """Substitute ``ARG`` defaults into ``FROM`` lines and drop ``--platform``.
+
+    Two parser limitations, handled together because both corrupt the base
+    image reference:
+
+    - The parser stores a ``FROM ${BASE_IMAGE}`` line verbatim, so the template
+      would be built from a literal image named ``${BASE_IMAGE}``. Only global
+      ARGs (declared before the first ``FROM``) are in scope for a ``FROM``
+      line, which is the form every in-repo Dockerfile uses.
+    - ``FROM --platform=linux/amd64 python:3.10-slim`` likewise keeps the flag
+      as part of the name. Novita builds for its own platform, so the flag is
+      dropped rather than propagated.
+
+    Substitutions are applied longest-name-first so a short unbraced prefix
+    cannot corrupt a longer sibling (e.g. ``$BASE`` must not rewrite
+    ``$BASE_IMAGE`` / ``${BASE_IMAGE}`` into ``python:3.12_IMAGE``).
+    """
+    arg_defaults: Dict[str, str] = {}
+    for line in content.split("\n"):
+        if _FROM_RE.match(line):
+            break
+        match = _ARG_DEFAULT_RE.match(line)
+        if match:
+            arg_defaults[match.group("name")] = match.group("value")
+
+    # Longest names first: `$BASE` must not rewrite `$BASE_IMAGE` mid-token.
+    arg_items = sorted(
+        arg_defaults.items(), key=lambda item: len(item[0]), reverse=True
+    )
+
+    out: List[str] = []
+    for line in content.split("\n"):
+        match = _FROM_RE.match(line)
+        if not match:
+            out.append(line)
+            continue
+
+        reference = match.group("rest").strip()
+        reference = re.sub(r"^(--platform=\S+\s*)+", "", reference).strip()
+        for name, value in arg_items:
+            reference = reference.replace(f"${{{name}}}", value)
+            reference = reference.replace(f"${name}", value)
+        out.append(f"FROM {reference}")
+    return "\n".join(out)
+
+
+def _split_stages(content: str) -> List[Dict[str, Any]]:
+    """Split a Dockerfile into its ``FROM``-delimited stages."""
+    stages: List[Dict[str, Any]] = []
+    current: Optional[Dict[str, Any]] = None
+    for line in content.split("\n"):
+        match = _FROM_RE.match(line)
+        if match:
+            rest = match.group("rest").strip()
+            alias_match = re.match(r"(?P<base>.+?)\s+[Aa][Ss]\s+(?P<alias>\S+)$", rest)
+            if alias_match:
+                base = alias_match.group("base").strip()
+                alias = alias_match.group("alias")
+            else:
+                base, alias = rest, None
+            current = {"base": base, "alias": alias, "body": []}
+            stages.append(current)
+        elif current is not None:
+            current["body"].append(line)
+    return stages
+
+
+def _dockerfile_declares_user(content: str) -> bool:
+    """Return whether *content* contains a Dockerfile ``USER`` instruction."""
+    return any(_USER_RE.match(line) for line in content.splitlines())
+
+
+def _reject_unflattened_copy_from(content: str) -> None:
+    """Reject COPY instructions with ``--from`` left after flattening."""
+    if any(_COPY_FROM_REMAINDER_RE.match(line) for line in content.splitlines()):
+        raise ValueError(
+            "Dockerfile contains a COPY --from instruction that Novita cannot "
+            "flatten. Use a simple `COPY --from=<stage> <src> <dst>` form or "
+            "build the image with `openenv build` and pass the registry image."
+        )
+
+
+def _flatten_multistage(content: str) -> str:
+    """Collapse a multi-stage Dockerfile into a single stage.
+
+    Novita's parser rejects multi-stage Dockerfiles outright, and that is the
+    layout every in-repo OpenEnv environment uses: a ``builder`` stage that runs
+    ``uv sync`` into ``/app/env/.venv``, then a runtime stage that copies the
+    result across. Both stages name the *same* base image, so the build can be
+    replayed linearly in one stage -- the builder's ``RUN`` steps execute, then
+    each ``COPY --from=builder`` becomes an in-place ``cp``.
+
+    A copy whose source and destination are equal is dropped: after flattening,
+    the file is already at its destination, and Novita rejects an absolute
+    source path.
+
+    Raises:
+        ValueError: If the stages do not share one base image, a copy pulls from
+            a stage other than the builder, or a ``COPY --from`` cannot be
+            flattened.
+    """
+    stages = _split_stages(content)
+    if len(stages) < 2:
+        _reject_unflattened_copy_from(content)
+        return content
+
+    bases = {stage["base"] for stage in stages}
+    if len(bases) != 1:
+        raise ValueError(
+            "Novita templates cannot express a multi-stage Dockerfile whose "
+            f"stages use different base images: {sorted(bases)}. Build the "
+            "image with `openenv build` and push it with "
+            "`openenv push --registry`, then pass the registry reference to "
+            "NovitaSandboxProvider(image=...)."
+        )
+
+    builder_alias = stages[0]["alias"]
+    if builder_alias is None:
+        raise ValueError(
+            "The first stage of a multi-stage Dockerfile must be named with "
+            "`AS <alias>` so its output can be referenced after flattening."
+        )
+
+    out: List[str] = [f"FROM {stages[0]['base']}"]
+    for index, stage in enumerate(stages):
+        for line in stage["body"]:
+            copy_match = _COPY_FROM_RE.match(line)
+            if copy_match is None:
+                out.append(line)
+                continue
+            if index == 0:
+                # A copy within the builder stage is already valid as written.
+                out.append(line)
+                continue
+            source_stage = copy_match.group("stage")
+            if source_stage != builder_alias:
+                raise ValueError(
+                    f"COPY --from={source_stage} does not reference the builder "
+                    f"stage ({builder_alias}), so it cannot be flattened into a "
+                    "single-stage template."
+                )
+            src = copy_match.group("src")
+            dst = copy_match.group("dst")
+            if src == dst:
+                continue
+            out.append(
+                f"RUN mkdir -p $(dirname {shlex.quote(dst)}) && "
+                f"cp -a {shlex.quote(src)} {shlex.quote(dst)}"
+            )
+
+    flattened = "\n".join(out)
+    _reject_unflattened_copy_from(flattened)
+    return flattened
+
+
+def _prepare_dockerfile(content: str) -> str:
+    """Rewrite an OpenEnv Dockerfile into what Novita's template parser accepts.
+
+    Applies, in order: BuildKit ``--mount`` stripping, ``ARG``/``--platform``
+    resolution in ``FROM``, then multi-stage flattening. Each transform
+    compensates for a parser limitation rather than changing the build's
+    meaning.
+    """
+    return _flatten_multistage(_resolve_from_references(_strip_mount_flags(content)))
+
+
+def _require_secure_url(url: str) -> str:
+    """Enforce https/wss transport (RFC 002 security invariant S1).
+
+    ``EnvClient`` derives its WebSocket URL from this base URL, so a plaintext
+    URL would become a cleartext ``ws://`` connection. The offending URL is
+    deliberately omitted from the error: a sandbox host is reachable only
+    through the account that owns it, so the address is treated as
+    account-scoped rather than pasted into logs.
+    """
+    if (
+        not isinstance(url, str)
+        or not url.lower().startswith("https://")
+        or "://" in url[len("https://") :]
+    ):
+        raise RuntimeError(
+            "Novita sandbox returned a non-HTTPS host URL. OpenEnv requires an "
+            "https/wss base_url so EnvClient traffic is encrypted. Refusing to "
+            "connect over plaintext."
+        )
+
+    # With NOVITA_DEBUG set, the SDK's `get_host` returns `localhost:<port>`
+    # instead of the sandbox host. Adding the scheme would yield a URL that
+    # looks valid and silently targets whatever is listening on the caller's
+    # machine, so reject it explicitly rather than connect to the wrong thing.
+    hostname = url.split("://", 1)[1].split("/", 1)[0].rsplit(":", 1)[0]
+    if hostname in ("localhost", "127.0.0.1", "::1"):
+        raise RuntimeError(
+            "Novita sandbox returned a loopback host URL, which means the SDK is "
+            "running with NOVITA_DEBUG enabled. Debug mode bypasses the sandbox "
+            "host and would connect to the local machine instead; unset "
+            "NOVITA_DEBUG to reach the sandbox."
+        )
+
+    return url
+
+
+def _require_bare_host(host: str) -> str:
+    """Require the SDK host value to contain no URL scheme."""
+    if not isinstance(host, str) or not host or "://" in host:
+        raise RuntimeError(
+            "Novita sandbox returned an invalid host. Expected a bare hostname "
+            "without a URL scheme."
+        )
+    return host
+
+
+def _raise_install_error(exc: ImportError) -> None:
+    raise RuntimeError(
+        "Novita sandbox support requires optional dependencies. "
+        "Install them with `pip install openenv[novita]`."
+    ) from exc
+
+
+class _DefaultNovitaAdapter:
+    """Thin adapter over the ``novita-sandbox`` SDK.
+
+    The provider talks to this private adapter instead of spreading SDK details
+    through its own logic; tests inject a duck-typed fake in its place. Keeping
+    the SDK surface here is what makes the provider testable without a Novita
+    account and localizes API churn (RFC 002 implementation-hygiene invariant).
+    """
+
+    def __init__(
+        self,
+        *,
+        api_key: Optional[str],
+        domain: Optional[str],
+    ):
+        try:
+            from novita_sandbox import Novita
+        except ImportError as exc:  # pragma: no cover - exercised via provider tests
+            _raise_install_error(exc)
+
+        # `api_key`/`domain` fall back to NOVITA_API_KEY / NOVITA_DOMAIN inside
+        # the SDK's ConnectionConfig, so `None` is the documented "use the env
+        # var" spelling rather than a missing value.
+        self._novita = Novita(api_key=api_key, domain=domain)
+
+    def create_sandbox(
+        self,
+        *,
+        image: Optional[str],
+        template: Optional[str],
+        env_vars: Optional[Dict[str, str]],
+        timeout: int,
+        metadata: Optional[Dict[str, str]],
+        secure: Optional[bool],
+    ) -> Any:
+        # Exactly one source: a registry image (resolved and cached by the SDK
+        # itself) or an already-built template id from `image_from_dockerfile`.
+        kwargs: Dict[str, Any] = {
+            "timeout": timeout,
+        }
+
+        if template is not None:
+            kwargs["template"] = template
+        else:
+            from novita_sandbox import wait_for_timeout
+
+            # `image` is an OCI reference; the SDK resolves it to a template
+            # itself (Template.from_image -> Template.build) and caches the
+            # result by image fingerprint, so repeated starts do not rebuild.
+            #
+            # The `build` block exists solely to pin the template's start
+            # command to a keepalive. `from_image` defaults to
+            # `inherit_config=True`, which restores the image's ENV/WORKDIR
+            # *and* its ENTRYPOINT/CMD as the template start command -- for an
+            # OpenEnv image that CMD is the uvicorn server, so leaving it in
+            # place would bind port 8000 before this provider launches its own
+            # copy. Setting `start_cmd` here wins over the inherited value (the
+            # caller always overrides the image). The SDK attempts to restore
+            # the image's ENV, but callers must use explicit paths when a
+            # registry config cannot be read during OCI template conversion.
+            ready_cmd = wait_for_timeout(_KEEPALIVE_READY_MS)
+            kwargs["image"] = image
+            kwargs["build"] = {
+                "cmd": _KEEPALIVE_CMD,
+                # Sandbox.create fingerprints this mapping with json.dumps
+                # before converting it into a TemplateBuilder. Keep the
+                # serialized command here; Template.set_start_cmd accepts the
+                # string as well as the ReadyCmd helper.
+                "ready_cmd": ready_cmd.get_cmd(),
+            }
+
+        if env_vars:
+            kwargs["envs"] = dict(env_vars)
+        if metadata:
+            kwargs["metadata"] = dict(metadata)
+        if secure is not None:
+            kwargs["secure"] = secure
+
+        return self._novita.sandbox.create(**kwargs)
+
+    def exec(
+        self,
+        sandbox: Any,
+        command: str,
+        *,
+        timeout: float = 10,
+        background: bool = False,
+        user: Optional[str] = None,
+    ) -> str:
+        """Run *command* through a login shell inside *sandbox*, returning stdout.
+
+        The SDK runs commands as ``/bin/bash -l -c <command>`` and raises
+        ``CommandExitException`` on a non-zero exit. The discovery and liveness
+        probes below are *expected* to exit non-zero ("no openenv.yaml", "process
+        is dead"), so the exit code is swallowed and stdout returned; callers
+        distinguish states by the marker the command echoes, not by the status.
+        """
+        from novita_sandbox import CommandExitException
+
+        run_kwargs: Dict[str, Any] = {
+            "timeout": timeout,
+            "background": background,
+        }
+        if user is not None:
+            run_kwargs["user"] = user
+
+        try:
+            result = sandbox.commands.run(command, **run_kwargs)
+        except CommandExitException as exc:
+            return exc.stdout or ""
+        except ImportError:  # pragma: no cover - SDK always exports this
+            result = sandbox.commands.run(command, **run_kwargs)
+
+        return getattr(result, "stdout", "") or ""
+
+    def host(self, sandbox: Any, port: int) -> str:
+        """Return the public host (no scheme) exposing *port* on *sandbox*."""
+        return str(sandbox.get_host(port))
+
+    def kill(self, sandbox: Any) -> None:
+        sandbox.kill()
+
+    def build_template(
+        self,
+        *,
+        dockerfile_content: str,
+        context_dir: str,
+        name: str,
+        cpu_count: int,
+        memory_mb: int,
+        on_build_logs: Optional[Callable[[Any], None]],
+    ) -> str:
+        """Build a Novita template from Dockerfile *content* and return its id.
+
+        Uses ``Template.from_dockerfile`` to translate the Dockerfile into
+        builder instructions, pins the template's start command to a keepalive
+        (so the image's own ``CMD`` does not bind port 8000 before the provider
+        launches the server), and builds it. ``file_context_path`` is what makes
+        relative ``COPY`` sources resolve against the build context.
+
+        Returns:
+            `str`: The built template id, to pass to ``Sandbox.create``.
+        """
+        from novita_sandbox import Template, wait_for_timeout
+
+        template = Template(file_context_path=context_dir)
+        builder = template.from_dockerfile(dockerfile_content)
+
+        # The parser rewrites USER to "user" when the Dockerfile declares none
+        # (dockerfile_parser.py: `if not user_changed: set_user("user")`). For an
+        # OpenEnv image that is wrong: `WORKDIR /app`, created by root, stays
+        # root-owned, so the server and every command the environment later runs
+        # in-process (`TB2_MODE=local`, and any env doing the same) fail with
+        # "Permission denied" on the install tree. Preserve an explicit Dockerfile
+        # USER, though, so the provider does not override the image author's
+        # requested runtime identity.
+        if not _dockerfile_declares_user(dockerfile_content):
+            builder.set_user(_SANDBOX_USER)
+
+        # `set_start_cmd` is the only way to pin the start command and it
+        # requires a readiness check, so the keepalive gets a short timer. The
+        # provider's own `wait_for_ready` does the real health polling.
+        builder.set_start_cmd(_KEEPALIVE_CMD, wait_for_timeout(_KEEPALIVE_READY_MS))
+
+        build_kwargs: Dict[str, Any] = {
+            "cpu_count": cpu_count,
+            "memory_mb": memory_mb,
+        }
+        if on_build_logs is not None:
+            build_kwargs["on_build_logs"] = on_build_logs
+
+        # Build through the configured namespace so the SDK merges the
+        # api_key/domain passed to Novita(...) into Template.build's connection
+        # config. Calling Template.build directly would fall back to an empty
+        # ConnectionConfig when those values are not present in the environment.
+        info = self._novita.template.build(builder, name, **build_kwargs)
+        return str(info.template_id)
+
+
+class NovitaSandboxProvider(ContainerProvider):
+    """
+    Container provider that runs environments in Novita AI sandboxes.
+
+    ``start_container`` accepts either form of source:
+
+    - A Docker/OCI registry reference (e.g. ``"ghcr.io/org/env:tag"``), which is
+      the form the ``ContainerProvider`` contract specifies. The SDK resolves it
+      into a Novita template internally and caches it by image fingerprint.
+    - A ``"template:<id>"`` reference returned by
+      :meth:`image_from_dockerfile`, which builds a template from a local
+      Dockerfile.
+
+    :meth:`image_from_dockerfile` rewrites the Dockerfile for Novita's template
+    parser, which does not accept multi-stage build definitions. The rewrite is
+    mechanical and preserves the build's meaning: BuildKit ``--mount`` flags are
+    stripped, ``ARG``/``--platform`` in ``FROM`` lines are resolved, and a
+    two-stage build whose stages share one base image is replayed as a single
+    stage (see `_prepare_dockerfile`). A Dockerfile that does not fit those
+    rules raises `ValueError` with the registry route as the alternative.
+
+    The environment runs untrusted code, so the provider is secure by default: it
+    enforces https/wss transport (S1) and never surfaces raw sandbox output
+    unless ``surface_server_logs=True`` (S4).
+
+    Only one sandbox is active per provider: calling ``start_container`` again
+    before ``stop_container()``/``close()`` raises ``RuntimeError`` rather than
+    orphaning the running sandbox.
+
+    Examples:
+
+        ```python
+        # From a pre-built registry image
+        with NovitaSandboxProvider(image="ghcr.io/org/echo-env:latest") as provider:
+            base_url = provider.start_container()
+            provider.wait_for_ready(base_url)
+
+        # From a local Dockerfile (builds a template on first use)
+        image = NovitaSandboxProvider.image_from_dockerfile(
+            "envs/echo_env/server/Dockerfile"
+        )
+        with NovitaSandboxProvider(image=image) as provider:
+            base_url = provider.start_container()
+        # sandbox killed on exit
+        ```
+    """
+
+    # Prepared Dockerfile content, keyed by "template:<abs_path>". Populated by
+    # `image_from_dockerfile` so `start_container` can build the template later,
+    # mirroring DaytonaProvider / ModalProvider.
+    _dockerfile_registry: Dict[str, Dict[str, Any]] = {}
+
+    @classmethod
+    def image_from_dockerfile(
+        cls,
+        dockerfile_path: str,
+        context_dir: Optional[str] = None,
+    ) -> str:
+        """Validate a Dockerfile and return a ``template:`` reference for
+        :meth:`start_container`.
+
+        Eagerly validates the Dockerfile (existence, COPY sources) and stores
+        the rewritten content in an internal registry. The Novita template is
+        built later, inside ``start_container``, by the adapter's
+        ``build_template`` — building needs credentials and network, which this
+        class-level helper deliberately does not.
+
+        The rewrite compensates for Novita's Dockerfile parser, which rejects
+        multi-stage build definitions (the layout every in-repo OpenEnv
+        environment uses). BuildKit ``--mount`` flags are stripped, ``ARG`` and
+        ``--platform`` in ``FROM`` lines are resolved, and a two-stage build
+        whose stages share one base image is replayed as a single stage.
+
+        Args:
+            dockerfile_path (`str`):
+                Path to the Dockerfile on disk.
+            context_dir (`str`, *optional*):
+                Build context directory, used to resolve relative ``COPY``
+                sources. Defaults to the Dockerfile's grandparent directory,
+                matching the `openenv init` convention where Dockerfiles live in
+                `<env>/server/Dockerfile` and the build context is `<env>/`.
+
+        Returns:
+            `str`: A `"template:<abs_path>"` reference to pass to
+            `start_container` or the constructor's `image`.
+
+        Raises:
+            FileNotFoundError: If *dockerfile_path* does not exist.
+            ValueError: If *context_dir* does not exist, if COPY sources cannot
+                be found under the resolved context directory, or if the
+                Dockerfile cannot be expressed as a Novita template (see
+                `_prepare_dockerfile`).
+
+        Examples:
+
+        ```python
+        image = NovitaSandboxProvider.image_from_dockerfile(
+            "envs/echo_env/server/Dockerfile"
+        )
+        provider = NovitaSandboxProvider(image=image)
+        base_url = provider.start_container()
+        ```
+        """
+        import pathlib
+
+        src = pathlib.Path(dockerfile_path).resolve()
+        if not src.is_file():
+            raise FileNotFoundError(f"Dockerfile not found: {dockerfile_path}")
+
+        if context_dir is not None:
+            ctx = pathlib.Path(context_dir)
+            if not ctx.is_dir():
+                raise ValueError(f"context_dir does not exist: {context_dir}")
+        else:
+            # Default: grandparent of the Dockerfile, matching the openenv init
+            # layout (<env>/server/Dockerfile -> <env>/).
+            ctx = src.parent.parent
+
+        prepared = _prepare_dockerfile(src.read_text())
+
+        # Validate that COPY sources exist under the context directory. This
+        # catches mismatches early (e.g. a Dockerfile expecting the repo root as
+        # context when we defaulted to the env directory).
+        for line in prepared.splitlines():
+            match = re.match(r"^\s*COPY\s+(?!--from=)(\S+)\s+", line, re.IGNORECASE)
+            if not match:
+                continue
+            copy_src = match.group(1)
+            if copy_src.startswith("/"):
+                continue
+            resolved = ctx / copy_src
+            if not resolved.exists() and not any(ctx.glob(copy_src)):
+                raise ValueError(
+                    f"Dockerfile COPY source '{copy_src}' not found under "
+                    f"context_dir '{ctx}'. This Dockerfile may expect a "
+                    "different build context (e.g. the repo root). Pass "
+                    "context_dir explicitly."
+                )
+
+        cls._dockerfile_registry[str(src)] = {
+            "content": prepared,
+            "context_dir": str(ctx),
+        }
+
+        return f"template:{src}"
+
+    def _resolve_template_ref(self, reference: str) -> str:
+        """Build the template for a ``template:<path>`` reference; return its id."""
+        path = reference[len("template:") :]
+        meta = self._dockerfile_registry.get(path)
+        if meta is None:
+            raise ValueError(
+                f"No registered Dockerfile metadata for {path}. Call "
+                "NovitaSandboxProvider.image_from_dockerfile() first."
+            )
+
+        return self._adapter.build_template(
+            dockerfile_content=meta["content"],
+            context_dir=meta["context_dir"],
+            name=self._template_name or self._default_template_name(path),
+            cpu_count=self._cpu_count,
+            memory_mb=self._memory_mb,
+            on_build_logs=self._on_build_logs,
+        )
+
+    @staticmethod
+    def _default_template_name(dockerfile_path: str) -> str:
+        """Derive a stable template name from the Dockerfile's env directory.
+
+        The env directory keeps the template recognizable in the Novita
+        dashboard; the content hash keeps two different Dockerfiles from
+        colliding on one name, while a rebuild of identical content still hits
+        the build cache.
+        """
+        import hashlib
+        import pathlib
+
+        src = pathlib.Path(dockerfile_path)
+        env_dir = src.parent.parent.name or src.parent.name
+        slug = re.sub(r"[^a-z0-9-]+", "-", env_dir.lower().replace("_", "-"))
+        digest = hashlib.sha256(src.read_bytes()).hexdigest()[:8]
+        return f"openenv-{slug.strip('-')}-{digest}"
+
+    def __init__(
+        self,
+        *,
+        image: Optional[str] = None,
+        env_vars: Optional[Dict[str, str]] = None,
+        api_key: Optional[str] = None,
+        domain: Optional[str] = None,
+        timeout: int = _DEFAULT_SANDBOX_TIMEOUT_S,
+        metadata: Optional[Dict[str, str]] = None,
+        secure: Optional[bool] = None,
+        cmd: Optional[str] = None,
+        working_directory: Optional[str] = None,
+        surface_server_logs: bool = False,
+        cpu_count: int = _DEFAULT_TEMPLATE_CPU,
+        memory_mb: int = _DEFAULT_TEMPLATE_MEMORY_MB,
+        template_name: Optional[str] = None,
+        on_build_logs: Optional[Callable[[Any], None]] = None,
+        _adapter: Any = None,
+    ):
+        """
+        Args:
+            image (`str`, *optional*):
+                Registry image reference (e.g. `"ghcr.io/org/env:latest"`) or a
+                `"template:<id>"` reference from
+                [`~image_from_dockerfile`], to use when `start_container()` is
+                called without an image.
+            env_vars (`dict`, *optional*):
+                Environment variables to use when `start_container()` is called
+                without explicit `env_vars`.
+            api_key (`str`, *optional*):
+                Novita API key. Falls back to the `NOVITA_API_KEY` environment
+                variable.
+            domain (`str`, *optional*):
+                Novita region domain. Falls back to the `NOVITA_DOMAIN`
+                environment variable, then to the SDK default (`us-phx-1`).
+                Features marked modern-domain-only (snapshots, secrets) are
+                unavailable on the legacy `us-virginia-1` domain.
+            timeout (`int`, *optional*, defaults to `3600`):
+                Hard sandbox lifetime in seconds. Counts down from creation
+                regardless of activity and always fires, so it must exceed the
+                longest expected episode.
+            metadata (`dict`, *optional*):
+                Arbitrary key/value labels on the sandbox. The `idle_timeout`
+                key is meaningful to Novita and takes seconds as a **string**
+                (e.g. `{"idle_timeout": "900"}`); omitting it disables the idle
+                timer, which is the conservative default for RL rollouts — an
+                automatic transition can drop a live WebSocket mid-episode
+                (RFC 002 invariant 6).
+            secure (`bool`, *optional*):
+                Whether envd requires its access token. `None` uses the SDK
+                default (secured on modern domains). This governs the sandbox
+                control plane, not the exposed server port.
+            cmd (`str`, *optional*):
+                Shell command to start the server inside the sandbox. When
+                omitted, the command is auto-discovered from `openenv.yaml`.
+            working_directory (`str`, *optional*):
+                Directory to `cd` into before running `cmd`.
+            surface_server_logs (`bool`, *optional*, defaults to `False`):
+                When `False` (default), captured sandbox output is withheld from
+                raised errors so secrets the workload printed cannot leak into
+                orchestrator/CI logs. When `True`, a best-effort redacted,
+                length-bounded excerpt is included in startup-crash errors (S4).
+            cpu_count (`int`, *optional*, defaults to `2`):
+                vCPU cores for sandboxes built from a Dockerfile via
+                [`~image_from_dockerfile`]. Ignored for registry images, whose
+                resources come from the image's own template.
+            memory_mb (`int`, *optional*, defaults to `1024`):
+                Memory in MiB for sandboxes built from a Dockerfile. Ignored for
+                registry images.
+            template_name (`str`, *optional*):
+                Template name to register when building from a Dockerfile.
+                Defaults to a name derived from the Dockerfile's directory plus a
+                content hash, so re-running with the same Dockerfile reuses the
+                build cache.
+            on_build_logs (`Callable`, *optional*):
+                Callback receiving build-log entries while a Dockerfile template
+                is built. Off by default: build logs are not redacted, so they
+                are withheld rather than printed (RFC 002 S4). Pass a callable to
+                surface them.
+            _adapter (`Any`, *optional*):
+                Injection seam for tests; a duck-typed fake replaces the SDK.
+        """
+        self._image = image
+        self._env_vars = env_vars
+        self._timeout = timeout
+        self._metadata = metadata
+        self._secure = secure
+        self._cmd = cmd
+        self._working_directory = working_directory
+        self.surface_server_logs = surface_server_logs
+        self._cpu_count = cpu_count
+        self._memory_mb = memory_mb
+        self._template_name = template_name
+        self._on_build_logs = on_build_logs
+
+        self._sandbox: Any = None
+        self._base_url: Optional[str] = None
+        self._server_user: Optional[str] = None
+        # Injected env-var values, used to scrub captured server output before
+        # it is ever surfaced in an error.
+        self._redact_values: set[str] = set()
+
+        if _adapter is None:
+            # Build the client eagerly so credential/region problems surface at
+            # construction rather than mid-episode, matching ModalProvider.
+            self._adapter: Any = _DefaultNovitaAdapter(api_key=api_key, domain=domain)
+        else:
+            self._adapter = _adapter
+
+    def _discover_server_cmd(self, port: int = _DEFAULT_NOVITA_PORT) -> str:
+        """Discover the server command from ``openenv.yaml`` inside the sandbox.
+
+        Finds the file, reads the ``app`` field, and constructs a command of the
+        form ``cd <env_root> && python -m uvicorn <app> --host 0.0.0.0 --port <port>``.
+
+        Raises:
+            ValueError: If ``openenv.yaml`` is not found or lacks an ``app`` field.
+        """
+        yaml_path = self._find_openenv_yaml()
+        if yaml_path is None:
+            raise ValueError(
+                "Could not find openenv.yaml inside the sandbox. Pass an "
+                "explicit cmd= to NovitaSandboxProvider or start_container()."
+            )
+
+        content = self._adapter.exec(self._sandbox, f"cat {shlex.quote(yaml_path)}")
+        app = self._parse_app_field(content)
+        if app is None:
+            raise ValueError(
+                f"openenv.yaml at {yaml_path} does not contain an 'app' field. "
+                "Pass an explicit cmd= to NovitaSandboxProvider or start_container()."
+            )
+
+        # The directory containing openenv.yaml is the env root
+        env_root = yaml_path.rsplit("/", 1)[0]
+        return (
+            f"cd {shlex.quote(env_root)} && "
+            f"python -m uvicorn {shlex.quote(app)} --host 0.0.0.0 --port {port}"
+        )
+
+    def _find_openenv_yaml(self) -> Optional[str]:
+        """Locate ``openenv.yaml`` inside the sandbox.
+
+        Tries the modern layout path ``/app/env/openenv.yaml`` first, then falls
+        back to a ``find`` command for the old layout.
+        """
+        # Fast path: modern Dockerfile layout
+        out = self._adapter.exec(
+            self._sandbox, "test -f /app/env/openenv.yaml && echo found"
+        )
+        if "found" in (out or ""):
+            return "/app/env/openenv.yaml"
+
+        # Fallback: search for it (redirect stderr so error messages like
+        # "No such file or directory" don't get mistaken for paths).
+        path = self._adapter.exec(
+            self._sandbox,
+            "find /app -maxdepth 4 -name openenv.yaml -print -quit 2>/dev/null",
+        ).strip()
+        if path and path.startswith("/"):
+            return path
+
+        return None
+
+    @staticmethod
+    def _parse_app_field(yaml_content: str) -> Optional[str]:
+        """Extract the ``app`` value from raw openenv.yaml content.
+
+        Uses PyYAML to handle comments, quotes, and nested keys correctly.
+        """
+        return parse_openenv_app_field(yaml_content)
+
+    def start_container(
+        self,
+        image: Optional[str] = None,
+        port: Optional[int] = None,
+        env_vars: Optional[Dict[str, str]] = None,
+        **kwargs: Any,
+    ) -> str:
+        """
+        Create a Novita sandbox and start an OpenEnv server inside it.
+
+        Args:
+            image (`str`, *optional*):
+                Either a registry image reference (e.g.
+                `"ghcr.io/org/env:latest"`) or a `"template:<path>"` reference
+                returned by [`~image_from_dockerfile`]. May be omitted when
+                supplied to the constructor.
+            port (`int`, *optional*):
+                Must be `None` or `8000`. Novita exposes port 8000 on the
+                sandbox host; other ports raise `ValueError`.
+            env_vars (`dict`, *optional*):
+                Environment variables forwarded to the sandbox, overriding the
+                constructor's.
+            **kwargs:
+                `cmd` (`str`) to override the server command. For a registry
+                OCI image, pass the command explicitly together with any
+                runtime environment the image needs: Novita may not be able to
+                restore the image's ``ENV`` while converting it to a template.
+                For example, the tbench2 image requires:
+
+                ``cmd="cd /app/env && PYTHONPATH=/app/env "
+                "/app/.venv/bin/python -m uvicorn server.app:app "
+                "--host 0.0.0.0 --port 8000"``
+
+                and any additional values can be supplied through
+                ``env_vars``. The provider runs this command as ``root`` for a
+                registry image; a ``template:<id>`` keeps the template's
+                configured user. Prefer
+                ``NovitaSandboxProvider.image_from_dockerfile(...)`` when a
+                local Dockerfile is available; its build preserves the
+                environment and startup configuration in the generated
+                template. Unknown options raise `ValueError` so typos cannot
+                silently change sandbox behavior.
+
+        Returns:
+            `str`: HTTPS sandbox URL for the exposed port (base_url).
+
+        Raises:
+            RuntimeError: If a sandbox is already active on this provider.
+            ValueError: If no image is available, the port is unsupported, or an
+                unknown option is passed.
+        """
+        if self._sandbox is not None:
+            raise RuntimeError(
+                "NovitaSandboxProvider already has an active sandbox. Call "
+                "stop_container() (or close()) before starting another — a "
+                "second start would orphan the running sandbox."
+            )
+
+        if port is not None and port != _DEFAULT_NOVITA_PORT:
+            raise ValueError(
+                f"NovitaSandboxProvider only supports port {_DEFAULT_NOVITA_PORT} "
+                f"(got {port}). The sandbox host routes to port "
+                f"{_DEFAULT_NOVITA_PORT} inside the sandbox."
+            )
+
+        effective_image = image if image is not None else self._image
+        if effective_image is None:
+            raise ValueError(
+                "NovitaSandboxProvider requires an image. Pass it to the "
+                "constructor or start_container()."
+            )
+        effective_env_vars = self._env_vars if env_vars is None else env_vars
+
+        cmd = kwargs.pop("cmd", None) or self._cmd
+        # `AutoEnv.from_env()` forwards this unconditionally (default 30.0) via
+        # `from_docker_image` -> `_bootstrap_container`, so rejecting it would
+        # make the entire AutoEnv path unusable for this provider. It is
+        # meaningless here: `_bootstrap_container` calls `wait_for_ready(base_url)`
+        # without a timeout, so no provider ever sees the value. Accept and drop
+        # it rather than fail; unknown options still raise, so a typo cannot
+        # silently change sandbox behavior.
+        kwargs.pop("wait_timeout", None)
+        if kwargs:
+            unknown = ", ".join(sorted(kwargs))
+            raise ValueError(
+                f"Unsupported NovitaSandboxProvider start options: {unknown}"
+            )
+
+        # A "template:<path>" reference builds a Novita template from a local
+        # Dockerfile; anything else is a registry image the SDK resolves itself.
+        # The build runs before `create_sandbox` so a build failure never leaves
+        # a sandbox behind, and outside the `_redact_values` window below so its
+        # recorded secrets do not outlive a failed build.
+        if effective_image.startswith("template:"):
+            template_id: Optional[str] = self._resolve_template_ref(effective_image)
+            sandbox_image: Optional[str] = None
+        else:
+            template_id = None
+            sandbox_image = effective_image
+
+        # Record injected secret values so captured server output can be scrubbed
+        # before it is ever surfaced in an error (security invariant S4).
+        self._redact_values = {
+            value for value in (effective_env_vars or {}).values() if value
+        }
+
+        # A create failure created nothing, so just drop the recorded secrets and
+        # re-raise (the single-sandbox guard above guarantees there is no
+        # pre-existing sandbox to delete).
+        try:
+            self._sandbox = self._adapter.create_sandbox(
+                image=sandbox_image,
+                template=template_id,
+                env_vars=effective_env_vars,
+                timeout=self._timeout,
+                metadata=self._metadata,
+                secure=self._secure,
+            )
+        except Exception:
+            self._redact_values = set()
+            raise
+
+        try:
+            # Discovery runs after creation, when the filesystem can be read.
+            if cmd is None:
+                cmd = self._discover_server_cmd()
+
+            # Registry images can lose their Dockerfile USER/ownership settings
+            # while Novita converts them to a cached OCI template. Run the
+            # provider-managed server command as root for that path. A template
+            # keeps its configured user, including an explicit non-root USER.
+            self._server_user = _SANDBOX_USER if sandbox_image is not None else None
+            self._launch_server(
+                cmd,
+                user=self._server_user,
+            )
+            host = _require_bare_host(
+                self._adapter.host(self._sandbox, _DEFAULT_NOVITA_PORT)
+            )
+            self._base_url = _require_secure_url(f"https://{host}")
+        except Exception:
+            # A cleanup failure here must not mask the original error: swallow
+            # any exception from stop_container() so the root cause propagates.
+            try:
+                self.stop_container()
+            except Exception:
+                pass
+            raise
+
+        return self._base_url
+
+    def _launch_server(self, cmd: str, *, user: Optional[str] = None) -> None:
+        """Start the OpenEnv server inside the sandbox as a background process.
+
+        This is the one place the provider assumes something about the image, so
+        it is kept to a single method.
+
+        The server is launched through the SDK's background command support and
+        its PID is written to ``/tmp/openenv-server.pid``, so ``wait_for_ready``
+        can tell "still booting" apart from "crashed" instead of waiting out
+        the full timeout.
+        The sandbox was created with a ``sleep infinity`` start command rather
+        than the image's own CMD (see ``_DefaultNovitaAdapter.create_sandbox``),
+        which is what leaves port 8000 free for this process.
+
+        An alternative is to let the platform own it: pass the server command as
+        the template's start command via the SDK's ``build={"cmd": ...}`` and
+        drop this method. That trades the PID-based crash detection
+        below for platform-managed lifecycle and readiness. It is not the
+        default because a start command can only be set when the image is first
+        resolved into a template, so it would not apply to an already-cached
+        image.
+
+        ``cmd`` is trusted orchestrator configuration, not agent or environment
+        input — callers must not pass agent-controlled text, and any dynamic
+        value interpolated into it must be quoted by the caller (S5). The
+        sandbox SDK runs commands as ``/bin/bash -l -c <command>``; the command
+        itself owns the PID file and log redirection while ``background=True``
+        keeps the remote process alive after this call.
+        """
+        if self._working_directory:
+            command = f"cd {shlex.quote(self._working_directory)} && {cmd}"
+        else:
+            command = cmd
+
+        launch_command = (
+            "echo $$ > /tmp/openenv-server.pid; "
+            f"exec sh -c {shlex.quote(command)} "
+            "> /tmp/openenv-server.log 2>&1"
+        )
+        escaped = shlex.quote(launch_command)
+        self._adapter.exec(
+            self._sandbox,
+            f"sh -c {escaped}",
+            timeout=0,
+            background=True,
+            user=user,
+        )
+
+    def stop_container(self) -> None:
+        """Kill the active Novita sandbox.
+
+        `kill()` permanently removes the sandbox: it cannot be resumed or
+        reconnected, and any snapshot must have been taken beforehand.
+        """
+        if self._sandbox is None:
+            # Still drop any injected secret values recorded by a failed start.
+            self._redact_values = set()
+            return
+
+        sandbox = self._sandbox
+        try:
+            self._adapter.kill(sandbox)
+        finally:
+            self._sandbox = None
+            self._base_url = None
+            self._server_user = None
+            self._redact_values = set()
+
+    def close(self) -> None:
+        """Stop the active sandbox and release provider-held resources.
+
+        Overrides the base no-op so a caller holding a bare `ContainerProvider`
+        reference can release the sandbox polymorphically (also invoked on
+        context-manager exit). The Novita client is stateless HTTP per call, so
+        this is equivalent to `stop_container()`.
+        """
+        self.stop_container()
+
+    @property
+    def base_url(self) -> str:
+        """URL returned by the last `start_container`."""
+        if self._base_url is None:
+            raise RuntimeError(
+                "NovitaSandboxProvider has no active base_url. Start the "
+                "provider before reading base_url."
+            )
+        return self._base_url
+
+    def _redact(self, text: str, *, max_chars: int = 2000) -> str:
+        """Scrub injected secret values and bound length before surfacing output.
+
+        Replaces any injected env-var value with `***` and keeps only the tail.
+        This is best-effort (exact-match only), which is why server output is
+        withheld entirely unless `surface_server_logs=True` (RFC 002 S4).
+        """
+        redacted = text or ""
+        for value in self._redact_values:
+            redacted = redacted.replace(value, "***")
+        if len(redacted) > max_chars:
+            redacted = "...(truncated)...\n" + redacted[-max_chars:]
+        return redacted
+
+    def _server_died_message(self) -> str:
+        """Build the startup-crash error, secure by default (RFC 002 S4).
+
+        Untrusted code can print secrets then force a crash to exfiltrate them
+        through the exception (which lands in orchestrator/CI logs), so sandbox
+        output is excluded unless `surface_server_logs=True`, in which case a
+        best-effort redacted, bounded excerpt is included.
+        """
+        base = (
+            "Novita sandbox server process died during startup. Server output is "
+            "not surfaced to avoid leaking secrets injected into the sandbox; "
+            "retrieve /tmp/openenv-server.log from the sandbox out of band, or "
+            "construct the provider with surface_server_logs=True to include a "
+            "redacted excerpt."
+        )
+        if not self.surface_server_logs:
+            return base
+        if self._sandbox is None:
+            return base
+
+        log = self._redact(
+            self._adapter.exec(
+                self._sandbox,
+                "cat /tmp/openenv-server.log 2>/dev/null",
+                user=self._server_user,
+            )
+        )
+        return (
+            "Novita sandbox server process died during startup. The excerpt below "
+            "is the sandbox server output with injected secret values redacted "
+            "(best-effort); it may still contain secrets the workload printed "
+            f"by other means.\nLog (redacted):\n{log}"
+        )
+
+    def wait_for_ready(self, base_url: str, timeout_s: float = 120.0) -> None:
+        """
+        Poll the /health endpoint until the sandbox is ready.
+
+        Uses a longer default timeout (120s) than local Docker providers because
+        a Novita sandbox is created from an image that may need to be resolved
+        and built on first use.
+
+        A `200` on `/health` proves HTTP reachability but **not** that the
+        exposed sandbox host proxies the `/ws` WebSocket upgrade `EnvClient`
+        needs; that requires a real `wss://` round-trip (RFC 002 invariant 2).
+
+        Args:
+            base_url (`str`):
+                Sandbox URL returned by `start_container()`.
+            timeout_s (`float`, *optional*, defaults to `120.0`):
+                Maximum seconds to wait.
+
+        Raises:
+            TimeoutError: If the sandbox doesn't become ready in time.
+            RuntimeError: If the server process died (detected via PID check).
+        """
+        # Imported lazily inside the method that needs it, matching the other
+        # providers in this package (e.g. DaytonaProvider).
+        import requests
+
+        deadline = time.time() + timeout_s
+        health_url = f"{base_url}/health"
+
+        while time.time() < deadline:
+            try:
+                response = requests.get(health_url, timeout=5.0)
+                if response.status_code == 200:
+                    return
+            except requests.RequestException:
+                pass
+
+            # Early exit: if the server process died, raise immediately instead
+            # of waiting for the full health-check timeout.
+            if self._sandbox is not None:
+                out = self._adapter.exec(
+                    self._sandbox,
+                    "if test -s /tmp/openenv-server.pid; then "
+                    "kill -0 $(cat /tmp/openenv-server.pid) 2>/dev/null "
+                    "&& echo RUNNING || echo DEAD; "
+                    "else echo STARTING; fi",
+                    user=self._server_user,
+                )
+                if "DEAD" in (out or ""):
+                    raise RuntimeError(self._server_died_message())
+
+            time.sleep(1.0)
+
+        raise TimeoutError(f"Novita sandbox did not become ready within {timeout_s}s.")
+
+
+__all__ = ["NovitaSandboxProvider"]

src/core/containers/runtime/providers.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 Container provider abstractions for running environment servers.
@@ -14,7 +10,9 @@ This module provides a pluggable architecture for different container providers
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Optional, Sequence
+from typing import Any, Dict, Optional, Sequence, TypeVar
+
+_ContainerProviderT = TypeVar("_ContainerProviderT", bound="ContainerProvider")
 
 
 class ContainerProvider(ABC):
@@ -30,12 +28,15 @@ class ContainerProvider(ABC):
     The provider manages a single container lifecycle and provides the base URL
     for connecting to it.
 
-    Example:
-        >>> provider = LocalDockerProvider()
-        >>> base_url = provider.start_container("echo-env:latest")
-        >>> print(base_url)  # http://localhost:8000
-        >>> # Use the environment via base_url
-        >>> provider.stop_container()
+    Examples:
+
+        ```python
+        provider = LocalDockerProvider()
+        base_url = provider.start_container("echo-env:latest")
+        print(base_url)  # http://localhost:8000
+        # Use the environment via base_url
+        provider.stop_container()
+        ```
     """
 
     @abstractmethod
@@ -50,16 +51,23 @@ class ContainerProvider(ABC):
         Start a container from the specified image.
 
         Args:
-            image: Container image name (e.g., "echo-env:latest")
-            port: Port to expose (if None, provider chooses)
-            env_vars: Environment variables to pass to container
-            **kwargs: Provider-specific options
+            image (`str`):
+                Provider-specific container *source* identifier. For
+                container-based providers this is a registry image name (e.g.
+                `"echo-env:latest"`); other providers may map it to a
+                provider-specific source (see the provider's documentation).
+            port (`int`, *optional*):
+                Port to expose. If `None`, the provider chooses.
+            env_vars (`dict`, *optional*):
+                Environment variables to pass to container.
+            **kwargs:
+                Provider-specific options.
 
         Returns:
-            Base URL to connect to the container (e.g., "http://localhost:8000")
+            `str`: Base URL to connect to the container (e.g., `"http://localhost:8000"`).
 
         Raises:
-            RuntimeError: If container fails to start
+            RuntimeError: If container fails to start.
         """
         pass
 
@@ -80,14 +88,34 @@ class ContainerProvider(ABC):
         This typically polls the /health endpoint until it returns 200.
 
         Args:
-            base_url: Base URL of the container
-            timeout_s: Maximum time to wait
+            base_url (`str`):
+                Base URL of the container.
+            timeout_s (`float`, *optional*, defaults to `30.0`):
+                Maximum time to wait in seconds.
 
         Raises:
-            TimeoutError: If container doesn't become ready in time
+            TimeoutError: If container doesn't become ready in time.
         """
         pass
 
+    def close(self) -> None:
+        """
+        Release provider-held resources (e.g. SDK clients, connections).
+
+        Defaults to a no-op so existing providers are unaffected. Providers that
+        hold external resources beyond the container itself (such as a cloud SDK
+        client) override this to release them; it is also invoked on context-
+        manager exit. Lightweight providers need not override it.
+        """
+        pass
+
+    def __enter__(self: _ContainerProviderT) -> _ContainerProviderT:
+        return self
+
+    def __exit__(self, exc_type, exc, tb) -> Optional[bool]:
+        self.close()
+        return None
+
 
 class LocalDockerProvider(ContainerProvider):
     """
@@ -96,11 +124,14 @@ class LocalDockerProvider(ContainerProvider):
     This provider runs containers on the local machine using Docker.
     Useful for development and testing.
 
-    Example:
-        >>> provider = LocalDockerProvider()
-        >>> base_url = provider.start_container("echo-env:latest")
-        >>> # Container running on http://localhost:<random-port>
-        >>> provider.stop_container()
+    Examples:
+
+        ```python
+        provider = LocalDockerProvider()
+        base_url = provider.start_container("echo-env:latest")
+        # Container running on http://localhost:<random-port>
+        provider.stop_container()
+        ```
     """
 
     def __init__(self):
@@ -138,13 +169,17 @@ class LocalDockerProvider(ContainerProvider):
         Start a Docker container locally.
 
         Args:
-            image: Docker image name
-            port: Port to expose (if None, finds available port)
-            env_vars: Environment variables for the container
-            **kwargs: Additional Docker run options
+            image (`str`):
+                Docker image name.
+            port (`int`, *optional*):
+                Port to expose. If `None`, finds an available port.
+            env_vars (`dict`, *optional*):
+                Environment variables for the container.
+            **kwargs:
+                Additional Docker run options.
 
         Returns:
-            Base URL to connect to the container
+            `str`: Base URL to connect to the container.
         """
         import subprocess
         import time
@@ -226,11 +261,13 @@ class LocalDockerProvider(ContainerProvider):
         Wait for container to be ready by polling /health endpoint.
 
         Args:
-            base_url: Base URL of the container
-            timeout_s: Maximum time to wait
+            base_url (`str`):
+                Base URL of the container.
+            timeout_s (`float`, *optional*, defaults to `30.0`):
+                Maximum time to wait in seconds.
 
         Raises:
-            TimeoutError: If container doesn't become ready
+            TimeoutError: If container doesn't become ready.
         """
         import time
 
@@ -261,7 +298,7 @@ class LocalDockerProvider(ContainerProvider):
         Find an available port on localhost.
 
         Returns:
-            An available port number
+            `int`: An available port number.
         """
         import socket
 
@@ -276,10 +313,11 @@ class LocalDockerProvider(ContainerProvider):
         Generate a unique container name based on image name and timestamp.
 
         Args:
-            image: Docker image name
+            image (`str`):
+                Docker image name.
 
         Returns:
-            A unique container name
+            `str`: A unique container name.
         """
         import time
 
@@ -306,12 +344,13 @@ class DockerSwarmProvider(ContainerProvider):
     ):
         """
         Args:
-            auto_init_swarm: Whether to call ``docker swarm init`` when Swarm
-                is not active. Otherwise, user must manually initialize Swarm.
-            overlay_network: Optional overlay network name for the service.
-                When provided, the network is created with
-                ``docker network create --driver overlay --attachable`` if it
-                does not already exist.
+            auto_init_swarm (`bool`, *optional*, defaults to `True`):
+                Whether to call `docker swarm init` when Swarm is not active.
+                Otherwise, the user must manually initialize Swarm.
+            overlay_network (`str`, *optional*):
+                Overlay network name for the service. When provided, the network
+                is created with `docker network create --driver overlay --attachable`
+                if it does not already exist.
         """
         self._service_name: Optional[str] = None
         self._service_id: Optional[str] = None
@@ -334,13 +373,28 @@ class DockerSwarmProvider(ContainerProvider):
         """
         Start (or scale) a Swarm service for the given image.
 
-        Supported kwargs:
-            replicas (int): Number of container replicas (default: 2).
-            cpu_limit (float | str): CPU limit passed to ``--limit-cpu``.
-            memory_limit (str): Memory limit passed to ``--limit-memory``.
-            constraints (Sequence[str]): Placement constraints.
-            labels (Dict[str, str]): Service labels.
-            command (Sequence[str] | str): Override container command.
+        Args:
+            image (`str`):
+                Docker image name.
+            port (`int`, *optional*):
+                Port to expose. If `None`, finds an available port.
+            env_vars (`dict`, *optional*):
+                Environment variables for the container.
+            replicas (`int`, *optional*, defaults to `2`):
+                Number of container replicas.
+            cpu_limit (`float` or `str`, *optional*):
+                CPU limit passed to `--limit-cpu`.
+            memory_limit (`str`, *optional*):
+                Memory limit passed to `--limit-memory`.
+            constraints (`Sequence[str]`, *optional*):
+                Placement constraints.
+            labels (`dict`, *optional*):
+                Service labels.
+            command (`Sequence[str]` or `str`, *optional*):
+                Override container command.
+
+        Returns:
+            `str`: Base URL to connect to the service.
         """
         import shlex
         import subprocess
@@ -464,7 +518,7 @@ class DockerSwarmProvider(ContainerProvider):
         """
         Wait for at least one replica to become healthy by polling /health.
 
-        Note: With Swarm's load balancer, requests round-robin across replicas,
+        With Swarm's load balancer, requests round-robin across replicas,
         so this only verifies that at least one replica is responding. Some
         replicas may still be starting when this returns.
         """
@@ -592,16 +646,12 @@ class DockerSwarmProvider(ContainerProvider):
 
 class KubernetesProvider(ContainerProvider):
     """
-    Container provider for Kubernetes clusters.
+    Planned container provider for Kubernetes clusters.
 
-    This provider creates pods in a Kubernetes cluster and exposes them
-    via services or port-forwarding.
-
-    Example:
-        >>> provider = KubernetesProvider(namespace="envtorch-dev")
-        >>> base_url = provider.start_container("echo-env:latest")
-        >>> # Pod running in k8s, accessible via service or port-forward
-        >>> provider.stop_container()
+    Not yet implemented: this is a placeholder for the planned Kubernetes
+    backend and does not implement the abstract `ContainerProvider` methods, so
+    it cannot be instantiated. Use `LocalDockerProvider`, `DockerSwarmProvider`,
+    `DaytonaProvider`, or `ACASandboxProvider` instead.
     """
 
     pass
@@ -616,11 +666,14 @@ class RuntimeProvider(ABC):
     The provider manages a single runtime lifecycle and provides the base URL
     for connecting to it.
 
-    Example:
-        >>> provider = UVProvider(project_path="/path/to/env")
-        >>> base_url = provider.start()
-        >>> print(base_url)  # http://localhost:8000
-        >>> provider.stop()
+    Examples:
+
+        ```python
+        provider = UVProvider(project_path="/path/to/env")
+        base_url = provider.start()
+        print(base_url)  # http://localhost:8000
+        provider.stop()
+        ```
     """
 
     @abstractmethod
@@ -631,13 +684,18 @@ class RuntimeProvider(ABC):
         **kwargs: Any,
     ) -> str:
         """
-        Start a runtime from the specified image.
+        Start the runtime.
 
         Args:
-            image: Runtime image name
-            port: Port to expose (if None, provider chooses)
-            env_vars: Environment variables for the runtime
-            **kwargs: Additional runtime options
+            port (`int`, *optional*):
+                Port to expose. If `None`, the provider chooses.
+            env_vars (`dict`, *optional*):
+                Environment variables for the runtime.
+            **kwargs:
+                Additional runtime options.
+
+        Returns:
+            `str`: Base URL to connect to the runtime.
         """
 
     @abstractmethod

src/core/containers/runtime/uv_provider.py

@@ -3,8 +3,10 @@
 from __future__ import annotations
 
 import os
+import shutil
 import socket
 import subprocess
+import tempfile
 import time
 from typing import Dict, Optional
 
@@ -12,6 +14,8 @@ import requests
 
 from .providers import RuntimeProvider
 
+GIT_URL_PREFIX = "git+"
+
 
 def _check_uv_installed() -> None:
     try:
@@ -22,6 +26,42 @@ def _check_uv_installed() -> None:
         ) from exc
 
 
+def _clone_git_project(git_url: str, timeout_s: float) -> str:
+    """Clone a `git+<url>` spec to a temp dir and return its local path.
+
+    `uv run --project` only discovers a project in a local directory (per
+    `uv run --help`) -- it has no support for remote git sources. Callers
+    that accept a `git+...` `project_path` (e.g. `EnvClient.from_env`)
+    must therefore clone it themselves before handing a path to `uv run`.
+    """
+    repo_url = git_url[len(GIT_URL_PREFIX) :]
+    clone_dir = tempfile.mkdtemp(prefix="openenv-uv-clone-")
+    try:
+        subprocess.run(
+            ["git", "clone", "--depth", "1", repo_url, clone_dir],
+            check=True,
+            capture_output=True,
+            text=True,
+            # A hung/slow remote would otherwise block start() (and the
+            # caller's own readiness timeout) indefinitely.
+            timeout=timeout_s,
+        )
+    except FileNotFoundError as exc:
+        shutil.rmtree(clone_dir, ignore_errors=True)
+        raise RuntimeError(
+            f"`git` executable not found; required to clone project_path {git_url!r}."
+        ) from exc
+    except subprocess.CalledProcessError as exc:
+        shutil.rmtree(clone_dir, ignore_errors=True)
+        raise RuntimeError(f"Failed to clone {repo_url!r}: {exc.stderr}") from exc
+    except subprocess.TimeoutExpired as exc:
+        shutil.rmtree(clone_dir, ignore_errors=True)
+        raise RuntimeError(
+            f"Timed out cloning {repo_url!r} after {timeout_s:.1f}s"
+        ) from exc
+    return clone_dir
+
+
 def _find_free_port() -> int:
     with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
         sock.bind(("", 0))
@@ -71,7 +111,11 @@ def _poll_health(health_url: str, timeout_s: float) -> None:
             if response.status_code == 200:
                 return
         except requests.RequestException:
-            continue
+            # Server not accepting connections yet. Fall through to the sleep
+            # below rather than `continue`-ing: a refused connection returns
+            # immediately, so retrying without a pause busy-spins a CPU core
+            # for the whole timeout window while the server boots.
+            pass
 
         time.sleep(0.5)
 
@@ -83,19 +127,29 @@ class UVProvider(RuntimeProvider):
     RuntimeProvider implementation backed by ``uv run``.
 
     Args:
-        project_path: Local path to a uv project (passed to ``uv run --project``)
-        app: ASGI application path for uvicorn (defaults to ``server.app:app``)
-        host: Host interface to bind to (defaults to ``0.0.0.0``)
-        reload: Whether to enable uvicorn's reload mode
-        env_vars: Environment variables to pass through to the spawned process
-        context_timeout_s: How long to wait for the environment to become ready
-
-    Example:
-        >>> provider = UVProvider(project_path="/path/to/env")
-        >>> base_url = provider.start()
-        >>> print(base_url)  # http://localhost:8000
-        >>> # Use the environment via base_url
-        >>> provider.stop()
+        project_path (`str`):
+            Local path to a uv project (passed to `uv run --project`), or a
+            `git+<url>` spec that is cloned to a temp directory on `start()`.
+        app (`str`, *optional*, defaults to `"server.app:app"`):
+            ASGI application path for uvicorn.
+        host (`str`, *optional*, defaults to `"0.0.0.0"`):
+            Host interface to bind to.
+        reload (`bool`, *optional*, defaults to `False`):
+            Whether to enable uvicorn's reload mode.
+        env_vars (`dict`, *optional*):
+            Environment variables to pass through to the spawned process.
+        context_timeout_s (`float`, *optional*, defaults to `60.0`):
+            How long to wait for the environment to become ready.
+
+    Examples:
+
+        ```python
+        provider = UVProvider(project_path="/path/to/env")
+        base_url = provider.start()
+        print(base_url)  # http://localhost:8000
+        # Use the environment via base_url
+        provider.stop()
+        ```
     """
 
     def __init__(
@@ -109,7 +163,14 @@ class UVProvider(RuntimeProvider):
         context_timeout_s: float = 60.0,
     ):
         """Initialize the UVProvider."""
-        self.project_path = os.path.abspath(project_path)
+        # `os.path.abspath` would mangle a `git+<url>` spec (e.g. collapsing
+        # `https://` to `https:/`) and is meaningless for one anyway -- only
+        # resolve it for genuine local paths. Git specs are cloned in `start()`.
+        self.project_path = (
+            project_path
+            if project_path.startswith(GIT_URL_PREFIX)
+            else os.path.abspath(project_path)
+        )
         self.app = app
         self.host = host
         self.reload = reload
@@ -118,6 +179,7 @@ class UVProvider(RuntimeProvider):
         _check_uv_installed()
         self._process = None
         self._base_url = None
+        self._clone_dir: Optional[str] = None
 
     def start(
         self,
@@ -130,28 +192,44 @@ class UVProvider(RuntimeProvider):
         Start the environment via `uv run`.
 
         Args:
-            port: The port to bind the environment to
-            env_vars: Environment variables to pass to the environment
-            workers: The number of workers to use
+            port (`int`, *optional*):
+                The port to bind the environment to.
+            env_vars (`dict`, *optional*):
+                Environment variables to pass to the environment.
+            workers (`int`, *optional*, defaults to `1`):
+                The number of workers to use.
 
         Returns:
-            The base URL of the environment
+            `str`: Base URL of the environment.
 
         Raises:
-            RuntimeError: If the environment is already running
+            RuntimeError: If the environment is already running.
         """
         if self._process is not None and self._process.poll() is None:
             raise RuntimeError("UVProvider is already running")
 
         bind_port = port or _find_free_port()
 
+        # A previous start() may have left a clone dir behind if the spawned
+        # process later died on its own -- the "already running" guard above
+        # only checks liveness, so a restart in that state would otherwise
+        # overwrite self._clone_dir and leak the old directory.
+        if self._clone_dir is not None:
+            shutil.rmtree(self._clone_dir, ignore_errors=True)
+            self._clone_dir = None
+
+        project_path = self.project_path
+        if project_path.startswith(GIT_URL_PREFIX):
+            self._clone_dir = _clone_git_project(project_path, self.context_timeout_s)
+            project_path = self._clone_dir
+
         command = _create_uv_command(
             host=self.host,
             port=bind_port,
             reload=self.reload,
             workers=workers,
             app=self.app,
-            project_path=self.project_path,
+            project_path=project_path,
         )
 
         env = os.environ.copy()
@@ -164,23 +242,31 @@ class UVProvider(RuntimeProvider):
         try:
             self._process = subprocess.Popen(command, env=env)
         except OSError as exc:
+            if self._clone_dir is not None:
+                shutil.rmtree(self._clone_dir, ignore_errors=True)
+                self._clone_dir = None
             raise RuntimeError(f"Failed to launch `uv run`: {exc}") from exc
 
         client_host = "127.0.0.1" if self.host in {"0.0.0.0", "::"} else self.host
         self._base_url = f"http://{client_host}:{bind_port}"
         return self._base_url
 
-    def wait_for_ready(self, timeout_s: float = 60.0) -> None:
+    def wait_for_ready(self, timeout_s: float | None = None) -> None:
         """
         Wait for the environment to become ready.
 
         Args:
-            timeout_s: The timeout to wait for the environment to become ready
+            timeout_s (`float`, *optional*):
+                Maximum time in seconds to wait for the environment to become
+                ready. Defaults to the provider's `context_timeout_s`.
 
         Raises:
-            RuntimeError: If the environment is not running
-            TimeoutError: If the environment does not become ready within the timeout
+            RuntimeError: If the environment is not running.
+            TimeoutError: If the environment does not become ready within the timeout.
         """
+        if timeout_s is None:
+            timeout_s = self.context_timeout_s
+
         if self._process and self._process.poll() is not None:
             code = self._process.returncode
             raise RuntimeError(f"uv process exited prematurely with code {code}")
@@ -190,23 +276,22 @@ class UVProvider(RuntimeProvider):
     def stop(self) -> None:
         """
         Stop the environment.
-
-        Raises:
-            RuntimeError: If the environment is not running
         """
-        if self._process is None:
-            return
-
-        if self._process.poll() is None:
-            self._process.terminate()
-            try:
-                self._process.wait(timeout=10.0)
-            except subprocess.TimeoutExpired:
-                self._process.kill()
-                self._process.wait(timeout=5.0)
-
-        self._process = None
-        self._base_url = None
+        if self._process is not None:
+            if self._process.poll() is None:
+                self._process.terminate()
+                try:
+                    self._process.wait(timeout=10.0)
+                except subprocess.TimeoutExpired:
+                    self._process.kill()
+                    self._process.wait(timeout=5.0)
+
+            self._process = None
+            self._base_url = None
+
+        if self._clone_dir is not None:
+            shutil.rmtree(self._clone_dir, ignore_errors=True)
+            self._clone_dir = None
 
     @property
     def base_url(self) -> str:
@@ -214,10 +299,10 @@ class UVProvider(RuntimeProvider):
         The base URL of the environment.
 
         Returns:
-            The base URL of the environment
+            `str`: Base URL of the running environment.
 
         Raises:
-            RuntimeError: If the environment is not running
+            RuntimeError: If the environment has not been started.
         """
         if self._base_url is None:
             raise RuntimeError("UVProvider has not been started")

src/core/containers/test_local_docker_provider.py

@@ -19,7 +19,7 @@ import requests
 from openenv.core.containers.runtime import LocalDockerProvider
 
 
-# TODO: Remove this test or make it a functional test sicne this will be tested in e2e test for echo env
+# TODO: Remove this test or make it a functional test since this will be tested in e2e test for echo env
 def test_local_docker_provider():
     """Test LocalDockerProvider end-to-end."""
     print("=" * 60)

src/core/env_client.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 Environment client for persistent sessions.
@@ -14,29 +10,43 @@ the overhead of HTTP request/response cycles.
 The client is async by default. For synchronous usage, use the `.sync()` method
 to get a `SyncEnvClient` wrapper.
 
-Example (async):
-    >>> async with GenericEnvClient(base_url="ws://localhost:8000") as env:
-    ...     result = await env.reset()
-    ...     result = await env.step({"code": "print('hello')"})
+Examples:
 
-Example (sync wrapper):
-    >>> env = GenericEnvClient(base_url="ws://localhost:8000").sync()
-    >>> with env:
-    ...     result = env.reset()
-    ...     result = env.step({"code": "print('hello')"})
+    Async usage:
+
+    ```python
+    async with GenericEnvClient(base_url="ws://localhost:8000") as env:
+        result = await env.reset()
+        result = await env.step({"code": "print('hello')"})
+    ```
+
+    Sync usage via `.sync()` wrapper:
+
+    ```python
+    env = GenericEnvClient(base_url="ws://localhost:8000").sync()
+    with env:
+        result = env.reset()
+        result = env.step({"code": "print('hello')"})
+    ```
 """
 
 from __future__ import annotations
 
 import asyncio
+import inspect
+import ipaddress
 import json
 import os
+import time
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Generic, Optional, Type, TYPE_CHECKING, TypeVar
+from collections.abc import Coroutine
+from contextlib import suppress
+from typing import Any, Callable, Dict, Generic, Optional, Type, TYPE_CHECKING, TypeVar
+from urllib.parse import urlsplit
 
 from .client_types import StateT, StepResult
 from .containers.runtime import LocalDockerProvider, UVProvider
-from .utils import convert_to_ws_url
+from .utils import convert_to_ws_url, hf_space_subdomain
 
 if TYPE_CHECKING:
     from websockets.asyncio.client import ClientConnection
@@ -45,10 +55,224 @@ if TYPE_CHECKING:
     from .sync_client import SyncEnvClient
 
 from websockets.asyncio.client import connect as ws_connect
+from websockets.exceptions import ConnectionClosed
+from websockets.protocol import State
 
 ActT = TypeVar("ActT")
 ObsT = TypeVar("ObsT")
 EnvClientT = TypeVar("EnvClientT", bound="EnvClient")
+ResultT = TypeVar("ResultT")
+
+_VALID_CLIENT_MODES = ("simulation", "production")
+
+
+class _AutoAsyncResult(Generic[ResultT]):
+    """Awaitable result that can also be resolved by synchronous access."""
+
+    def __init__(
+        self,
+        client: "EnvClient[Any, Any, Any]",
+        coro_factory: Callable[[], Any],
+    ):
+        self._client = client
+        self._coro_factory = coro_factory
+        self._resolved = False
+        self._result: ResultT | None = None
+
+    def __await__(self):
+        async def _await_result():
+            self._client._claim_execution_mode("async")
+            return await self._coro_factory()
+
+        return _await_result().__await__()
+
+    def _sync_result(self) -> ResultT:
+        if not self._resolved:
+            self._result = self._client._run_sync(self._coro_factory)
+            self._resolved = True
+        return self._result
+
+    def __getattr__(self, name: str) -> Any:
+        return getattr(self._sync_result(), name)
+
+    def __bool__(self) -> bool:
+        return bool(self._sync_result())
+
+    def __repr__(self) -> str:
+        if self._resolved:
+            return repr(self._result)
+        return f"<{type(self).__name__} pending>"
+
+
+class _BootstrapResult(Coroutine, Generic[EnvClientT]):
+    """Bootstrap handle returned by the client factory methods.
+
+    Returned by [`~openenv.core.EnvClient.from_docker_image`] and
+    [`~openenv.core.EnvClient.from_env`]. Resolving the handle is what actually
+    starts the container / Space and connects the WebSocket, so a single factory
+    call serves both execution modes:
+
+    - `await handle` connects on the running event loop and returns the connected
+      async client (unchanged async behavior).
+    - `handle.sync()` connects on the sync background loop and returns a
+      `SyncEnvClient`, mirroring the instance-level [`~openenv.core.EnvClient.sync`].
+
+    The handle is a full coroutine (it implements `send` / `throw` / `close`), so
+    it is a drop-in for the previous `async def` factories: `asyncio.run(...)`,
+    `run_async_safely(...)`, and bare `await` all accept it, in addition to the
+    new `.sync()` chain. Bootstrap is lazy — the underlying coroutine (and the
+    container/Space start) is created only when the handle is driven or `.sync()`
+    is called.
+    """
+
+    def __init__(self, bootstrap: Callable[[], EnvClientT]):
+        self._bootstrap = bootstrap
+        self._coro: Optional[Coroutine[Any, Any, EnvClientT]] = None
+        self._used = False
+
+    def _consume(self) -> EnvClientT:
+        """Run the bootstrap exactly once; a second resolve is a programming error.
+
+        Mirrors native coroutine semantics (a coroutine cannot be awaited twice)
+        so that a handle re-driven via a second `.sync()`, or `await` followed by
+        `.sync()`, raises instead of silently starting a second container/Space.
+        """
+        if self._used:
+            raise RuntimeError(
+                "This bootstrap handle has already been resolved; call the "
+                "factory again to start a new environment."
+            )
+        self._used = True
+        return self._bootstrap()
+
+    async def _resolve_async(self) -> EnvClientT:
+        client = self._consume()
+        try:
+            await client.connect()
+        except asyncio.CancelledError:
+            # _connect_async releases the provider on ordinary failures, but
+            # cancellation bypasses its `except Exception`. No client reaches
+            # the caller, so release the provider here before propagating the
+            # original cancellation. A repeated cancellation or a stop failure
+            # during cleanup must not replace it.
+            with suppress(Exception, asyncio.CancelledError):
+                await client.close()
+            raise
+        return client
+
+    def _ensure_coro(self) -> "Coroutine[Any, Any, EnvClientT]":
+        if self._coro is None:
+            self._coro = self._resolve_async()
+        return self._coro
+
+    def __await__(self):
+        return self._ensure_coro().__await__()
+
+    def send(self, value: Any) -> Any:
+        return self._ensure_coro().send(value)
+
+    def throw(self, *args: Any, **kwargs: Any) -> Any:
+        return self._ensure_coro().throw(*args, **kwargs)
+
+    def close(self) -> None:
+        if self._coro is not None:
+            self._coro.close()
+
+    def sync(self) -> "SyncEnvClient":
+        client = self._consume()
+        try:
+            client._run_sync(client._connect_async)
+        except Exception:
+            # _consume() already started the provider. On the async path
+            # _connect_async releases it, but here the failure may be the sync
+            # loop setup itself (before _connect_async runs), so stop the
+            # provider directly rather than routing through the broken loop.
+            client._stop_provider_best_effort()
+            raise
+        return client.sync()
+
+    def __repr__(self) -> str:
+        return f"<{type(self).__name__} pending (await, run, or call .sync())>"
+
+
+def _normalize_mode(mode: Optional[str]) -> str:
+    """Resolve and validate the client communication mode."""
+    raw_mode = (
+        os.environ.get("OPENENV_CLIENT_MODE", "simulation") if mode is None else mode
+    )
+    normalized_mode = raw_mode.lower()
+    if normalized_mode not in _VALID_CLIENT_MODES:
+        raise ValueError(
+            f"Invalid mode: '{normalized_mode}'. Must be 'simulation' or 'production'. "
+            f"Set via constructor parameter or OPENENV_CLIENT_MODE environment variable."
+        )
+    return normalized_mode
+
+
+def _is_localhost_ws_url(ws_url: str) -> bool:
+    """Return True when the WebSocket URL targets the local loopback interface.
+
+    The hostname is parsed from the URL so that only the actual host is matched.
+    Substring matching is avoided because remote hosts such as
+    ``my-localhost-proxy.example.com`` or ``127.0.0.1.example.com`` must not be
+    treated as local.
+    """
+    hostname = urlsplit(ws_url).hostname
+    if hostname is None:
+        return False
+    if hostname == "localhost":
+        return True
+    try:
+        return ipaddress.ip_address(hostname).is_loopback
+    except ValueError:
+        return False
+
+
+def _required_start_container_parameters(provider: Any) -> list[str]:
+    """Return required arguments for a bound provider.start_container()."""
+    try:
+        signature = inspect.signature(provider.start_container)
+    except (TypeError, ValueError):
+        return []
+    return [
+        name
+        for name, parameter in signature.parameters.items()
+        if parameter.default is inspect.Parameter.empty
+        and parameter.kind
+        in (
+            inspect.Parameter.POSITIONAL_ONLY,
+            inspect.Parameter.POSITIONAL_OR_KEYWORD,
+            inspect.Parameter.KEYWORD_ONLY,
+        )
+    ]
+
+
+async def _best_effort_close(ws: ClientConnection) -> None:
+    """Close a socket without letting a slow handshake or failure propagate.
+
+    Scheduled as a background task (never awaited directly) so a dropped
+    socket's close handshake can't hold up the caller that triggered the
+    drop -- see `EnvClient._receive()`. `EnvClient._close_async()` awaits
+    any still-running instance of this task before it returns, so a real
+    `close()` call does wait for the handshake; only the caller that
+    happened to trigger the drop is spared. `CancelledError` is caught
+    alongside `Exception` because that awaiting is exactly what would
+    otherwise cancel this mid-handshake if a caller called `close()`
+    concurrently.
+    """
+    try:
+        await ws.close()
+    except (Exception, asyncio.CancelledError):
+        pass  # Best effort
+
+
+async def _best_effort_graceful_close(ws: ClientConnection) -> None:
+    """Notify the server, then close a socket without propagating failures."""
+    try:
+        await ws.send(json.dumps({"type": "close"}))
+    except (Exception, asyncio.CancelledError):
+        pass  # Best effort
+    await _best_effort_close(ws)
 
 
 class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
@@ -68,75 +292,209 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
     - Better suited for long-running episodes
     - Async by default for modern Python async/await patterns
 
-    Example (async):
-        >>> from envs.coding_env.client import CodingEnv
-        >>>
-        >>> # Connect to a server using async context manager
-        >>> async with CodingEnv(base_url="ws://localhost:8000") as env:
-        ...     result = await env.reset(seed=42)
-        ...     while not result.done:
-        ...         action = agent.predict(result.observation)
-        ...         result = await env.step(action)
-
-    Example (sync wrapper):
-        >>> env = CodingEnv(base_url="ws://localhost:8000").sync()
-        >>> with env:
-        ...     result = env.reset(seed=42)
-        ...     result = env.step(action)
+    Examples:
+
+        Async usage:
+
+        ```python
+        from envs.coding_env.client import CodingEnv
+
+        async with CodingEnv(base_url="ws://localhost:8000") as env:
+            result = await env.reset(seed=42)
+            while not result.done:
+                action = agent.predict(result.observation)
+                result = await env.step(action)
+        ```
+
+        Sync usage via `.sync()` wrapper:
+
+        ```python
+        env = CodingEnv(base_url="ws://localhost:8000").sync()
+        with env:
+            result = env.reset(seed=42)
+            result = env.step(action)
+        ```
     """
 
     def __init__(
         self,
-        base_url: str,
+        base_url: Optional[str] = None,
         connect_timeout_s: float = 10.0,
         message_timeout_s: float = 60.0,
         max_message_size_mb: float = 100.0,
+        websocket_ping_interval_s: Optional[float] = 20.0,
+        websocket_ping_timeout_s: Optional[float] = 20.0,
         provider: Optional["ContainerProvider | RuntimeProvider"] = None,
         mode: Optional[str] = None,
+        headers: Optional[Dict[str, str]] = None,
     ):
         """
         Initialize environment client.
 
         Args:
-            base_url: Base URL of the environment server (http:// or ws://).
-                     Will be converted to ws:// if http:// is provided.
-            connect_timeout_s: Timeout for establishing WebSocket connection
-            message_timeout_s: Timeout for receiving responses to messages
-            max_message_size_mb: Maximum WebSocket message size in megabytes.
-                                Default 100MB to handle large observations (screenshots, DOM, etc.)
-            provider: Optional container/runtime provider for lifecycle management.
-                     Can be a ContainerProvider (Docker) or RuntimeProvider (UV).
-            mode: Communication mode: 'simulation' for Gym-style API (default) or
-                 'production' for MCP JSON-RPC protocol. Can also be set via the
-                 OPENENV_CLIENT_MODE environment variable. Constructor parameter
-                 takes precedence over environment variable. Case-insensitive.
+            base_url (`str`, *optional*):
+                Base URL of the environment server (http:// or ws://). Will be converted to
+                ws:// if http:// is provided. May be omitted when the provider
+                has enough constructor state to start itself.
+            connect_timeout_s (`float`, *optional*, defaults to `10.0`):
+                Timeout for establishing WebSocket connection.
+            message_timeout_s (`float`, *optional*, defaults to `60.0`):
+                Timeout for receiving responses to messages.
+            max_message_size_mb (`float`, *optional*, defaults to `100.0`):
+                Maximum WebSocket message size in megabytes. Default 100MB to handle large
+                observations (screenshots, DOM, etc.).
+            websocket_ping_interval_s (`float` or `None`, *optional*, defaults to `20.0`):
+                WebSocket keepalive ping interval. Pass `None` to disable.
+            websocket_ping_timeout_s (`float` or `None`, *optional*, defaults to `20.0`):
+                WebSocket keepalive pong timeout. Pass `None` to disable.
+            provider (`ContainerProvider` or `RuntimeProvider`, *optional*):
+                Container/runtime provider for lifecycle management.
+            mode (`str`, *optional*):
+                Communication mode: `'simulation'` for Gym-style API (default) or
+                `'production'` for MCP JSON-RPC protocol. Can also be set via the
+                `OPENENV_CLIENT_MODE` environment variable. Constructor parameter takes
+                precedence over environment variable. Case-insensitive.
+            headers (`dict[str, str]`, *optional*):
+                Extra HTTP headers sent when opening the connection, e.g.
+                `{"Authorization": f"Bearer {hf_token}"}` for a private Hugging Face Space.
         """
-        # Determine mode (constructor > env var > default)
-        if mode is None:
-            mode = os.environ.get("OPENENV_CLIENT_MODE", "simulation")
-
-        # Normalize and validate mode
-        mode = mode.lower()
-        if mode not in ("simulation", "production"):
-            raise ValueError(
-                f"Invalid mode: '{mode}'. Must be 'simulation' or 'production'. "
-                f"Set via constructor parameter or OPENENV_CLIENT_MODE environment variable."
-            )
+        if base_url is None and provider is None:
+            raise ValueError("EnvClient requires either base_url or provider.")
 
         # Store mode (use object.__setattr__ to bypass immutability)
-        object.__setattr__(self, "_mode", mode)
-
-        # Convert HTTP URL to WebSocket URL
-        ws_url = convert_to_ws_url(base_url)
+        object.__setattr__(self, "_mode", _normalize_mode(mode))
 
-        self._ws_url = f"{ws_url}/ws"
+        self._base_url: Optional[str] = None
+        self._ws_url: Optional[str] = None
         self._connect_timeout = connect_timeout_s
         self._message_timeout = message_timeout_s
         self._max_message_size = int(
             max_message_size_mb * 1024 * 1024
         )  # Convert MB to bytes
+        self._websocket_ping_interval_s = websocket_ping_interval_s
+        self._websocket_ping_timeout_s = websocket_ping_timeout_s
+        self._headers = headers
         self._provider = provider
+        self._provider_stopped = False
+        self._provider_cleanup_pending = False
+        self._start_provider_on_connect = base_url is None
+        self._child_clients: list[EnvClient[Any, Any, Any]] = []
         self._ws: Optional[ClientConnection] = None
+        self._execution_mode: Optional[str] = None
+        self._sync_client: Optional["SyncEnvClient[ActT, ObsT, StateT]"] = None
+        self._ws_loop: Optional[asyncio.AbstractEventLoop] = None
+        # Strong references for fire-and-forget socket closes (see _receive),
+        # so the task isn't garbage-collected mid-close. Discarded on done.
+        self._pending_close_tasks: set[asyncio.Task] = set()
+        if base_url is not None:
+            self._set_base_url(base_url)
+
+    @property
+    def base_url(self) -> Optional[str]:
+        """Public read-only URL of the environment server this client targets.
+
+        Returns the normalized `http(s)`/`ws(s)` base URL (no trailing slash),
+        or `None` when the client was created with a provider but has not yet
+        started its container (the URL is assigned lazily on `connect()`).
+
+        This is the public counterpart to the private `self._base_url`; the
+        previously-public `base_url` attribute was dropped in the `core`
+        refactor, leaving sync consumers with no way to read the URL back.
+        """
+        return self._base_url
+
+    def _set_base_url(self, base_url: str) -> None:
+        self._base_url = base_url.rstrip("/")
+        ws_url = convert_to_ws_url(base_url)
+        self._ws_url = f"{ws_url}/ws"
+
+    def _start_provider_if_needed(self) -> None:
+        # A missing URL does not mean the previous resource was released.
+        # Retry its cleanup before start can overwrite the provider's handle.
+        if self._provider_cleanup_pending:
+            if not self._start_provider_on_connect:
+                raise RuntimeError(
+                    "Provider cleanup is pending for this client with an existing "
+                    "base URL. Retry close() to finish cleanup, then create a new "
+                    "client instead of reconnecting to the old URL."
+                )
+            self._stop_provider()
+        if self._ws_url is not None:
+            return
+        if self._provider is None:
+            raise RuntimeError("EnvClient has no base URL or provider.")
+        if hasattr(self._provider, "start_container"):
+            required_parameters = _required_start_container_parameters(self._provider)
+            if required_parameters:
+                required = ", ".join(required_parameters)
+                raise ValueError(
+                    f"{type(self._provider).__name__} does not support "
+                    "provider-owned startup because start_container() requires "
+                    f"{required}. Start the provider manually and pass base_url, "
+                    "or configure a provider with a constructor-owned image/source."
+                )
+            self._provider_stopped = False
+            base_url = self._provider.start_container()
+            self._provider.wait_for_ready(base_url)
+        elif hasattr(self._provider, "start"):
+            self._provider_stopped = False
+            base_url = self._provider.start()
+            self._provider.wait_for_ready()
+        else:
+            raise TypeError("provider must define start_container() or start().")
+        self._set_base_url(base_url)
+
+    def _create_session_client(self) -> "EnvClient[Any, Any, Any]":
+        # Match _start_provider_if_needed's startup condition. A provider with
+        # an existing URL may already be serving the parent or other children.
+        starts_provider_here = self._provider is not None and self._ws_url is None
+        try:
+            self._start_provider_if_needed()
+            if self._base_url is None:
+                raise RuntimeError("EnvClient has no base URL.")
+
+            signature = inspect.signature(type(self))
+            accepts_kwargs = any(
+                parameter.kind == inspect.Parameter.VAR_KEYWORD
+                for parameter in signature.parameters.values()
+            )
+            candidate_kwargs = {
+                "base_url": self._base_url,
+                "connect_timeout_s": self._connect_timeout,
+                "message_timeout_s": self._message_timeout,
+                "max_message_size_mb": self._max_message_size / (1024 * 1024),
+                "websocket_ping_interval_s": self._websocket_ping_interval_s,
+                "websocket_ping_timeout_s": self._websocket_ping_timeout_s,
+                "mode": self._mode,
+                "headers": self._headers,
+            }
+            constructor_kwargs = {}
+            for name, value in candidate_kwargs.items():
+                if accepts_kwargs or name in signature.parameters:
+                    constructor_kwargs[name] = value
+
+            return type(self)(**constructor_kwargs)
+        except Exception:
+            if starts_provider_here and not self._provider_cleanup_pending:
+                self._stop_provider_best_effort()
+            raise
+
+    async def new_session(self) -> "EnvClient[Any, Any, Any]":
+        """
+        Create and connect a new session against the same environment server.
+
+        Returns:
+            `EnvClient`: A connected child client of the same concrete type.
+
+        The child session is tracked by this parent and closed when the parent
+        is closed. Server-side capacity still applies: when the server is at
+        `MAX_CONCURRENT_ENVS`, opening the child WebSocket can fail and is
+        surfaced as a connection error.
+        """
+        client = self._create_session_client()
+        await client.connect()
+        self._child_clients.append(client)
+        return client
 
     def __setattr__(self, name: str, value: Any) -> None:
         """Prevent modification of _mode after initialization."""
@@ -144,7 +502,56 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
             raise AttributeError("Cannot modify mode after initialization")
         super().__setattr__(name, value)
 
-    async def connect(self) -> "EnvClient":
+    def _claim_execution_mode(self, mode: str) -> None:
+        """Lock the client to sync or async execution on first use."""
+        if self._execution_mode is None:
+            self._execution_mode = mode
+        elif self._execution_mode != mode:
+            raise RuntimeError(
+                f"EnvClient is already being used in {self._execution_mode} mode. "
+                "Create a separate client instance when mixing sync and async code."
+            )
+
+    def _run_sync(
+        self,
+        coro_factory: Callable[[], Any],
+        *,
+        allow_async_handoff: bool = False,
+    ) -> Any:
+        """Run an async operation through the sync wrapper."""
+        if allow_async_handoff and self._execution_mode == "async":
+            self._execution_mode = "sync"
+        else:
+            self._claim_execution_mode("sync")
+        if self._sync_client is None:
+            self._sync_client = self.sync()
+        return self._sync_client._run(coro_factory())
+
+    def _dispatch(self, coro_factory: Callable[[], Any]) -> Any:
+        """Return an awaitable in async code and a concrete result in sync code."""
+        try:
+            running_loop = asyncio.get_running_loop()
+        except RuntimeError:
+            running_loop = None
+
+        if self._execution_mode == "sync":
+            sync_loop = (
+                self._sync_client._loop if self._sync_client is not None else None
+            )
+            if running_loop is sync_loop:
+                return coro_factory()
+            return self._run_sync(coro_factory)
+        if running_loop is not None:
+            if self._execution_mode == "async":
+                return coro_factory()
+            return _AutoAsyncResult(self, coro_factory)
+
+        return self._run_sync(coro_factory, allow_async_handoff=True)
+
+    def connect(self) -> Any:
+        return self._dispatch(self._connect_async)
+
+    async def _connect_async(self) -> "EnvClient":
         """
         Establish WebSocket connection to the server.
 
@@ -155,59 +562,132 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
             ConnectionError: If connection cannot be established
         """
         if self._ws is not None:
-            return self
-
-        # Bypass proxy for localhost connections
-        ws_url_lower = self._ws_url.lower()
-        is_localhost = "localhost" in ws_url_lower or "127.0.0.1" in ws_url_lower
-
-        old_no_proxy = os.environ.get("NO_PROXY")
-        if is_localhost:
-            # Set NO_PROXY to bypass proxy for localhost
-            current_no_proxy = old_no_proxy or ""
-            if "localhost" not in current_no_proxy.lower():
-                os.environ["NO_PROXY"] = (
-                    f"{current_no_proxy},localhost,127.0.0.1"
-                    if current_no_proxy
-                    else "localhost,127.0.0.1"
-                )
+            if self._ws.state in (State.CLOSING, State.CLOSED):
+                # Closed by the far end: a keepalive timeout, a tunnel dropping
+                # the socket, the server restarting. The object is still not
+                # None, so without this check it stays cached and every later
+                # call raises `ConnectionClosed` for the rest of the process's
+                # life. Only a demonstrably closed socket is dropped, so one
+                # still CONNECTING is left alone.
+                self._ws = None
+                self._ws_loop = None
+            elif self._ws_loop is asyncio.get_running_loop():
+                return self
+            # Connected from a different event loop than the one running
+            # now -- e.g. `client = await Client.from_env(...)` inside
+            # `asyncio.run(...)`, then `client.sync()` drives every later
+            # call on `SyncEnvClient`'s own dedicated background loop. The
+            # websocket object is bound to internals of the original loop,
+            # which is typically already closed by the time we get here, so
+            # it cannot be reused (or even cleanly closed) from this loop.
+            # Drop the stale reference and reconnect fresh below rather than
+            # silently no-op-ing onto a dead connection.
+            self._ws = None
+            self._ws_loop = None
+
+        # A timed-out request drops its socket immediately but closes it in the
+        # background so the timeout itself remains prompt. Wait for that close
+        # before opening a replacement: the old server-side session continues
+        # occupying a capacity slot until the close handshake finishes, and
+        # many environments allow only one session.
+        await self._drain_pending_close_tasks()
+
+        try:
+            self._start_provider_if_needed()
+        except Exception:
+            # A failed cleanup retry must propagate without attempting stop
+            # again. For a new startup failure, preserve its original error.
+            if not self._provider_cleanup_pending:
+                with suppress(Exception):
+                    await self.close()
+            raise
+
+        assert self._ws_url is not None
+
+        # Disable the proxy for localhost connections via the per-connection
+        # `proxy` argument rather than mutating the process-global NO_PROXY
+        # env var: concurrent connect() calls (e.g. asyncio.gather over many
+        # env clients) would otherwise race on os.environ and leak state.
+        connect_kwargs: Dict[str, Any] = {}
+        if _is_localhost_ws_url(self._ws_url):
+            connect_kwargs["proxy"] = None
 
         try:
             self._ws = await ws_connect(
                 self._ws_url,
                 open_timeout=self._connect_timeout,
                 max_size=self._max_message_size,
+                ping_interval=self._websocket_ping_interval_s,
+                ping_timeout=self._websocket_ping_timeout_s,
+                additional_headers=self._headers,
+                **connect_kwargs,
             )
+            self._ws_loop = asyncio.get_running_loop()
         except Exception as e:
+            await self.close()
             raise ConnectionError(f"Failed to connect to {self._ws_url}: {e}") from e
-        finally:
-            # Restore original NO_PROXY value
-            if is_localhost:
-                if old_no_proxy is None:
-                    os.environ.pop("NO_PROXY", None)
-                else:
-                    os.environ["NO_PROXY"] = old_no_proxy
 
         return self
 
-    async def disconnect(self) -> None:
+    def disconnect(self) -> Any:
+        return self._dispatch(self._disconnect_async)
+
+    def _schedule_socket_close(
+        self, ws: ClientConnection, *, notify_server: bool = False
+    ) -> asyncio.Task[None]:
+        """Schedule and track a socket close until its handshake finishes."""
+        close_coro = (
+            _best_effort_graceful_close(ws) if notify_server else _best_effort_close(ws)
+        )
+        close_task = asyncio.create_task(close_coro)
+        self._pending_close_tasks.add(close_task)
+        close_task.add_done_callback(self._pending_close_tasks.discard)
+        return close_task
+
+    async def _disconnect_async(self) -> None:
         """Close the WebSocket connection."""
         if self._ws is not None:
-            try:
-                # Send close message
-                await self._send({"type": "close"})
-            except Exception:
-                pass  # Best effort
-            try:
-                await self._ws.close()
-            except Exception:
-                pass
+            ws = self._ws
+            ws_loop = self._ws_loop
+            # Detach first so cancellation during the close handshake cannot
+            # leave a stale socket cached for a later operation.
             self._ws = None
+            self._ws_loop = None
+            same_loop = ws_loop is asyncio.get_running_loop()
+            if same_loop:
+                # Keep ownership of the handshake if this caller is cancelled.
+                # A later close/reconnect drains the task before tearing down
+                # the loop or opening a replacement server session.
+                close_task = self._schedule_socket_close(ws, notify_server=True)
+                await asyncio.shield(close_task)
+
+    async def _drain_pending_close_tasks(self) -> None:
+        """Wait for background socket closes owned by the current event loop.
+
+        Shielding keeps cancellation of the caller from cancelling the close
+        tasks themselves. This matters both before reconnecting, when the old
+        server session must release its capacity slot, and during explicit
+        client shutdown.
+        """
+        loop = asyncio.get_running_loop()
+        tasks = [
+            task
+            for task in tuple(self._pending_close_tasks)
+            if not task.done() and task.get_loop() is loop
+        ]
+        if tasks:
+            await asyncio.shield(asyncio.gather(*tasks, return_exceptions=True))
 
     async def _ensure_connected(self) -> None:
-        """Ensure WebSocket connection is established."""
-        if self._ws is None:
-            await self.connect()
+        """Ensure WebSocket connection is established on the current loop.
+
+        Always delegates to `_connect_async()` rather than pre-checking
+        `self._ws is None`: `_connect_async()` itself is the one that knows
+        whether an existing `_ws` is reusable (same event loop) or stale (a
+        different one, e.g. from a prior `from_env()` call now being driven
+        through `.sync()`'s own loop).
+        """
+        await self._connect_async()
 
     async def _send(self, message: Dict[str, Any]) -> None:
         """Send a message over the WebSocket."""
@@ -216,14 +696,49 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
         await self._ws.send(json.dumps(message))
 
     async def _receive(self) -> Dict[str, Any]:
-        """Receive and parse a message from the WebSocket."""
+        """Receive the response on the same WebSocket that carried the request.
+
+        Reconnecting here would wait on a fresh socket where the request was
+        never sent. The next complete operation may reconnect in `_send()`.
+        """
         assert self._ws is not None
-        raw = await asyncio.wait_for(self._ws.recv(), timeout=self._message_timeout)
+        try:
+            raw = await asyncio.wait_for(self._ws.recv(), timeout=self._message_timeout)
+        except (asyncio.TimeoutError, asyncio.CancelledError):
+            # The server may still write the response for this request to the
+            # socket after we give up waiting on it. If we left the socket
+            # open, the next call's `_receive()` would read that stale frame
+            # and silently pair it with an unrelated request -- and every
+            # response after that would be shifted by one for the life of
+            # the connection. Drop the socket so the next `_send()` opens a
+            # fresh one instead of reusing a desynced one.
+            #
+            # An outer `asyncio.wait_for` around the whole call (e.g. a
+            # caller-imposed deadline on step()) cancels this await with
+            # CancelledError, not our own message_timeout's TimeoutError --
+            # same desync, different exception, so both are caught here.
+            ws = self._ws
+            self._ws = None
+            self._ws_loop = None
+            # Fire-and-forget: `close()` waits out the library's close
+            # handshake (`close_timeout`, 10s by default) if the server is
+            # slow to ack, and *awaiting* that here would hold up this
+            # exception -- so an outer `asyncio.wait_for(..., timeout=50ms)`
+            # would actually block for up to 10s before its deadline was
+            # honored. Scheduling it lets the exception propagate
+            # immediately while the close still happens in the background.
+            self._schedule_socket_close(ws)
+            raise
         return json.loads(raw)
 
     async def _send_and_receive(self, message: Dict[str, Any]) -> Dict[str, Any]:
         """Send a message and wait for response."""
-        await self._send(message)
+        try:
+            await self._send(message)
+        except ConnectionClosed:
+            # A refused session (e.g. at capacity) sends an error frame then closes. If
+            # the close beats our send, read the queued frame (recv() re-raises if none).
+            pass
         response = await self._receive()
 
         # Check for error response
@@ -237,40 +752,78 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
         return response
 
     @classmethod
-    async def from_docker_image(
+    def _bootstrap_container(
         cls: Type[EnvClientT],
         image: str,
         provider: Optional["ContainerProvider"] = None,
         **kwargs: Any,
     ) -> EnvClientT:
+        """Start a Docker container and build an *unconnected* client for it.
+
+        Invoked lazily by the `from_docker_image` bootstrap handle when it is
+        awaited or `.sync()`'d: the container start / readiness wait is plain
+        blocking code, so the only thing that differs between the async and sync
+        resolution paths is how the WebSocket is connected (awaited vs. run sync).
+        """
+        if provider is None:
+            provider = LocalDockerProvider()
+
+        try:
+            # Start container, wait for readiness, then build the client.
+            base_url = provider.start_container(image, **kwargs)
+            provider.wait_for_ready(base_url)
+            return cls(base_url=base_url, provider=provider)
+        except Exception:
+            # No EnvClient exists yet for the caller to close(), so release the
+            # container here if start / readiness / construction fails.
+            provider.stop_container()
+            raise
+
+    @classmethod
+    def from_docker_image(
+        cls: Type[EnvClientT],
+        image: str,
+        provider: Optional["ContainerProvider"] = None,
+        **kwargs: Any,
+    ) -> "_BootstrapResult[EnvClientT]":
         """
         Create an environment client by spinning up a Docker container.
 
+        Returns a bootstrap handle so the same call works from both async and
+        synchronous code: `await` it for the connected async client, or chain
+        `.sync()` for a connected `SyncEnvClient` (e.g. a TRL GRPO rollout loop
+        that cannot `await`). Bootstrap is lazy — the container starts when the
+        handle is resolved.
+
         Args:
-            image: Docker image name to run (e.g., "coding-env:latest")
-            provider: Container provider to use (defaults to LocalDockerProvider)
-            **kwargs: Additional arguments to pass to provider.start_container()
+            image (`str`):
+                Docker image name to run (e.g., `"coding-env:latest"`).
+            provider (`ContainerProvider`, *optional*):
+                Container provider to use. Defaults to `LocalDockerProvider`.
+            **kwargs:
+                Additional arguments to pass to `provider.start_container()`.
 
         Returns:
-            Connected client instance
-        """
-        if provider is None:
-            provider = LocalDockerProvider()
+            `_BootstrapResult`: `await` for a connected async client, or call
+            `.sync()` for a connected `SyncEnvClient`.
 
-        # Start container
-        base_url = provider.start_container(image, **kwargs)
+        Examples:
 
-        # Wait for server to be ready
-        provider.wait_for_ready(base_url)
+        ```python
+        # Async
+        env = await MyEnv.from_docker_image("coding-env:latest")
 
-        # Create and connect client
-        client = cls(base_url=base_url, provider=provider)
-        await client.connect()
-
-        return client
+        # Sync
+        env = MyEnv.from_docker_image("coding-env:latest").sync()
+        result = env.reset()
+        ```
+        """
+        return _BootstrapResult(
+            lambda: cls._bootstrap_container(image, provider, **kwargs)
+        )
 
     @classmethod
-    async def from_env(
+    def _bootstrap_env(
         cls: Type[EnvClientT],
         repo_id: str,
         *,
@@ -278,40 +831,13 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
         provider: Optional["ContainerProvider | RuntimeProvider"] = None,
         **provider_kwargs: Any,
     ) -> EnvClientT:
-        """
-        Create a client from a Hugging Face Space.
+        """Start a HF Space (docker or uv) and build an *unconnected* client.
 
-        Args:
-            repo_id: Hugging Face space identifier ``{org}/{space}``.
-            use_docker: When ``True`` (default) pull from the HF registry and
-                launch via :class:`LocalDockerProvider`. When ``False`` run the
-                space locally with :class:`UVProvider`.
-            provider: Optional provider instance to reuse. Must be a
-                :class:`ContainerProvider` when ``use_docker=True`` and a
-                :class:`RuntimeProvider` otherwise.
-            provider_kwargs: Additional keyword arguments forwarded to
-                either the container provider's ``start_container`` (docker)
-                or to the ``UVProvider`` constructor/start (uv). When
-                ``use_docker=False``, the ``project_path`` argument can be
-                used to override the default git URL
-                (``git+https://huggingface.co/spaces/{repo_id}``).
-
-        Returns:
-            Connected client instance
-
-        Examples:
-            >>> # Pull and run from HF Docker registry
-            >>> env = await MyEnv.from_env("openenv/echo-env")
-            >>>
-            >>> # Run locally with UV (clones the space)
-            >>> env = await MyEnv.from_env("openenv/echo-env", use_docker=False)
-            >>>
-            >>> # Run from a local checkout
-            >>> env = await MyEnv.from_env(
-            ...     "openenv/echo-env",
-            ...     use_docker=False,
-            ...     project_path="/path/to/local/checkout"
-            ... )
+        Invoked lazily by the `from_env` bootstrap handle; see
+        `_bootstrap_container` for the same split rationale. On a startup
+        failure the spawned process (and, for a git+ `project_path`, the temp
+        clone directory) is released before re-raising; a later *connection*
+        failure is cleaned up by `_connect_async` via `close()`.
         """
         # Extract start args that apply to both providers
         start_args = {}
@@ -323,15 +849,18 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
             # Docker mode: pull from HF registry
             docker_provider = provider or LocalDockerProvider()
             tag = provider_kwargs.pop("tag", "latest")
-            image = f"registry.hf.space/{repo_id.replace('/', '-')}:{tag}"
-            base_url = docker_provider.start_container(
-                image, **start_args, **provider_kwargs
-            )
-            docker_provider.wait_for_ready(base_url)
-
-            client = cls(base_url=base_url, provider=docker_provider)
-            await client.connect()
-            return client
+            image = f"registry.hf.space/{hf_space_subdomain(repo_id)}:{tag}"
+            try:
+                base_url = docker_provider.start_container(
+                    image, **start_args, **provider_kwargs
+                )
+                docker_provider.wait_for_ready(base_url)
+                return cls(base_url=base_url, provider=docker_provider)
+            except Exception:
+                # No EnvClient exists yet for the caller to close(), so release
+                # the container here if start / readiness / construction fails.
+                docker_provider.stop_container()
+                raise
         else:
             # UV mode: clone and run with uv
             if provider is None:
@@ -347,12 +876,92 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
                         "provider_kwargs cannot be used when supplying a provider instance"
                     )
 
-            base_url = provider.start(**start_args)
-            provider.wait_for_ready()
+            try:
+                context_timeout_s = getattr(provider, "context_timeout_s", None)
+                deadline = (
+                    time.monotonic() + context_timeout_s
+                    if context_timeout_s is not None
+                    else None
+                )
+                base_url = provider.start(**start_args)
+                if deadline is None:
+                    provider.wait_for_ready()
+                else:
+                    provider.wait_for_ready(
+                        timeout_s=max(0.0, deadline - time.monotonic())
+                    )
+                return cls(base_url=base_url, provider=provider)
+            except Exception:
+                # No EnvClient exists yet for the caller to close(), so this is
+                # the only chance to release the spawned process and (for a
+                # git+ project_path) the temp clone directory. Covers start,
+                # readiness, and client construction (e.g. an invalid mode).
+                provider.stop()
+                raise
 
-            client = cls(base_url=base_url, provider=provider)
-            await client.connect()
-            return client
+    @classmethod
+    def from_env(
+        cls: Type[EnvClientT],
+        repo_id: str,
+        *,
+        use_docker: bool = True,
+        provider: Optional["ContainerProvider | RuntimeProvider"] = None,
+        **provider_kwargs: Any,
+    ) -> "_BootstrapResult[EnvClientT]":
+        """
+        Create a client from a Hugging Face Space.
+
+        Returns a bootstrap handle: `await` it for the connected async client, or
+        chain `.sync()` for a connected `SyncEnvClient`. Bootstrap is lazy — the
+        Space starts when the handle is resolved.
+
+        Args:
+            repo_id (`str`):
+                Hugging Face space identifier `{org}/{space}`.
+            use_docker (`bool`, *optional*, defaults to `True`):
+                When `True`, pull from the HF registry and launch via `LocalDockerProvider`.
+                When `False`, run the space locally with `UVProvider`.
+            provider (`ContainerProvider` or `RuntimeProvider`, *optional*):
+                Provider instance to reuse. Must be a `ContainerProvider` when
+                `use_docker=True` and a `RuntimeProvider` otherwise.
+            **provider_kwargs:
+                Additional keyword arguments forwarded to either the container provider's
+                `start_container` (docker) or to the `UVProvider` constructor/start (uv).
+                When `use_docker=False`, the `project_path` argument can be used to override
+                the default git URL (`git+https://huggingface.co/spaces/{repo_id}`).
+
+        Returns:
+            `_BootstrapResult`: `await` for a connected async client, or call
+            `.sync()` for a connected `SyncEnvClient`.
+
+        Examples:
+
+            ```python
+            # Async: pull and run from HF Docker registry
+            env = await MyEnv.from_env("openenv/echo-env")
+
+            # Sync: chain .sync()
+            env = MyEnv.from_env("openenv/echo-env").sync()
+
+            # Run locally with UV (clones the space)
+            env = await MyEnv.from_env("openenv/echo-env", use_docker=False)
+
+            # Run from a local checkout
+            env = await MyEnv.from_env(
+                "openenv/echo-env",
+                use_docker=False,
+                project_path="/path/to/local/checkout"
+            )
+            ```
+        """
+        return _BootstrapResult(
+            lambda: cls._bootstrap_env(
+                repo_id,
+                use_docker=use_docker,
+                provider=provider,
+                **provider_kwargs,
+            )
+        )
 
     @abstractmethod
     def _step_payload(self, action: ActT) -> Dict[str, Any]:
@@ -369,15 +978,16 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
         """Convert a JSON response from the state endpoint to a State object."""
         raise NotImplementedError
 
-    async def reset(self, **kwargs: Any) -> StepResult[ObsT]:
+    def reset(self, **kwargs: Any) -> Any:
+        return self._dispatch(lambda: self._reset_async(**kwargs))
+
+    async def _reset_async(self, **kwargs: Any) -> StepResult[ObsT]:
         """
         Reset the environment with optional parameters.
 
         Args:
-            **kwargs: Optional parameters passed to the environment's reset method.
-                     Common parameters include:
-                     - seed: Random seed for reproducibility
-                     - episode_id: Custom episode identifier
+            **kwargs:
+                Optional parameters passed to the environment's reset method.
 
         Returns:
             StepResult containing initial observation
@@ -389,13 +999,18 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
         response = await self._send_and_receive(message)
         return self._parse_result(response.get("data", {}))
 
-    async def step(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:
+    def step(self, action: ActT, **kwargs: Any) -> Any:
+        return self._dispatch(lambda: self._step_async(action, **kwargs))
+
+    async def _step_async(self, action: ActT, **kwargs: Any) -> StepResult[ObsT]:
         """
         Execute an action in the environment.
 
         Args:
-            action: The action to execute
-            **kwargs: Optional parameters (currently ignored)
+            action:
+                The action to execute.
+            **kwargs:
+                Optional parameters (currently ignored).
 
         Returns:
             StepResult containing observation, reward, and done status
@@ -407,7 +1022,10 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
         response = await self._send_and_receive(message)
         return self._parse_result(response.get("data", {}))
 
-    async def state(self) -> StateT:
+    def state(self) -> Any:
+        return self._dispatch(self._state_async)
+
+    async def _state_async(self) -> StateT:
         """
         Get the current environment state from the server.
 
@@ -418,24 +1036,120 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
         response = await self._send_and_receive(message)
         return self._parse_state(response.get("data", {}))
 
-    async def close(self) -> None:
+    def close(self) -> Any:
+        return self._dispatch(self._close_async)
+
+    async def _close_child_clients(self) -> asyncio.CancelledError | None:
+        """Close every captured child while deferring caller cancellation."""
+        children = list(self._child_clients)
+        self._child_clients.clear()
+        close_tasks = []
+        deferred_cancellation: asyncio.CancelledError | None = None
+
+        for child in children:
+            try:
+                close_tasks.append(asyncio.ensure_future(child.close()))
+            except asyncio.CancelledError as exc:
+                if deferred_cancellation is None:
+                    deferred_cancellation = exc
+            except Exception:
+                pass  # Best effort
+
+        if close_tasks:
+            close_group = asyncio.gather(*close_tasks, return_exceptions=True)
+            while not close_group.done():
+                try:
+                    await asyncio.shield(close_group)
+                except asyncio.CancelledError as exc:
+                    if deferred_cancellation is None:
+                        deferred_cancellation = exc
+
+            for result in close_group.result():
+                if (
+                    isinstance(result, asyncio.CancelledError)
+                    and deferred_cancellation is None
+                ):
+                    deferred_cancellation = result
+
+        return deferred_cancellation
+
+    async def _close_async(self) -> None:
         """
         Close the WebSocket connection and clean up resources.
 
         If this client was created via from_docker_image() or from_env(),
         this will also stop and remove the associated container/process.
         """
-        await self.disconnect()
+        deferred_cancellation: asyncio.CancelledError | None = None
+        try:
+            try:
+                try:
+                    child_cancellation = await self._close_child_clients()
+                except asyncio.CancelledError as exc:
+                    child_cancellation = exc
+                if child_cancellation is not None:
+                    deferred_cancellation = child_cancellation
+
+                # A real close waits out backgrounded closes, while shielding
+                # their socket handshakes from cancellation.
+                try:
+                    await self._drain_pending_close_tasks()
+                except asyncio.CancelledError as exc:
+                    if deferred_cancellation is None:
+                        deferred_cancellation = exc
+            finally:
+                # Run even when child or pending-close cleanup is cancelled. A
+                # client may already have reconnected, and that current socket
+                # must not remain cached or open during teardown.
+                try:
+                    await self._disconnect_async()
+                except asyncio.CancelledError as exc:
+                    if deferred_cancellation is None:
+                        deferred_cancellation = exc
+        finally:
+            self._stop_provider()
+
+        if deferred_cancellation is not None:
+            raise deferred_cancellation
 
-        if self._provider is not None:
-            # Handle both ContainerProvider and RuntimeProvider
-            if hasattr(self._provider, "stop_container"):
-                self._provider.stop_container()
-            elif hasattr(self._provider, "stop"):
-                self._provider.stop()
+    def _stop_provider(self) -> None:
+        """Stop the provider once, retaining it for retries and later startup."""
+        try:
+            provider = self._provider
+            if provider is None or self._provider_stopped:
+                return
+
+            if hasattr(provider, "stop_container"):
+                stop = provider.stop_container
+            elif hasattr(provider, "stop"):
+                stop = provider.stop
+            else:
+                return
+
+            self._provider_cleanup_pending = True
+            stop()
+            # Only a successful stop discharges our cleanup responsibility.
+            self._provider_cleanup_pending = False
+            self._provider_stopped = True
+        finally:
+            if self._start_provider_on_connect:
+                self._base_url = None
+                self._ws_url = None
+
+    def _stop_provider_best_effort(self) -> None:
+        """Stop the underlying provider directly, ignoring any errors.
+
+        Releases a started container/process when there is no connected client
+        to `close()` through the normal path — e.g. sync bootstrap setup fails
+        after the provider started but before the connection is established, so
+        routing cleanup through the (possibly broken) sync loop is not an option.
+        """
+        with suppress(Exception):
+            self._stop_provider()
 
     async def __aenter__(self) -> "EnvClient":
         """Enter async context manager, ensuring connection is established."""
+        self._claim_execution_mode("async")
         await self.connect()
         return self
 
@@ -444,17 +1158,13 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
         await self.close()
 
     def __enter__(self) -> "EnvClient":
-        """Sync context manager entry - raises error suggesting async usage."""
-        raise TypeError(
-            "EnvClient is async by default. Use 'async with' instead of 'with', "
-            "or call .sync() to get a synchronous wrapper:\n"
-            "  async with client:  # async usage\n"
-            "  with client.sync():  # sync wrapper"
-        )
+        """Enter sync context manager, ensuring connection is established."""
+        self._run_sync(self._connect_async)
+        return self
 
     def __exit__(self, exc_type, exc_val, exc_tb) -> None:
-        """Sync context manager exit - should not be reached."""
-        pass  # pragma: no cover
+        """Exit sync context manager, closing connection."""
+        self._run_sync(self._close_async)
 
     def sync(self) -> "SyncEnvClient":
         """
@@ -469,16 +1179,19 @@ class EnvClient(ABC, Generic[ActT, ObsT, StateT]):
         Returns:
             SyncEnvClient wrapper that provides synchronous methods
 
-        Example:
-            >>> # Create async client and get sync wrapper
-            >>> async_client = GenericEnvClient(base_url="http://localhost:8000")
-            >>> sync_client = async_client.sync()
-            >>>
-            >>> # Use synchronous API
-            >>> with sync_client:
-            ...     result = sync_client.reset()
-            ...     result = sync_client.step({"code": "print('hello')"})
+        Examples:
+
+            ```python
+            async_client = GenericEnvClient(base_url="http://localhost:8000")
+            sync_client = async_client.sync()
+
+            with sync_client:
+                result = sync_client.reset()
+                result = sync_client.step({"code": "print('hello')"})
+            ```
         """
         from .sync_client import SyncEnvClient
 
-        return SyncEnvClient(self)
+        if self._sync_client is None:
+            self._sync_client = SyncEnvClient(self)
+        return self._sync_client

src/core/env_server/__init__.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """Core environment interfaces and types."""
 
@@ -16,7 +12,7 @@ from .exceptions import (
     SessionNotFoundError,
 )
 from .http_server import create_app, create_fastapi_app, HTTPEnvServer
-from .interfaces import Environment, Message, ModelTokenizer, Transform
+from .interfaces import Environment, Message, ModelTokenizer, TaskProvider, Transform
 
 try:
     from .mcp_environment import MCPEnvironment
@@ -51,8 +47,12 @@ from .types import (
     Action,
     BaseMessage,
     ConcurrencyConfig,
+    GetTaskRangeRequest,
+    GetTaskRequest,
     HealthResponse,
     HealthStatus,
+    ListTasksRequest,
+    NumTasksRequest,
     Observation,
     SchemaResponse,
     ServerCapacityStatus,
@@ -80,6 +80,7 @@ __all__ = [
     # Core interfaces
     "Environment",
     "Transform",
+    "TaskProvider",
     "Message",
     "ModelTokenizer",
     # Types
@@ -88,6 +89,10 @@ __all__ = [
     "State",
     "SchemaResponse",
     "HealthResponse",
+    "ListTasksRequest",
+    "NumTasksRequest",
+    "GetTaskRequest",
+    "GetTaskRangeRequest",
     # Enums
     "HealthStatus",
     "ServerMode",

src/core/env_server/_utils.py

@@ -0,0 +1,10 @@
+"""Shared env server implementation helpers."""
+
+from __future__ import annotations
+
+from typing import Any
+
+
+def overrides_method(method: Any, base_method: Any) -> bool:
+    """Return whether a bound method differs from the base implementation."""
+    return getattr(method, "__func__", method) is not base_method

src/core/env_server/base_transforms.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """Base transform implementations for composing environment-specific transforms."""
 

src/core/env_server/exceptions.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """Custom exceptions for environment server operations."""
 

src/core/env_server/gradio_theme.py

@@ -1,42 +1,11 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
-"""Unified terminal-style theme for OpenEnv Gradio UI (light/dark)."""
+"""Theme for the OpenEnv Gradio UI: white, hairline borders, one green accent."""
 
 from __future__ import annotations
 
 import gradio as gr
 
-_MONO_FONTS = (
-    "JetBrains Mono",
-    "Fira Code",
-    "Cascadia Code",
-    "Consolas",
-    "ui-monospace",
-    "monospace",
-)
-
-_CORE_FONT = (
-    "Lato",
-    "Inter",
-    "Arial",
-    "Helvetica",
-    "sans-serif",
-)
-
-_ZERO_RADIUS = gr.themes.Size(
-    xxs="0px",
-    xs="0px",
-    sm="0px",
-    md="0px",
-    lg="0px",
-    xl="0px",
-    xxl="0px",
-)
-
 _GREEN_HUE = gr.themes.Color(
     c50="#e6f4ea",
     c100="#ceead6",
@@ -52,77 +21,122 @@ _GREEN_HUE = gr.themes.Color(
 )
 
 _NEUTRAL_HUE = gr.themes.Color(
-    c50="#f6f8fa",
-    c100="#eaeef2",
-    c200="#d0d7de",
-    c300="#afb8c1",
-    c400="#8c959f",
-    c500="#6e7781",
-    c600="#57606a",
-    c700="#424a53",
-    c800="#32383f",
-    c900="#24292f",
-    c950="#1b1f24",
+    c50="#fafafa",
+    c100="#f4f4f5",
+    c200="#e4e4e7",
+    c300="#d4d4d8",
+    c400="#a1a1aa",
+    c500="#71717a",
+    c600="#52525b",
+    c700="#3f3f46",
+    c800="#27272a",
+    c900="#18181b",
+    c950="#0a0a0a",
 )
 
 OPENENV_GRADIO_THEME = gr.themes.Base(
     primary_hue=_GREEN_HUE,
     secondary_hue=_NEUTRAL_HUE,
     neutral_hue=_NEUTRAL_HUE,
-    font=_CORE_FONT,
-    font_mono=_MONO_FONTS,
-    radius_size=_ZERO_RADIUS,
+    font=(gr.themes.GoogleFont("Geist"), "system-ui", "sans-serif"),
+    font_mono=(gr.themes.GoogleFont("Geist Mono"), "ui-monospace", "monospace"),
+    radius_size=gr.themes.sizes.radius_md,
 ).set(
     body_background_fill="#ffffff",
+    body_text_color="#0a0a0a",
     background_fill_primary="#ffffff",
-    background_fill_secondary="#f6f8fa",
+    background_fill_secondary="#fafafa",
     block_background_fill="#ffffff",
-    block_border_color="#ffffff",
-    block_label_text_color="#57606a",
-    block_title_text_color="#24292f",
-    border_color_primary="#d0d7de",
+    block_border_width="0px",
+    block_label_text_color="#52525b",
+    block_label_background_fill="transparent",
+    border_color_primary="#ecedef",
     input_background_fill="#ffffff",
-    input_border_color="#d0d7de",
-    button_primary_background_fill="#1a7f37",
-    button_primary_background_fill_hover="#116329",
+    input_border_color="#e4e4e7",
+    input_border_color_focus="#1a7f37",
+    button_primary_background_fill="#0a0a0a",
+    button_primary_background_fill_hover="#27272a",
     button_primary_text_color="#ffffff",
-    button_secondary_background_fill="#f6f8fa",
-    button_secondary_background_fill_hover="#eaeef2",
-    button_secondary_text_color="#24292f",
-    button_secondary_border_color="#d0d7de",
-    body_background_fill_dark="#0d1117",
-    background_fill_primary_dark="#0d1117",
-    background_fill_secondary_dark="#0d1117",
-    block_background_fill_dark="#0d1117",
-    block_border_color_dark="#0d1117",
-    block_label_text_color_dark="#8b949e",
-    block_title_text_color_dark="#c9d1d9",
-    border_color_primary_dark="#30363d",
-    input_background_fill_dark="#0d1117",
-    input_border_color_dark="#30363d",
-    button_primary_background_fill_dark="#30363d",
-    button_primary_background_fill_hover_dark="#484f58",
-    button_primary_text_color_dark="#c9d1d9",
-    button_secondary_background_fill_dark="#21262d",
-    button_secondary_background_fill_hover_dark="#30363d",
-    button_secondary_text_color_dark="#c9d1d9",
-    button_secondary_border_color_dark="#30363d",
+    button_primary_border_color="#0a0a0a",
+    button_secondary_background_fill="#ffffff",
+    button_secondary_background_fill_hover="#fafafa",
+    button_secondary_text_color="#0a0a0a",
+    button_secondary_border_color="#e4e4e7",
+    button_border_width="1px",
+    checkbox_background_color_selected="#1a7f37",
+    checkbox_border_color_selected="#1a7f37",
+    color_accent="#1a7f37",
+    color_accent_soft="#f0f7f2",
+    color_accent_soft_dark="#0d2416",
+    body_text_color_subdued="#52525b",
+    body_text_color_subdued_dark="#a1a1aa",
+    body_background_fill_dark="#0a0a0a",
+    body_text_color_dark="#fafafa",
+    background_fill_primary_dark="#0a0a0a",
+    background_fill_secondary_dark="#141414",
+    block_background_fill_dark="#0a0a0a",
+    block_label_text_color_dark="#a1a1aa",
+    border_color_primary_dark="#27272a",
+    input_background_fill_dark="#0a0a0a",
+    input_border_color_dark="#3f3f46",
+    button_primary_background_fill_dark="#fafafa",
+    button_primary_background_fill_hover_dark="#e4e4e7",
+    button_primary_text_color_dark="#0a0a0a",
+    button_primary_border_color_dark="#fafafa",
+    button_secondary_background_fill_dark="#0a0a0a",
+    button_secondary_background_fill_hover_dark="#18181b",
+    button_secondary_text_color_dark="#fafafa",
+    button_secondary_border_color_dark="#3f3f46",
 )
 
 OPENENV_GRADIO_CSS = """
-* { border-radius: 0 !important; }
-.col-left { padding: 16px !important; }
-.col-right { padding: 16px !important; }
-.prose, .markdown-text, .md,
-.prose > *, .markdown-text > * {
-    background: transparent !important;
-    border: none !important;
-    box-shadow: none !important;
-}
-.dark .col-left {
-    border-left-color: rgba(139, 148, 158, 0.4) !important;
-}
-.dark .col-right {
-    border-left-color: rgba(201, 209, 217, 0.3) !important;
-}
+.gradio-container { max-width: 1120px !important; margin: 0 auto !important; }
+.oe-header { padding: 24px 0 12px; }
+.oe-header h1 { margin: 6px 0 4px; font-size: 40px; line-height: 1.1; font-weight: 700; letter-spacing: -.03em; }
+.oe-header p, .oe-step p, .oe-episode small, .oe-muted { color: var(--body-text-color-subdued); }
+.oe-header p { margin: 0; font-size: 16px; max-width: 640px; }
+.oe-eyebrow { font-family: var(--font-mono); font-size: 13px; font-weight: 500; color: var(--color-accent); }
+.oe-card { border: 1px solid var(--border-color-primary) !important; border-radius: 12px !important; padding: 18px 20px !important; background: var(--background-fill-primary) !important; gap: 12px !important; }
+.oe-row { align-items: center !important; }
+.oe-step h2, .oe-episode h2 { margin: 0; font-size: 17px; font-weight: 600; }
+.oe-step p { margin: 2px 0 0; font-size: 14px; }
+.oe-step code, .oe-episode code, .oe-result code { font-family: var(--font-mono); background: none; padding: 0; }
+.oe-run { flex: none !important; align-self: flex-start !important; width: auto !important; min-width: 140px !important; padding: 0 22px !important; margin-left: var(--block-padding, 10px) !important; }
+.oe-card textarea, .oe-card input[type=number], .oe-card input[type=password], .oe-card input[type=text] { border: 1px solid var(--input-border-color) !important; border-radius: 8px !important; padding: 10px 12px !important; font-family: var(--font-mono) !important; }
+.oe-tools .wrap, .oe-actions .wrap { gap: 8px; }
+.oe-tools .wrap { display: grid !important; grid-template-columns: repeat(auto-fit, minmax(220px, 1fr)); }
+.oe-actions .wrap { display: flex !important; flex-wrap: wrap; }
+.oe-tools .wrap label, .oe-actions .wrap label { border: 1px solid var(--input-border-color) !important; border-radius: 10px !important; background: var(--background-fill-primary) !important; font-family: var(--font-mono); cursor: pointer; }
+.oe-tools .wrap label { padding: 12px 14px !important; font-size: 13px; }
+.oe-actions .wrap label { padding: 10px 16px !important; font-size: 14px; }
+.oe-tools .wrap label.selected { border: 1.5px solid var(--color-accent) !important; background: var(--color-accent-soft) !important; }
+.oe-actions input[type=radio] { display: none; }
+.oe-result { border: 1px solid var(--border-color-primary); border-radius: 10px; overflow: hidden; }
+.oe-output, .oe-fields { background: var(--background-fill-secondary); border-bottom: 1px solid var(--border-color-primary); }
+.oe-output { margin: 0; padding: 14px 16px; font-family: var(--font-mono); font-size: 15px; white-space: pre-wrap; word-break: break-word; }
+.oe-fields { padding: 10px 16px; }
+.oe-fields div { display: grid; grid-template-columns: minmax(96px, 28%) minmax(0, 1fr); gap: 12px; padding: 3px 0; font-size: 13px; }
+.oe-fields code { white-space: pre-wrap; word-break: break-word; }
+.oe-fields span, .oe-stats span, .oe-episode li > span { font-family: var(--font-mono); color: var(--body-text-color-subdued); }
+.oe-fields span { overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.oe-stats { display: grid; grid-template-columns: repeat(3, minmax(0, 1fr)); }
+.oe-stats div { padding: 8px 16px; border-right: 1px solid var(--border-color-primary); }
+.oe-stats div:last-child { border-right: none; }
+.oe-stats span { display: block; font-size: 12px; }
+.oe-stats code { display: block; overflow: hidden; text-overflow: ellipsis; white-space: nowrap; }
+.oe-error { color: #b42318; }
+.dark .oe-error { color: #ff8a80; }
+.oe-episode ol { list-style: none; margin: 8px 0 0; padding: 0; }
+.oe-episode li { display: flex; gap: 12px; padding: 8px 0; border-top: 1px solid var(--border-color-primary); }
+.oe-episode li > span { flex: none; width: 20px; font-size: 12px; }
+.oe-episode code { display: block; font-size: 13px; word-break: break-word; }
+.oe-episode small { display: block; font-size: 13px; word-break: break-word; }
+.oe-muted { margin: 8px 0 0; font-size: 14px; }
+.oe-code pre code { white-space: pre-wrap !important; word-break: break-word; }
+.oe-json .code_wrap, .oe-json .cm-editor { max-height: 360px; overflow: auto; }
+.oe-code pre { background: var(--background-fill-secondary) !important; border: 1px solid var(--border-color-primary); border-radius: 8px; }
+.oe-raw { border: none !important; }
+.oe-visual { flex: none !important; width: auto !important; min-width: 0 !important; }
+.dark { --color-accent: #3fb950; }
+footer { display: none !important; }
 """

src/core/env_server/gradio_ui.py

@@ -1,62 +1,37 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 Gradio-based web UI for OpenEnv environments.
 
-Replaces the legacy HTML/JavaScript interface when ENABLE_WEB_INTERFACE is set.
-Mount at /web via gr.mount_gradio_app() from create_web_interface_app().
+Mounted at /web via gr.mount_gradio_app() from create_web_interface_app() when
+ENABLE_WEB_INTERFACE is set. The page follows the loop an agent runs: reset,
+take an action, read the result, with the same call in Python next to it.
 """
 
 from __future__ import annotations
 
+import html
 import json
+import math
 import re
-from typing import Any, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import gradio as gr
 
+from .mcp_environment import MCPEnvironment
+from .mcp_types import ListToolsAction
 from .types import EnvironmentMetadata
 
-
-def _escape_md(text: str) -> str:
-    """Escape Markdown special characters in user-controlled content."""
-    return re.sub(r"([\\`*_\{\}\[\]()#+\-.!|~>])", r"\\\1", str(text))
-
-
-def _format_observation(data: Dict[str, Any]) -> str:
-    """Format reset/step response for Markdown display."""
-    lines: List[str] = []
-    obs = data.get("observation", {})
-    if isinstance(obs, dict):
-        if obs.get("prompt"):
-            lines.append(f"**Prompt:**\n\n{_escape_md(obs['prompt'])}\n")
-        messages = obs.get("messages", [])
-        if messages:
-            lines.append("**Messages:**\n")
-            for msg in messages:
-                sender = _escape_md(str(msg.get("sender_id", "?")))
-                content = _escape_md(str(msg.get("content", "")))
-                cat = _escape_md(str(msg.get("category", "")))
-                lines.append(f"- `[{cat}]` Player {sender}: {content}")
-            lines.append("")
-    reward = data.get("reward")
-    done = data.get("done")
-    if reward is not None:
-        lines.append(f"**Reward:** `{reward}`")
-    if done is not None:
-        lines.append(f"**Done:** `{done}`")
-    return "\n".join(lines) if lines else "*No observation data*"
-
-
-def _readme_section(metadata: Optional[EnvironmentMetadata]) -> str:
-    """README content for the left panel."""
-    if not metadata or not metadata.readme_content:
-        return "*No README available.*"
-    return metadata.readme_content
+_SECRET_NAMES = ("api_key", "token", "secret", "password")
+_LONG_TEXT_NAMES = (
+    "code",
+    "message",
+    "response",
+    "text",
+    "content",
+    "answer",
+    "prompt",
+)
 
 
 def get_gradio_display_title(
@@ -68,6 +43,314 @@ def get_gradio_display_title(
     return f"OpenEnv Agentic Environment: {name}"
 
 
+def _readme_body(readme: str) -> str:
+    """The README without the Hugging Face Space front matter."""
+    return re.sub(r"\A---\n.*?\n---\n", "", readme, flags=re.S)
+
+
+def _description(metadata: Optional[EnvironmentMetadata]) -> str:
+    """The metadata description, or the README's first paragraph when it's the default."""
+    if not metadata:
+        return ""
+    if (
+        not re.fullmatch(r"\S+ environment", metadata.description or "")
+        or not metadata.readme_content
+    ):
+        return metadata.description
+    for paragraph in re.split(r"\n\s*\n", _readme_body(metadata.readme_content)):
+        text = paragraph.strip()
+        if text and not text.startswith(("#", "<", "[!", "!", "|", "```", ">")):
+            return re.sub(r"\[([^\]]+)\]\([^)]+\)", r"\1", " ".join(text.split()))
+    return metadata.description
+
+
+def _header_html(metadata: Optional[EnvironmentMetadata], kind: str) -> str:
+    name = metadata.name if metadata else "environment"
+    title = name.removesuffix("_env").replace("_", " ").title()
+    return (
+        f'<div class="oe-header"><span class="oe-eyebrow">{html.escape(name)} · {kind}</span>'
+        f"<h1>{html.escape(title)}</h1><p>{html.escape(_description(metadata))}</p></div>"
+    )
+
+
+def _step_title(number: int, title: str, text: str) -> str:
+    return f'<div class="oe-step"><h2>{number}. {title}</h2><p>{text}</p></div>'
+
+
+def _rounded(value: Any) -> Any:
+    if isinstance(value, float):
+        return round(value, 4)
+    if isinstance(value, list):
+        return [_rounded(v) for v in value]
+    if isinstance(value, dict):
+        return {k: _rounded(v) for k, v in value.items()}
+    return value
+
+
+def _folded(value: Any) -> Any:
+    """
+    The value with each array of more than 64 numbers (a grid, a screen, a screenshot)
+    shown as its shape, so the page stays light. Smaller boards and lists keep their values.
+    """
+    if isinstance(value, dict):
+        return {k: _folded(v) for k, v in value.items()}
+    if not isinstance(value, list):
+        return value
+    shape, first = [], value
+    while isinstance(first, list) and first:
+        shape.append(len(first))
+        first = first[0]
+    if isinstance(first, (int, float)) and math.prod(shape) > 64:
+        return f"list · {'×'.join(map(str, shape))}"
+    return [_folded(v) for v in value]
+
+
+def _short(value: Any, limit: int = 160) -> str:
+    value = _rounded(value)
+    if isinstance(value, list) and len(value) > 20:
+        text = f"{len(value)} items: {json.dumps(value[:8], default=str)[:-1]}, …]"
+    else:
+        text = value if isinstance(value, str) else json.dumps(value, default=str)
+    return text if len(text) <= limit else text[: limit - 1] + "…"
+
+
+def _error(obs: Dict[str, Any]) -> Optional[str]:
+    """The error an observation reports: a failed tool call, or an env error in metadata."""
+    error = obs.get("error") or (obs.get("metadata") or {}).get("error")
+    if not error:
+        return None
+    return str(error.get("message", error) if isinstance(error, dict) else error)
+
+
+def _output(obs: Dict[str, Any]) -> Any:
+    """What a tool call returned."""
+    result = obs.get("result")
+    return result.get("data", result) if isinstance(result, dict) else result
+
+
+def _result_html(data: Dict[str, Any], step_count: int) -> str:
+    """The observation (a tool's output or its fields), then reward, done and step."""
+    obs = data.get("observation", {}) or {}
+    error = _error(obs)
+    if error:
+        body = (
+            f'<pre class="oe-output oe-error">Error: {html.escape(error[:2000])}</pre>'
+        )
+    elif obs.get("result") is not None:
+        output = _output(obs)
+        text = (
+            output
+            if isinstance(output, str)
+            else json.dumps(_rounded(output), indent=2, default=str)
+        )
+        body = f'<pre class="oe-output">{html.escape(text[:4000])}</pre>'
+    else:
+        rows = "".join(
+            f'<div><span title="{html.escape(k)}">{html.escape(k)}</span>'
+            f"<code>{html.escape(_short(v, 600 if isinstance(v, str) else 160))}</code></div>"
+            for k, v in obs.items()
+            if k not in ("metadata", "done", "reward", "result", "error")
+            and v not in (None, "", [], {})
+        )
+        body = f'<div class="oe-fields">{rows}</div>' if rows else ""
+    reward = data.get("reward")
+    stats = "".join(
+        f"<div><span>{label}</span><code>{html.escape(str(value))}</code></div>"
+        for label, value in (
+            ("reward", "–" if reward is None else _rounded(reward)),
+            ("done", str(bool(data.get("done"))).lower()),
+            ("step", step_count),
+        )
+    )
+    return f'<div class="oe-result">{body}<div class="oe-stats">{stats}</div></div>'
+
+
+def _error_html(message: str, step_count: int) -> str:
+    return _result_html({"observation": {"error": message}}, step_count)
+
+
+def _episode_html(entries: List[List[str]]) -> str:
+    if not entries:
+        items = '<p class="oe-muted">Nothing yet. Reset, then run a step.</p>'
+    else:
+        items = (
+            "<ol>"
+            + "".join(
+                f"<li><span>{html.escape(badge)}</span><div><code>{html.escape(call)}</code>"
+                f"<small>{html.escape(result)}</small></div></li>"
+                for badge, call, result in entries
+            )
+            + "</ol>"
+        )
+    return f'<div class="oe-episode"><h2>Episode</h2>{items}</div>'
+
+
+def _list_tools(env: Any) -> List[Any]:
+    return list(env.step(ListToolsAction()).tools)
+
+
+def _tool_label(tool: Any) -> str:
+    """The tool's name and the first line of its description, if it has one."""
+    first_line = (tool.description or "").strip().split("\n")[0]
+    return f"{tool.name}: {first_line}" if first_line else tool.name
+
+
+def _params(schema: Dict[str, Any]) -> List[Tuple[str, Dict[str, Any], bool]]:
+    """
+    The (name, schema, required) of each property of a JSON schema, with `$ref`s
+    resolved and `X | None` unwrapped to `X`.
+    """
+    defs = schema.get("$defs", {})
+
+    def resolve(prop: Dict[str, Any]) -> Dict[str, Any]:
+        if "$ref" in prop:
+            prop = {
+                **defs.get(prop["$ref"].split("/")[-1], {}),
+                **{k: v for k, v in prop.items() if k != "$ref"},
+            }
+        for key in ("anyOf", "oneOf", "allOf"):
+            options = [o for o in prop.get(key, []) if o.get("type") != "null"]
+            if options:
+                merged = {k: v for k, v in prop.items() if k != key}
+                return resolve({**resolve(options[0]), **merged})
+        return prop
+
+    required = set(schema.get("required", []))
+    return [
+        (name, resolve(prop), name in required)
+        for name, prop in schema.get("properties", {}).items()
+        if name != "metadata"
+    ]
+
+
+def _call_args(arguments: Dict[str, Any]) -> str:
+    """Arguments as `name=value` for the episode log, with secrets masked and long values cut."""
+    return ", ".join(
+        f"{k}={'***' if k.lower().endswith(_SECRET_NAMES) else _short(json.dumps(v), 60)}"
+        for k, v in arguments.items()
+    )
+
+
+def _blank(value: Any) -> bool:
+    return value is None or (isinstance(value, str) and not value.strip())
+
+
+def _param_input(
+    name: str, schema: Dict[str, Any], required: bool
+) -> Tuple[Any, Callable[[Any], Any]]:
+    """
+    A form input for one parameter, and a function that turns its value into
+    the one to send (`None` to leave the parameter out).
+    """
+    kind = schema.get("type", "string")
+    default = schema.get("default")
+    is_json = kind in ("array", "object")
+    label = f"{name} ({kind}{', JSON' if is_json else ''}){' · required' if required else ''}"
+    info = (schema.get("description") or "").strip()[:200] or None
+
+    if "enum" in schema:
+        # A required choice starts on its first option, so the form can run as is.
+        value = schema["enum"][0] if required else default
+        return (
+            gr.Dropdown(choices=schema["enum"], value=value, label=label, info=info),
+            lambda v: None if _blank(v) else v,
+        )
+    if kind == "boolean":
+        return gr.Checkbox(label=label, value=bool(default), info=info), bool
+    if kind in ("integer", "number"):
+        cast = int if kind == "integer" else float
+
+        def number(value: Any) -> Any:
+            if _blank(value):
+                return None
+            try:
+                return cast(float(value))
+            except ValueError:
+                raise ValueError(f"{name} must be a number")
+
+        # Gradio shows an empty number box as 0, so optional numbers without a default are text.
+        if default is None and not required:
+            return gr.Textbox(label=label, info=info, placeholder="optional"), number
+        return gr.Number(
+            label=label,
+            value=default,
+            info=info,
+            precision=0 if kind == "integer" else None,
+        ), number
+    if is_json:
+
+        def parse(value: Any) -> Any:
+            if _blank(value):
+                return None
+            try:
+                return json.loads(value)
+            except json.JSONDecodeError as e:
+                raise ValueError(f"{name} must be JSON: {e}")
+
+        return (
+            gr.Textbox(
+                label=label,
+                info=info,
+                lines=3,
+                value=None if default is None else json.dumps(default),
+                placeholder="[ ... ]" if kind == "array" else "{ ... }",
+            ),
+            parse,
+        )
+    long_text = (
+        any(n in name.lower() for n in _LONG_TEXT_NAMES)
+        or schema.get("maxLength", 0) > 100
+    )
+    return (
+        gr.Textbox(
+            label=label,
+            info=info,
+            value=default if isinstance(default, str) else None,
+            lines=4 if long_text else 1,
+            type="password" if name.lower().endswith(_SECRET_NAMES) else "text",
+        ),
+        lambda v: None if _blank(v) else v,
+    )
+
+
+class _Form:
+    """Inputs for a set of parameters, and the dict to send built from their values."""
+
+    def __init__(self, params: List[Tuple[str, Dict[str, Any], bool]]):
+        required = [p for p in params if p[2]]
+        optional = [p for p in params if not p[2]]
+        self.names: List[Tuple[str, bool]] = []
+        self.inputs: List[Any] = []
+        self.converters: List[Callable[[Any], Any]] = []
+        self._add(required)
+        if required and optional:
+            with gr.Accordion("More fields", open=False, elem_classes="oe-raw"):
+                self._add(optional)
+        else:
+            self._add(optional)
+
+    def _add(self, params: List[Tuple[str, Dict[str, Any], bool]]) -> None:
+        for name, schema, required in params:
+            component, convert = _param_input(name, schema, required)
+            self.names.append((name, required))
+            self.inputs.append(component)
+            self.converters.append(convert)
+
+    def values(self, raw: List[Any]) -> Dict[str, Any]:
+        """The values to send. Raises `ValueError` for a missing or invalid one."""
+        values, missing = {}, []
+        for (name, required), convert, value in zip(self.names, self.converters, raw):
+            converted = convert(value)
+            if converted is None:
+                if required:
+                    missing.append(name)
+                continue
+            values[name] = converted
+        if missing:
+            raise ValueError(f"Fill in {', '.join(missing)}.")
+        return values
+
+
 def build_gradio_app(
     web_manager: Any,
     action_fields: List[Dict[str, Any]],
@@ -81,160 +364,347 @@ def build_gradio_app(
 
     Args:
         web_manager: WebInterfaceManager (reset/step_environment, get_state).
-        action_fields: Field dicts from _extract_action_fields(action_cls).
+        action_fields: Field dicts from _extract_action_fields(action_cls). The form is
+            built from the action's JSON schema, so this is kept for `gradio_builder`s.
         metadata: Environment metadata for README/name.
-        is_chat_env: If True, single message textbox; else form from action_fields.
+        is_chat_env: If True, a single message box; else a form from the action schema.
         title: App title (overridden by metadata.name when present; see get_gradio_display_title).
         quick_start_md: Optional Quick Start markdown (class names already replaced).
 
     Returns:
         gr.Blocks to mount with gr.mount_gradio_app(app, blocks, path="/web").
     """
-    readme_content = _readme_section(metadata)
-    display_title = get_gradio_display_title(metadata, fallback=title)
+    env = web_manager.env
+    tools = _list_tools(env) if isinstance(env, MCPEnvironment) else []
+    is_mcp = bool(tools)
+    kind = "MCP environment" if is_mcp else "environment"
 
-    async def reset_env():
+    def render(data: Dict[str, Any]):
+        """The env's drawing and its one-click actions for this observation."""
+        obs = {
+            **(data.get("observation") or {}),
+            "reward": data.get("reward"),
+            "done": data.get("done"),
+        }
+        drawing = env.render_web(obs)
+        actions = env.web_actions(obs) if not obs["done"] else []
+        return (
+            gr.update(value=drawing or "", visible=bool(drawing)),
+            gr.update(
+                choices=[(label, str(i)) for i, (label, _) in enumerate(actions)],
+                value=None,
+                visible=bool(actions),
+            ),
+            [action for _, action in actions],
+        )
+
+    def steps_in(entries: List[List[str]]) -> int:
+        return sum(1 for badge, _, _ in entries if badge != "R")
+
+    def unchanged(message: str, entries: List[List[str]]):
+        """Show an error without touching the episode."""
+        return (
+            gr.update(),
+            gr.update(),
+            gr.update(),
+            _error_html(message, steps_in(entries)),
+            gr.update(),
+            gr.update(),
+            entries,
+        )
+
+    async def reset_env(entries):
         try:
             data = await web_manager.reset_environment()
-            obs_md = _format_observation(data)
+        except Exception as e:
+            # Keep the current episode and the last step's result.
+            (
+                visual_update,
+                quick_update,
+                actions,
+                result_html,
+                raw,
+                episode_html,
+                kept,
+            ) = unchanged(f"reset() failed: {e}", entries)
             return (
-                obs_md,
-                json.dumps(data, indent=2),
-                "Environment reset successfully.",
+                visual_update,
+                quick_update,
+                actions,
+                result_html,
+                raw,
+                episode_html,
+                kept,
+                gr.update(),
             )
-        except Exception as e:
-            return ("", "", f"Error: {e}")
+        error = _error(data.get("observation", {}) or {})
+        entries = [
+            ["R", "reset()", f"error: {_short(error)}" if error else "new episode"]
+        ]
+        visual_update, quick_update, actions = render(data)
+        shown = _folded(data)
+        return (
+            visual_update,
+            quick_update,
+            actions,
+            _result_html(shown, 0),
+            json.dumps(shown, indent=2, default=str, ensure_ascii=False),
+            _episode_html(entries),
+            entries,
+            "",
+        )
 
-    def _step_with_action(action_data: Dict[str, Any]):
-        async def _run():
+    async def run_step(action: Dict[str, Any], call: str, entries):
+        if not entries:
             try:
-                data = await web_manager.step_environment(action_data)
-                obs_md = _format_observation(data)
-                return (
-                    obs_md,
-                    json.dumps(data, indent=2),
-                    "Step complete.",
-                )
+                await web_manager.reset_environment()
             except Exception as e:
-                return ("", "", f"Error: {e}")
-
-        return _run
-
-    async def step_chat(message: str):
-        if not (message or str(message).strip()):
-            return ("", "", "Please enter an action message.")
-        action = {"message": str(message).strip()}
-        return await _step_with_action(action)()
-
-    def get_state_sync():
+                return unchanged(f"reset() failed: {e}", entries)
+            entries = [["R", "reset()", "new episode"]]
         try:
-            data = web_manager.get_state()
-            return json.dumps(data, indent=2)
+            data = await web_manager.step_environment(action)
         except Exception as e:
-            return f"Error: {e}"
+            entries = entries + [
+                [str(steps_in(entries) + 1), call, f"error: {_short(str(e))}"]
+            ]
+            return (
+                gr.update(),
+                gr.update(),
+                gr.update(),
+                _error_html(str(e), steps_in(entries)),
+                gr.update(),
+                _episode_html(entries),
+                entries,
+            )
+        obs = data.get("observation", {}) or {}
+        error = _error(obs)
+        summary = (
+            f"error: {error}"
+            if error
+            else _output(obs)
+            if obs.get("result") is not None
+            else f"reward {data.get('reward')}"
+        )
+        entries = entries + [[str(steps_in(entries) + 1), call, _short(summary)]]
+        visual_update, quick_update, actions = render(data)
+        shown = _folded(data)
+        return (
+            visual_update,
+            quick_update,
+            actions,
+            _result_html(shown, steps_in(entries)),
+            json.dumps(shown, indent=2, default=str, ensure_ascii=False),
+            _episode_html(entries),
+            entries,
+        )
 
-    with gr.Blocks(title=display_title) as demo:
-        with gr.Row():
-            with gr.Column(scale=1, elem_classes="col-left"):
-                if quick_start_md:
-                    with gr.Accordion("Quick Start", open=True):
-                        gr.Markdown(quick_start_md)
-                with gr.Accordion("README", open=False):
-                    gr.Markdown(readme_content)
-
-            with gr.Column(scale=2, elem_classes="col-right"):
-                obs_display = gr.Markdown(
-                    value=("# Playground\n\nClick **Reset** to start a new episode."),
-                )
-                with gr.Group():
-                    if is_chat_env:
-                        action_input = gr.Textbox(
-                            label="Action message",
-                            placeholder="e.g. Enter your message...",
+    with gr.Blocks(title=get_gradio_display_title(metadata, fallback=title)) as demo:
+        entries_state = gr.State([])
+        gr.HTML(_header_html(metadata, kind))
+        with gr.Row(equal_height=False):
+            with gr.Column(scale=3, min_width=360):
+                with gr.Column(elem_classes="oe-card"):
+                    with gr.Row(elem_classes="oe-row", equal_height=True):
+                        gr.HTML(
+                            _step_title(
+                                1,
+                                "Start an episode",
+                                "<code>reset()</code> returns the first observation.",
+                            )
+                        )
+                        reset_btn = gr.Button(
+                            "Reset", variant="secondary", scale=0, min_width=110
+                        )
+                    reset_result = gr.HTML()
+
+                with gr.Column(elem_classes="oe-card"):
+                    if is_mcp:
+                        gr.HTML(
+                            _step_title(
+                                2,
+                                "Call a tool",
+                                "<code>step()</code> sends a tool call. Pick one:",
+                            )
                         )
-                        step_inputs = [action_input]
-                        step_fn = step_chat
                     else:
-                        step_inputs = []
-                        for field in action_fields:
-                            name = field["name"]
-                            field_type = field.get("type", "text")
-                            label = name.replace("_", " ").title()
-                            placeholder = field.get("placeholder", "")
-                            if field_type == "checkbox":
-                                inp = gr.Checkbox(label=label)
-                            elif field_type == "number":
-                                inp = gr.Number(label=label)
-                            elif field_type == "select":
-                                choices = field.get("choices") or []
-                                inp = gr.Dropdown(
-                                    choices=choices,
-                                    label=label,
-                                    allow_custom_value=False,
+                        gr.HTML(
+                            _step_title(
+                                2,
+                                "Take an action",
+                                "<code>step()</code> sends an action.",
+                            )
+                        )
+                    with gr.Row(equal_height=False):
+                        visual = gr.HTML(visible=False, elem_classes="oe-visual")
+                        with gr.Column(min_width=260):
+                            quick = gr.Radio(
+                                choices=[],
+                                visible=False,
+                                label="Legal actions",
+                                elem_classes="oe-actions",
+                            )
+                            if is_mcp:
+                                tool_choice = gr.Radio(
+                                    choices=[(_tool_label(t), t.name) for t in tools],
+                                    value=tools[0].name,
+                                    show_label=False,
+                                    elem_classes="oe-tools",
                                 )
-                            elif field_type in ("textarea", "tensor"):
-                                inp = gr.Textbox(
-                                    label=label,
-                                    placeholder=placeholder,
-                                    lines=3,
+                                forms, groups = {}, []
+                                for i, tool in enumerate(tools):
+                                    with gr.Column(visible=i == 0) as group:
+                                        forms[tool.name] = _Form(
+                                            _params(tool.input_schema)
+                                        )
+                                    groups.append(group)
+                                tool_choice.change(
+                                    lambda name: [
+                                        gr.update(visible=t.name == name) for t in tools
+                                    ],
+                                    inputs=tool_choice,
+                                    outputs=groups,
                                 )
+                                all_inputs = [
+                                    i for f in forms.values() for i in f.inputs
+                                ]
+
+                                async def step_fn(entries, name, *values):
+                                    start = 0
+                                    for tool_name, form in forms.items():
+                                        if tool_name == name:
+                                            break
+                                        start += len(form.inputs)
+                                    form = forms[name]
+                                    try:
+                                        arguments = form.values(
+                                            list(
+                                                values[start : start + len(form.inputs)]
+                                            )
+                                        )
+                                    except ValueError as e:
+                                        return unchanged(str(e), entries)
+                                    action = {
+                                        "type": "call_tool",
+                                        "tool_name": name,
+                                        "arguments": arguments,
+                                    }
+                                    return await run_step(
+                                        action,
+                                        f"{name}({_call_args(arguments)})",
+                                        entries,
+                                    )
+
+                                step_inputs = [entries_state, tool_choice, *all_inputs]
                             else:
-                                inp = gr.Textbox(
-                                    label=label,
-                                    placeholder=placeholder,
-                                )
-                            step_inputs.append(inp)
-
-                        async def step_form(*values):
-                            if not action_fields:
-                                return await _step_with_action({})()
-                            action_data = {}
-                            for i, field in enumerate(action_fields):
-                                if i >= len(values):
-                                    break
-                                name = field["name"]
-                                val = values[i]
-                                if field.get("type") == "checkbox":
-                                    action_data[name] = bool(val)
-                                elif val is not None and val != "":
-                                    action_data[name] = val
-                            return await _step_with_action(action_data)()
-
-                        step_fn = step_form
-
-                    with gr.Row():
-                        step_btn = gr.Button("Step", variant="primary")
-                        reset_btn = gr.Button("Reset", variant="secondary")
-                        state_btn = gr.Button("Get state", variant="secondary")
-                    with gr.Row():
-                        status = gr.Textbox(
-                            label="Status",
-                            interactive=False,
+                                if is_chat_env:
+                                    params = [
+                                        (
+                                            "message",
+                                            {
+                                                "type": "string",
+                                                "description": "Your message to the model.",
+                                            },
+                                            True,
+                                        )
+                                    ]
+                                else:
+                                    params = [
+                                        p
+                                        for p in _params(
+                                            web_manager.action_cls.model_json_schema()
+                                        )
+                                        if p[0] != "type"
+                                    ]
+                                form = _Form(params)
+
+                                async def step_fn(entries, *values):
+                                    try:
+                                        action = form.values(list(values))
+                                    except ValueError as e:
+                                        return unchanged(str(e), entries)
+                                    return await run_step(
+                                        action, f"step({_call_args(action)})", entries
+                                    )
+
+                                step_inputs = [entries_state, *form.inputs]
+
+                            step_btn = gr.Button(
+                                "Run step", variant="primary", elem_classes="oe-run"
+                            )
+                    result = gr.HTML()
+                    with gr.Accordion(
+                        "raw JSON", open=False, elem_classes="oe-raw oe-json"
+                    ):
+                        raw_json = gr.Code(
+                            language="json", interactive=False, show_label=False
                         )
-                    raw_json = gr.Code(
-                        label="Raw JSON response",
-                        language="json",
-                        interactive=False,
-                    )
 
+            with gr.Column(scale=2, min_width=320):
+                with gr.Column(elem_classes="oe-card"):
+                    episode = gr.HTML(_episode_html([]))
+                    with gr.Accordion(
+                        "State", open=False, elem_classes="oe-raw oe-json"
+                    ):
+                        state_json = gr.Code(
+                            language="json", interactive=False, show_label=False
+                        )
+                if quick_start_md:
+                    with gr.Column(elem_classes="oe-card oe-code"):
+                        gr.HTML(
+                            '<div class="oe-step"><h2>Same call in Python</h2></div>'
+                        )
+                        gr.Markdown(
+                            quick_start_md.replace(
+                                "### Connect to this environment\n\n", ""
+                            )
+                        )
+
+        if metadata and metadata.readme_content:
+            with gr.Accordion("README", open=False, elem_classes="oe-raw"):
+                gr.Markdown(_readme_body(metadata.readme_content))
+
+        quick_state = gr.State([])
+
+        def show_state():
+            try:
+                return json.dumps(web_manager.get_state(), indent=2, default=str)
+            except Exception:
+                return ""
+
+        async def quick_step(entries, actions, choice):
+            action = actions[int(choice)]
+            return await run_step(action, f"step({_call_args(action)})", entries)
+
+        outputs = [visual, quick, quick_state, result, raw_json, episode, entries_state]
+        reset_outputs = [
+            visual,
+            quick,
+            quick_state,
+            reset_result,
+            raw_json,
+            episode,
+            entries_state,
+            result,
+        ]
+        # One call to the environment at a time, so a slow reset can't overlap a step.
+        busy = dict(concurrency_id="env", concurrency_limit=1, trigger_mode="once")
         reset_btn.click(
-            fn=reset_env,
-            outputs=[obs_display, raw_json, status],
-        )
-        step_btn.click(
-            fn=step_fn,
+            reset_env, inputs=[entries_state], outputs=reset_outputs, **busy
+        ).then(show_state, outputs=state_json)
+        # Enter submits a one-line text field, Shift+Enter a multi-line one.
+        gr.on(
+            [step_btn.click]
+            + [i.submit for i in step_inputs if isinstance(i, gr.Textbox)],
+            step_fn,
             inputs=step_inputs,
-            outputs=[obs_display, raw_json, status],
-        )
-        if is_chat_env:
-            action_input.submit(
-                fn=step_fn,
-                inputs=step_inputs,
-                outputs=[obs_display, raw_json, status],
-            )
-        state_btn.click(
-            fn=get_state_sync,
-            outputs=[raw_json],
-        )
+            outputs=outputs,
+            **busy,
+        ).then(show_state, outputs=state_json)
+        quick.input(
+            quick_step,
+            inputs=[entries_state, quick_state, quick],
+            outputs=outputs,
+            **busy,
+        ).then(show_state, outputs=state_json)
 
     return demo

src/core/env_server/http_server.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 HTTP server wrapper for Environment instances.
@@ -22,9 +18,20 @@ import time
 import uuid
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import AsyncExitStack
-from typing import Any, AsyncContextManager, Callable, cast, Dict, Optional, Type
+from typing import (
+    Any,
+    AsyncContextManager,
+    Awaitable,
+    Callable,
+    cast,
+    Dict,
+    Optional,
+    Type,
+    TypeVar,
+)
 
 _MISSING = object()
+_MCPResult = TypeVar("_MCPResult")
 
 from fastapi import (
     Body,
@@ -37,12 +44,17 @@ from fastapi import (
 )
 from pydantic import ValidationError
 
+from ._utils import overrides_method
 from .interfaces import Environment
 from .mcp_environment import get_server_tools
 from .mcp_types import (
+    CallToolAction,
+    CallToolObservation,
     JsonRpcErrorCode,
     JsonRpcRequest,
     JsonRpcResponse,
+    ListToolsAction,
+    ListToolsObservation,
     McpMethod,
     WSMCPMessage,
     WSMCPResponse,
@@ -53,8 +65,12 @@ from .types import (
     Action,
     ConcurrencyConfig,
     EnvironmentMetadata,
+    GetTaskRangeRequest,
+    GetTaskRequest,
     HealthResponse,
     HealthStatus,
+    ListTasksRequest,
+    NumTasksRequest,
     Observation,
     ResetRequest,
     ResetResponse,
@@ -83,10 +99,11 @@ def _make_json_serializable(obj: Any) -> Any:
     Handles Pydantic models, dataclasses, and other common types.
 
     Args:
-        obj: The object to convert
+        obj (`Any`):
+            The object to convert.
 
     Returns:
-        A JSON-serializable representation of the object
+        `Any`: A JSON-serializable representation of the object.
     """
     if obj is None:
         return None
@@ -106,6 +123,13 @@ def _make_json_serializable(obj: Any) -> Any:
     return str(obj)
 
 
+async def _maybe_await(value: Any) -> Any:
+    """Await values returned by async task APIs while preserving sync APIs."""
+    if inspect.isawaitable(value):
+        return await value
+    return value
+
+
 from .exceptions import (
     ConcurrencyConfigurationError,
     EnvironmentFactoryError,
@@ -124,23 +148,26 @@ class HTTPEnvServer:
     - Action deserialization: Converts JSON dict to Action subclass
     - Observation serialization: Converts Observation subclass to JSON dict
 
-    Example:
-        >>> from core.env_server import HTTPEnvServer
-        >>> from envs.coding_env.server import CodeExecutionEnvironment
-        >>> from envs.coding_env.models import CodeAction, CodeObservation
-        >>>
-        >>> # Pass environment class (factory pattern)
-        >>> server = HTTPEnvServer(
-        ...     env=CodeExecutionEnvironment,
-        ...     action_cls=CodeAction,
-        ...     observation_cls=CodeObservation,
-        ...     max_concurrent_envs=4,
-        ... )
-        >>>
-        >>> # Register routes with FastAPI
-        >>> from fastapi import FastAPI
-        >>> app = FastAPI()
-        >>> server.register_routes(app)
+    Examples:
+
+        ```python
+        from core.env_server import HTTPEnvServer
+        from envs.coding_env.server import CodeExecutionEnvironment
+        from envs.coding_env.models import CodeAction, CodeObservation
+
+        # Pass environment class (factory pattern)
+        server = HTTPEnvServer(
+            env=CodeExecutionEnvironment,
+            action_cls=CodeAction,
+            observation_cls=CodeObservation,
+            max_concurrent_envs=4,
+        )
+
+        # Register routes with FastAPI
+        from fastapi import FastAPI
+        app = FastAPI()
+        server.register_routes(app)
+        ```
     """
 
     def __init__(
@@ -150,24 +177,42 @@ class HTTPEnvServer:
         observation_cls: Type[Observation],
         max_concurrent_envs: Optional[int] = None,
         concurrency_config: Optional[ConcurrencyConfig] = None,
+        env_name: Optional[str] = None,
+        state_cls: Type[State] = State,
+        *,
+        reset_observation_cls: Optional[Type[Observation]] = None,
     ):
         """
         Initialize HTTP server wrapper.
 
         Args:
-            env: Environment factory (callable) that creates new instances.
-                 Will be called to create a new environment for each WebSocket session.
-            action_cls: The Action subclass this environment expects
-            observation_cls: The Observation subclass this environment returns
-            max_concurrent_envs: Maximum number of concurrent WebSocket sessions.
-                                 Mutually exclusive with concurrency_config.
-            concurrency_config: Optional ConcurrencyConfig for advanced concurrency settings.
-                                Mutually exclusive with max_concurrent_envs.
+            env (`Callable[[], Environment]`):
+                Environment factory (callable) that creates new instances. Will be called
+                to create a new environment for each WebSocket session.
+            action_cls (`Type[Action]`):
+                The `Action` subclass this environment expects.
+            observation_cls (`Type[Observation]`):
+                The `Observation` subclass returned by step.
+            reset_observation_cls (`Type[Observation]`, *optional*):
+                The reset observation model published in `/schema`. Defaults to
+                `observation_cls`; declare a distinct model when reset differs.
+            max_concurrent_envs (`int`, *optional*):
+                Maximum number of concurrent WebSocket sessions. Mutually exclusive with
+                `concurrency_config`.
+            concurrency_config (`ConcurrencyConfig`, *optional*):
+                Advanced concurrency settings. Mutually exclusive with
+                `max_concurrent_envs`.
+            env_name (`str`, *optional*):
+                Public environment name used by task/split endpoints.
+            state_cls (`Type[State]`, *optional*, defaults to `State`):
+                The `State` subclass this environment reports. Used for the `/state`
+                response model and the `state` entry of `/schema`, so that fields
+                declared by the subclass are published and serialized.
 
         Raises:
-            ValueError: If both max_concurrent_envs and concurrency_config are provided.
-            ConcurrencyConfigurationError: If max_concurrent_envs > 1 for an
-                environment that is not marked as SUPPORTS_CONCURRENT_SESSIONS.
+            `ValueError`: If both `max_concurrent_envs` and `concurrency_config` are provided.
+            `ConcurrencyConfigurationError`: If `max_concurrent_envs` > 1 for an environment
+                that is not marked as `SUPPORTS_CONCURRENT_SESSIONS`.
         """
         # Validate that env is callable
         if not callable(env):
@@ -206,18 +251,28 @@ class HTTPEnvServer:
 
         self.action_cls = action_cls
         self.observation_cls = observation_cls
+        self.reset_observation_cls = reset_observation_cls or observation_cls
+        self.state_cls = state_cls
+        self.env_name = env_name or self._default_env_name()
 
         # Session management for WebSocket connections
         self._sessions: Dict[str, Optional[Environment]] = {}
         self._session_executors: Dict[str, ThreadPoolExecutor] = {}
         self._session_stacks: Dict[str, AsyncExitStack] = {}
         self._session_info: Dict[str, SessionInfo] = {}
+        self._session_websocket_attachments: set[str] = set()
+        self._session_pending_closes: set[str] = set()
         self._session_lock = asyncio.Lock()
 
         # Create thread pool for running sync code in async context
         # This is needed for environments using sync libraries (e.g., Playwright)
         self._executor = ThreadPoolExecutor(max_workers=32)
 
+        self._requires_single_thread_executor = self._detect_single_thread_requirement()
+        self._shared_session_executor: Optional[ThreadPoolExecutor] = None
+        if self._requires_single_thread_executor:
+            self._shared_session_executor = ThreadPoolExecutor(max_workers=1)
+
         # Idle session reaper configuration.
         # Timeout is taken from ConcurrencyConfig.session_timeout;
         # None means no timeout (default — reaper is a no-op).
@@ -226,19 +281,31 @@ class HTTPEnvServer:
         )
         self._reaper_task: Optional[asyncio.Task[None]] = None
 
+    def _default_env_name(self) -> str:
+        factory = self._env_factory
+        if inspect.isclass(factory):
+            return factory.__name__
+        return getattr(factory, "__name__", "environment")
+
     def _validate_concurrency_safety(self) -> None:
         """
         Validate that the environment supports the configured concurrency level.
 
         Raises:
-            ConcurrencyConfigurationError: If max_concurrent_envs > 1 for an
-                environment that is not marked as SUPPORTS_CONCURRENT_SESSIONS.
+            `ConcurrencyConfigurationError`: If `max_concurrent_envs` > 1 for an environment
+                that is not marked as `SUPPORTS_CONCURRENT_SESSIONS`.
         """
+        import functools
+
         if self._max_concurrent_envs <= 1:
             return
 
-        if inspect.isclass(self._env_factory):
-            env_cls = self._env_factory
+        factory = self._env_factory
+        if isinstance(factory, functools.partial):
+            factory = factory.func
+
+        if inspect.isclass(factory):
+            env_cls = factory
         else:
             _temp_env = self._env_factory()
             env_cls = type(_temp_env)
@@ -251,12 +318,22 @@ class HTTPEnvServer:
                 max_concurrent_envs=self._max_concurrent_envs,
             )
 
+    def _detect_single_thread_requirement(self) -> bool:
+        import functools
+
+        factory = self._env_factory
+        if isinstance(factory, functools.partial):
+            factory = factory.func
+        if inspect.isclass(factory):
+            return getattr(factory, "REQUIRES_SINGLE_THREAD_EXECUTOR", False)
+        return False
+
     def get_capacity_status(self) -> ServerCapacityStatus:
         """
         Get the current capacity status of the server.
 
         Returns:
-            ServerCapacityStatus with current session counts and availability.
+            [`ServerCapacityStatus`] with current session counts and availability.
         """
         return ServerCapacityStatus.from_counts(
             active=len(self._sessions),
@@ -298,11 +375,11 @@ class HTTPEnvServer:
         Create a new WebSocket session with its own environment instance.
 
         Returns:
-            Tuple of (session_id, environment)
+            `tuple[str, Environment]`: Tuple of (session_id, environment).
 
         Raises:
-            SessionCapacityError: If max concurrent sessions reached
-            EnvironmentFactoryError: If the factory fails to create an environment
+            `SessionCapacityError`: If max concurrent sessions reached.
+            `EnvironmentFactoryError`: If the factory fails to create an environment.
         """
         async with self._session_lock:
             if len(self._sessions) >= self._max_concurrent_envs:
@@ -316,7 +393,10 @@ class HTTPEnvServer:
 
             # Create executor and reserve slot so capacity is not exceeded while
             # we create the env outside the lock (avoids blocking other sessions)
-            executor = ThreadPoolExecutor(max_workers=1)
+            if self._shared_session_executor is not None:
+                executor = self._shared_session_executor
+            else:
+                executor = ThreadPoolExecutor(max_workers=1)
             self._session_executors[session_id] = executor
             self._sessions[session_id] = None  # placeholder until env is ready
 
@@ -326,7 +406,8 @@ class HTTPEnvServer:
             env = await loop.run_in_executor(executor, self._env_factory)
         except Exception as e:
             async with self._session_lock:
-                executor.shutdown(wait=False)
+                if executor is not self._shared_session_executor:
+                    executor.shutdown(wait=False)
                 self._session_executors.pop(session_id, None)
                 self._sessions.pop(session_id, None)
             factory_name = getattr(
@@ -376,13 +457,16 @@ class HTTPEnvServer:
         Destroy a WebSocket session and cleanup resources.
 
         Args:
-            session_id: The session ID to destroy
+            session_id (`str`):
+                The session ID to destroy.
         """
         async with self._session_lock:
             env = self._sessions.pop(session_id, None)
             executor = self._session_executors.pop(session_id, None)
             stack = self._session_stacks.pop(session_id, None)
             self._session_info.pop(session_id, None)
+            self._session_websocket_attachments.discard(session_id)
+            self._session_pending_closes.discard(session_id)
 
         await self._cleanup_session_resources(env, executor, stack)
 
@@ -421,8 +505,7 @@ class HTTPEnvServer:
                 except Exception:
                     pass  # Best effort cleanup
 
-        # Shutdown executor after close is done
-        if executor is not None:
+        if executor is not None and executor is not self._shared_session_executor:
             executor.shutdown(wait=False)
 
     def _update_session_activity(
@@ -432,8 +515,10 @@ class HTTPEnvServer:
         Update session activity timestamp and optionally increment step count.
 
         Args:
-            session_id: The session ID to update
-            increment_step: If True, increment the step count
+            session_id (`str`):
+                The session ID to update.
+            increment_step (`bool`, *optional*, defaults to `False`):
+                If `True`, increment the step count.
         """
         if session_id in self._session_info:
             self._session_info[session_id].last_activity_at = time.time()
@@ -453,7 +538,10 @@ class HTTPEnvServer:
                 stale_ids: list[str] = []
                 async with self._session_lock:
                     for sid, info in self._session_info.items():
-                        if now - info.last_activity_at > timeout:
+                        if (
+                            sid not in self._session_websocket_attachments
+                            and now - info.last_activity_at > timeout
+                        ):
                             stale_ids.append(sid)
                 for sid in stale_ids:
                     # Re-check under lock: activity may have arrived since
@@ -463,7 +551,11 @@ class HTTPEnvServer:
                     now = time.time()
                     async with self._session_lock:
                         info = self._session_info.get(sid)
-                        if info is None or (now - info.last_activity_at) <= timeout:
+                        if (
+                            info is None
+                            or sid in self._session_websocket_attachments
+                            or (now - info.last_activity_at) <= timeout
+                        ):
                             continue
                     await self._destroy_session(sid)
             except asyncio.CancelledError:
@@ -490,10 +582,11 @@ class HTTPEnvServer:
         Get information about a specific session.
 
         Args:
-            session_id: The session ID to query
+            session_id (`str`):
+                The session ID to query.
 
         Returns:
-            SessionInfo if the session exists, None otherwise
+            [`SessionInfo`] if the session exists, `None` otherwise.
         """
         return self._session_info.get(session_id)
 
@@ -505,6 +598,25 @@ class HTTPEnvServer:
         loop = asyncio.get_event_loop()
         return await loop.run_in_executor(executor, lambda: func(*args, **kwargs))
 
+    async def _run_mcp_client_operation(
+        self,
+        mcp_client: Any,
+        mcp_session_factory: Any,
+        managed_session_id: Optional[str],
+        operation: Callable[[], Awaitable[_MCPResult]],
+    ) -> _MCPResult:
+        """Run an MCP client operation with the appropriate transport lifecycle."""
+        if managed_session_id and mcp_client.is_connected():
+            return await operation()
+
+        if callable(mcp_session_factory):
+            mcp_session_cm = cast(AsyncContextManager[Any], mcp_session_factory())
+            async with mcp_session_cm:
+                return await operation()
+
+        async with mcp_client:
+            return await operation()
+
     @property
     def active_sessions(self) -> int:
         """Return the number of active WebSocket sessions."""
@@ -534,6 +646,212 @@ class HTTPEnvServer:
         """Return the concurrency configuration."""
         return self._concurrency_config
 
+    def _factory_produces_harness_env(self) -> bool:
+        """Return whether the env factory produces a HarnessEnvironment."""
+        import inspect
+
+        # Lazy import: openenv.core.harness imports env_server modules, so a
+        # top-level import here would be circular.
+        from ..harness.environment import HarnessEnvironment
+
+        if inspect.isclass(self._env_factory):
+            return issubclass(self._env_factory, HarnessEnvironment)
+        _temp_env = self._env_factory()
+        try:
+            return isinstance(_temp_env, HarnessEnvironment)
+        finally:
+            _temp_env.close()
+
+    def _register_harness_route(self, app: FastAPI) -> None:
+        """
+        Register the production `/harness` WebSocket route (RFC 005).
+
+        Each connection gets its own environment session: connecting resets
+        the environment (which starts the harness process and injects tools),
+        and each `{"type": "message", "content": ...}` frame runs one
+        conversational turn, streamed back as `HarnessEvent` JSON frames
+        ending with a `turn_complete` event. Malformed client frames receive a
+        recoverable `protocol_error` response without starting a turn; the
+        connection remains usable. Terminal failures use `error`.
+        """
+        # Lazy import to avoid a circular import with openenv.core.harness.
+        from ..harness.adapter import HarnessNotRunningError
+        from ..harness.events import (
+            HarnessClientMessage,
+            HarnessEvent,
+            HarnessEventType,
+            HarnessProtocolError,
+        )
+
+        @app.websocket("/harness")
+        async def harness_websocket_endpoint(websocket: WebSocket):
+            await websocket.accept()
+
+            session_id = None
+            session_env = None
+
+            async def send_error(message: str, code: WSErrorCode) -> None:
+                error_response = WSErrorResponse(
+                    data={"message": message, "code": code}
+                )
+                await websocket.send_text(error_response.model_dump_json())
+
+            async def send_protocol_error(message: str, code: WSErrorCode) -> None:
+                error_response = HarnessProtocolError(
+                    data={"message": message, "code": code}
+                )
+                await websocket.send_text(error_response.model_dump_json())
+
+            async def send_harness_error(message: str) -> None:
+                """Emit a terminal ERROR event in the harness event stream."""
+                error_event = HarnessEvent(
+                    type=HarnessEventType.ERROR,
+                    data={"message": message, "recoverable": False},
+                )
+                await websocket.send_text(error_event.model_dump_json())
+
+            try:
+                session_id, session_env = await self._create_session()
+                # Protect the live harness from idle reaping and HTTP session
+                # close, including startup and turns that emit no events.
+                self._session_websocket_attachments.add(session_id)
+
+                async with AsyncExitStack() as stack:
+                    mcp_session_factory = getattr(session_env, "mcp_session", None)
+                    if callable(mcp_session_factory):
+                        mcp_session_cm = cast(
+                            AsyncContextManager[Any], mcp_session_factory()
+                        )
+                        await stack.enter_async_context(mcp_session_cm)
+
+                    # Starts the harness process and injects environment tools
+                    await session_env.reset_async()
+                    await websocket.send_text(
+                        json.dumps(
+                            {
+                                "type": "session_started",
+                                "data": {
+                                    "session_id": session_id,
+                                    "harness": session_env.adapter.config.name,
+                                },
+                            }
+                        )
+                    )
+
+                    while True:
+                        raw_message = await websocket.receive_text()
+
+                        try:
+                            message_dict = json.loads(raw_message)
+                        except json.JSONDecodeError as e:
+                            await send_protocol_error(
+                                f"Invalid JSON: {e}", WSErrorCode.INVALID_JSON
+                            )
+                            continue
+                        try:
+                            client_message = HarnessClientMessage(**message_dict)
+                        except (ValidationError, TypeError) as e:
+                            await send_protocol_error(
+                                f"Invalid message: {e}",
+                                WSErrorCode.VALIDATION_ERROR,
+                            )
+                            continue
+
+                        self._update_session_activity(session_id, increment_step=True)
+
+                        async def stream_turn(content: str) -> bool:
+                            """Stream one turn; True if it ended with TURN_COMPLETE."""
+                            saw_terminal = False
+                            adapter = session_env.adapter
+                            if not await adapter.is_alive():
+                                raise HarnessNotRunningError(
+                                    "harness process is not running"
+                                )
+                            async for event in adapter.send_message_streaming(content):
+                                await websocket.send_text(event.model_dump_json())
+                                # Record progress throughout the turn.
+                                self._update_session_activity(session_id)
+                                saw_terminal = (
+                                    event.type is HarnessEventType.TURN_COMPLETE
+                                )
+                            return saw_terminal
+
+                        # Bound the turn in wall-clock time, matching what
+                        # simulation mode does in HarnessEnvironment._run_turn.
+                        # Without this a hung harness holds the session open
+                        # forever, and the server sits at capacity.
+                        turn_timeout_s = session_env.adapter.config.session_timeout_s
+                        try:
+                            completed = await asyncio.wait_for(
+                                stream_turn(client_message.content),
+                                turn_timeout_s,
+                            )
+                        except asyncio.TimeoutError:
+                            await send_harness_error(
+                                f"harness turn exceeded {turn_timeout_s} seconds"
+                            )
+                            break
+                        except HarnessNotRunningError:
+                            await send_harness_error("harness process is not running")
+                            break
+                        except Exception:
+                            # Harness state after a crash is undefined; end
+                            # the session so a reconnect gets a fresh one.
+                            # Adapter exceptions can contain credentials or
+                            # subprocess output; do not expose them to clients.
+                            await send_harness_error("harness turn failed")
+                            break
+
+                        if not completed:
+                            # send_message() raises HarnessError here; the
+                            # socket equivalent is to say so and end the
+                            # session, rather than leaving a client that
+                            # blocks on the terminal event waiting forever.
+                            await send_harness_error(
+                                "harness event stream ended without a "
+                                "TURN_COMPLETE event"
+                            )
+                            break
+
+            except WebSocketDisconnect:
+                pass
+            except SessionCapacityError as e:
+                await send_error(str(e), WSErrorCode.CAPACITY_REACHED)
+            except EnvironmentFactoryError as e:
+                await send_error(str(e), WSErrorCode.FACTORY_ERROR)
+            except Exception:
+                try:
+                    await send_error(
+                        "harness session failed", WSErrorCode.SESSION_ERROR
+                    )
+                except (RuntimeError, WebSocketDisconnect):
+                    pass
+            finally:
+                if session_id:
+                    # Release ownership without an await so cancellation cannot
+                    # leave a session permanently exempt from idle reaping.
+                    self._session_websocket_attachments.discard(session_id)
+                    cleanup = asyncio.create_task(self._destroy_session(session_id))
+                    try:
+                        await asyncio.shield(cleanup)
+                    except asyncio.CancelledError:
+                        # ASGI cancellation must not orphan a running harness.
+                        # Finish teardown before propagating cancellation, even
+                        # when the request's cancel scope cancels us repeatedly.
+                        while not cleanup.done():
+                            try:
+                                await asyncio.shield(cleanup)
+                            except asyncio.CancelledError:
+                                pass
+                        cleanup.result()
+                        raise
+                try:
+                    await websocket.close()
+                except (RuntimeError, WebSocketDisconnect):
+                    # TestClient raises RuntimeError, real ASGI servers raise
+                    # WebSocketDisconnect when the client is already gone.
+                    pass
+
     def register_routes(
         self, app: FastAPI, mode: ServerMode | str = ServerMode.SIMULATION
     ) -> None:
@@ -541,14 +859,15 @@ class HTTPEnvServer:
         Register HTTP routes on a FastAPI application.
 
         Args:
-            app: FastAPI application instance
-            mode: Server mode - either SIMULATION or PRODUCTION (or string equivalents).
-                  In production mode, simulation control endpoints (/reset, /step, /state)
-                  are NOT registered. Only safe endpoints (/health, /schema, /metadata, /ws)
-                  are available. Defaults to SIMULATION for backwards compatibility.
+            app (`FastAPI`):
+                FastAPI application instance.
+            mode (`ServerMode` or `str`, *optional*, defaults to `ServerMode.SIMULATION`):
+                Server mode. In production mode, simulation control endpoints (/reset, /step,
+                /state) are NOT registered. Only safe endpoints (/health, /schema, /metadata,
+                /ws) are available.
 
         Raises:
-            ValueError: If mode is not a valid ServerMode or string equivalent.
+            `ValueError`: If `mode` is not a valid `ServerMode` or string equivalent.
         """
         # Convert string to ServerMode enum for backwards compatibility
         if isinstance(mode, str):
@@ -568,6 +887,8 @@ class HTTPEnvServer:
 
         async def _stop_session_reaper() -> None:
             server_ref._stop_reaper()
+            if server_ref._shared_session_executor is not None:
+                server_ref._shared_session_executor.shutdown(wait=True)
 
         if not getattr(app.router, "_openenv_reaper_registered", False):
             app.router.on_startup.append(_start_session_reaper)
@@ -584,7 +905,7 @@ class HTTPEnvServer:
             try:
                 kwargs = request.model_dump(exclude_unset=True)
 
-                is_async = _env.reset_async.__func__ is not Environment.reset_async
+                is_async = overrides_method(_env.reset_async, Environment.reset_async)
 
                 if is_async:
                     sig = inspect.signature(_env.reset_async)
@@ -619,7 +940,7 @@ class HTTPEnvServer:
             try:
                 kwargs = request.model_dump(exclude_unset=True, exclude={"action"})
 
-                is_async = _env.step_async.__func__ is not Environment.step_async
+                is_async = overrides_method(_env.step_async, Environment.step_async)
 
                 if is_async:
                     sig = inspect.signature(_env.step_async)
@@ -650,9 +971,8 @@ class HTTPEnvServer:
             Handle MCP JSON-RPC requests.
 
             Supports tools/list and tools/call methods in JSON-RPC 2.0 format,
-            plus OpenEnv session lifecycle methods for HTTP MCP:
-            - openenv/session/create
-            - openenv/session/close
+            plus OpenEnv session lifecycle methods for HTTP MCP (openenv/session/create,
+            openenv/session/close).
             """
             method = request.method
             request_id = request.id
@@ -713,15 +1033,33 @@ class HTTPEnvServer:
                     )
 
                 async with self._session_lock:
-                    env = self._sessions.pop(target_session_id, _MISSING)
-                    if env is not _MISSING:
+                    if target_session_id in self._session_websocket_attachments:
+                        env = _MISSING
+                        attached = True
+                        self._session_pending_closes.add(target_session_id)
+                        executor = None
+                        stack = None
+                    else:
+                        attached = False
+                        env = self._sessions.pop(target_session_id, _MISSING)
+                    if not attached and env is not _MISSING:
                         executor = self._session_executors.pop(target_session_id, None)
                         stack = self._session_stacks.pop(target_session_id, None)
                         self._session_info.pop(target_session_id, None)
-                    else:
+                    elif not attached:
                         executor = None
                         stack = None
 
+                if attached:
+                    return JsonRpcResponse.success(
+                        result={
+                            "session_id": target_session_id,
+                            "closed": False,
+                            "closing": True,
+                        },
+                        request_id=request_id,
+                    )
+
                 if env is _MISSING:
                     return JsonRpcResponse.error_response(
                         JsonRpcErrorCode.INVALID_PARAMS,
@@ -786,9 +1124,45 @@ class HTTPEnvServer:
                 mcp_server = getattr(_env, "mcp_server", None)
                 mcp_session_factory = getattr(_env, "mcp_session", None)
 
+                async def call_mcp_style_step(action: Action) -> Observation:
+                    is_async = overrides_method(_env.step_async, Environment.step_async)
+                    if is_async:
+                        return await _env.step_async(action)
+                    if managed_session_id:
+                        return await self._run_in_session_executor(
+                            managed_session_id,
+                            _env.step,
+                            action,
+                        )
+                    return await self._run_sync_in_thread_pool(_env.step, action)
+
+                supports_mcp_style_actions = self.action_cls in {
+                    CallToolAction,
+                    ListToolsAction,
+                }
+
                 if method == McpMethod.TOOLS_LIST:
                     # Check if environment is MCP-enabled
                     if mcp_client is None and mcp_server is None:
+                        if supports_mcp_style_actions:
+                            observation = await call_mcp_style_step(ListToolsAction())
+                            if isinstance(observation, ListToolsObservation):
+                                return JsonRpcResponse.success(
+                                    result={
+                                        "tools": [
+                                            tool.model_dump()
+                                            for tool in observation.tools
+                                        ]
+                                    },
+                                    request_id=request_id,
+                                )
+                            return JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.INTERNAL_ERROR,
+                                "MCP-style tools/list step returned "
+                                f"{type(observation).__name__}, expected "
+                                "ListToolsObservation",
+                                request_id=request_id,
+                            )
                         return JsonRpcResponse.error_response(
                             JsonRpcErrorCode.INTERNAL_ERROR,
                             "Environment does not support MCP",
@@ -796,21 +1170,12 @@ class HTTPEnvServer:
                         )
 
                     if mcp_client:
-                        if managed_session_id and mcp_client.is_connected():
-                            # Session-managed with live transport — call
-                            # directly, no redundant re-entry.
-                            tools = await mcp_client.list_tools()
-                        elif callable(mcp_session_factory):
-                            # Stateless request, or session-managed but the
-                            # background transport was lost: (re-)open.
-                            mcp_session_cm = cast(
-                                AsyncContextManager[Any], mcp_session_factory()
-                            )
-                            async with mcp_session_cm:
-                                tools = await mcp_client.list_tools()
-                        else:
-                            async with mcp_client:
-                                tools = await mcp_client.list_tools()
+                        tools = await self._run_mcp_client_operation(
+                            mcp_client,
+                            mcp_session_factory,
+                            managed_session_id,
+                            mcp_client.list_tools,
+                        )
 
                         return JsonRpcResponse.success(
                             result={
@@ -849,41 +1214,48 @@ class HTTPEnvServer:
                     tool_name = params.get("name")
                     arguments = params.get("arguments", {})
 
-                    if mcp_client is None and mcp_server is None:
+                    if not tool_name:
                         return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INTERNAL_ERROR,
-                            "Environment does not support MCP",
+                            JsonRpcErrorCode.INVALID_PARAMS,
+                            "Missing 'name' in params",
                             request_id=request_id,
                         )
 
-                    if not tool_name:
+                    if mcp_client is None and mcp_server is None:
+                        if supports_mcp_style_actions:
+                            observation = await call_mcp_style_step(
+                                CallToolAction(
+                                    tool_name=tool_name,
+                                    arguments=arguments,
+                                )
+                            )
+                            if isinstance(observation, CallToolObservation):
+                                return JsonRpcResponse.success(
+                                    result=_make_json_serializable(observation),
+                                    request_id=request_id,
+                                )
+                            return JsonRpcResponse.error_response(
+                                JsonRpcErrorCode.INTERNAL_ERROR,
+                                "MCP-style tools/call step returned "
+                                f"{type(observation).__name__}, expected "
+                                "CallToolObservation",
+                                request_id=request_id,
+                            )
                         return JsonRpcResponse.error_response(
-                            JsonRpcErrorCode.INVALID_PARAMS,
-                            "Missing 'name' in params",
+                            JsonRpcErrorCode.INTERNAL_ERROR,
+                            "Environment does not support MCP",
                             request_id=request_id,
                         )
 
                     if mcp_client:
-                        if managed_session_id and mcp_client.is_connected():
-                            # Session-managed with live transport.
-                            result = await mcp_client.call_tool(
+                        result = await self._run_mcp_client_operation(
+                            mcp_client,
+                            mcp_session_factory,
+                            managed_session_id,
+                            lambda: mcp_client.call_tool(
                                 name=tool_name, arguments=arguments
-                            )
-                        elif callable(mcp_session_factory):
-                            # Stateless request, or session-managed but the
-                            # background transport was lost: (re-)open.
-                            mcp_session_cm = cast(
-                                AsyncContextManager[Any], mcp_session_factory()
-                            )
-                            async with mcp_session_cm:
-                                result = await mcp_client.call_tool(
-                                    name=tool_name, arguments=arguments
-                                )
-                        else:
-                            async with mcp_client:
-                                result = await mcp_client.call_tool(
-                                    name=tool_name, arguments=arguments
-                                )
+                            ),
+                        )
                     elif mcp_server:
                         server_tools = get_server_tools(mcp_server)
                         if tool_name in server_tools:
@@ -935,6 +1307,96 @@ class HTTPEnvServer:
                 if should_close:
                     _env.close()
 
+        def _check_env_name(env_name: str) -> None:
+            if env_name.lower() != self.env_name.lower():
+                raise HTTPException(
+                    status_code=status.HTTP_404_NOT_FOUND,
+                    detail=f"unknown environment {env_name!r}",
+                )
+
+        def _normalize_split(split: Any) -> Dict[str, Any]:
+            if hasattr(split, "model_dump"):
+                return cast(Dict[str, Any], split.model_dump())
+            if isinstance(split, dict):
+                return _make_json_serializable(split)
+            if isinstance(split, str):
+                split_type = (
+                    split if split in {"train", "validation", "test"} else "validation"
+                )
+                return {"name": split, "type": split_type}
+            return {"name": str(split), "type": "validation"}
+
+        async def _call_task_method(method_name: str, *args: Any) -> Any:
+            _env = self._env_factory()
+            try:
+                method = getattr(_env, method_name, None)
+                if not callable(method):
+                    raise NotImplementedError
+                return await _maybe_await(method(*args))
+            except NotImplementedError as e:
+                raise HTTPException(
+                    status_code=status.HTTP_501_NOT_IMPLEMENTED,
+                    detail=f"{method_name} is not supported for this environment",
+                ) from e
+            except IndexError as e:
+                raise HTTPException(
+                    status_code=status.HTTP_400_BAD_REQUEST,
+                    detail="Invalid task index",
+                ) from e
+            finally:
+                _env.close()
+
+        @app.get("/list_environments", tags=["Task API"])
+        async def list_environments() -> list[str]:
+            return [self.env_name]
+
+        @app.get("/{env_name}/splits", tags=["Task API"])
+        async def list_splits(env_name: str) -> list[Dict[str, Any]]:
+            _check_env_name(env_name)
+            splits = await _call_task_method("list_splits")
+            return [_normalize_split(split) for split in splits]
+
+        @app.post("/{env_name}/tasks", tags=["Task API"])
+        async def list_tasks(
+            env_name: str,
+            request: ListTasksRequest,
+        ) -> Dict[str, Any]:
+            _check_env_name(env_name)
+            tasks = await _call_task_method("list_tasks", request.split)
+            return {
+                "tasks": _make_json_serializable(tasks),
+                "env_name": self.env_name,
+            }
+
+        @app.post("/{env_name}/num_tasks", tags=["Task API"])
+        async def num_tasks(
+            env_name: str,
+            request: NumTasksRequest,
+        ) -> Dict[str, Any]:
+            _check_env_name(env_name)
+            count = await _call_task_method("num_tasks", request.split)
+            return {"num_tasks": count}
+
+        @app.post("/{env_name}/task", tags=["Task API"])
+        async def get_task(
+            env_name: str,
+            request: GetTaskRequest,
+        ) -> Dict[str, Any]:
+            _check_env_name(env_name)
+            task = await _call_task_method("get_task", request.split, request.index)
+            return {"task": _make_json_serializable(task)}
+
+        @app.post("/{env_name}/task_range", tags=["Task API"])
+        async def get_task_range(
+            env_name: str,
+            request: GetTaskRangeRequest,
+        ) -> Dict[str, Any]:
+            _check_env_name(env_name)
+            tasks = await _call_task_method(
+                "get_task_range", request.split, request.start, request.stop
+            )
+            return {"tasks": _make_json_serializable(tasks)}
+
         # Register MCP WebSocket endpoint (available in both production and simulation modes)
         @app.websocket("/mcp")
         async def mcp_websocket_endpoint(websocket: WebSocket):
@@ -942,10 +1404,8 @@ class HTTPEnvServer:
             WebSocket endpoint for MCP JSON-RPC requests.
 
             Each WebSocket connection gets its own environment instance for MCP operations.
-
-            Message Protocol:
-            - Client sends: JSON-RPC 2.0 request (tools/list, tools/call)
-            - Server responds: JSON-RPC 2.0 response (result or error)
+            The client sends JSON-RPC 2.0 requests (tools/list, tools/call) and the server
+            responds with JSON-RPC 2.0 responses (result or error).
             """
             await websocket.accept()
 
@@ -1039,9 +1499,14 @@ class HTTPEnvServer:
                     await self._destroy_session(session_id)
                 try:
                     await websocket.close()
-                except RuntimeError:
+                except (RuntimeError, WebSocketDisconnect):
                     pass
 
+        # In production mode, a harness environment is exposed directly to
+        # clients via a streaming WebSocket (RFC 005).
+        if mode == ServerMode.PRODUCTION and self._factory_produces_harness_env():
+            self._register_harness_route(app)
+
         # Register simulation control routes only in simulation mode
         if mode == ServerMode.SIMULATION:
 
@@ -1176,7 +1641,7 @@ version, author, and documentation links.
                 GetEndpointConfig(
                     path="/state",
                     handler=get_state_handler,
-                    response_model=State,
+                    response_model=self.state_cls,
                     tag="State Management",
                     summary="Get current environment state",
                     description="""
@@ -1200,7 +1665,8 @@ Get JSON schemas for actions, observations, and state in a single response.
 
 Returns a combined schema object containing:
 - **action**: JSON schema for actions accepted by this environment
-- **observation**: JSON schema for observations returned by this environment
+- **observation**: JSON schema for observations returned by step
+- **reset_observation**: JSON schema for observations returned by reset
 - **state**: JSON schema for environment state objects
 
 This is more efficient than calling individual schema endpoints and provides
@@ -1235,14 +1701,14 @@ all schema information needed to interact with the environment.
             return SchemaResponse(
                 action=self.action_cls.model_json_schema(),
                 observation=self.observation_cls.model_json_schema(),
-                state=State.model_json_schema(),
+                reset_observation=self.reset_observation_cls.model_json_schema(),
+                state=self.state_cls.model_json_schema(),
             )
 
         # Register MCP endpoint for production mode (direct MCP access)
         @app.post("/mcp")
         async def mcp_endpoint(request_raw: Request) -> Dict[str, Any]:
-            """
-            MCP JSON-RPC endpoint for production mode.
+            """MCP JSON-RPC endpoint for production mode.
 
             Bypasses step() overhead and provides direct access to MCP tools.
             Supports tools/list and tools/call methods.
@@ -1275,20 +1741,48 @@ all schema information needed to interact with the environment.
             """
             WebSocket endpoint for persistent environment sessions.
 
-            Each WebSocket connection gets its own environment instance.
-
-            Message Protocol:
-            - Client sends: WSResetMessage | WSStepMessage | WSStateMessage | WSCloseMessage
-            - Server responds: WSObservationResponse | WSStateResponse | WSErrorResponse
+            By default, each WebSocket connection gets its own environment instance.
+            A client can instead attach to an existing HTTP MCP session by passing its
+            session ID in the query string. The client sends WSResetMessage,
+            WSStepMessage, WSStateMessage, or WSCloseMessage; the server responds with
+            WSObservationResponse, WSStateResponse, or WSErrorResponse.
             """
             await websocket.accept()
 
             session_id = None
             session_env = None
+            owns_session = False
+            attached_session = False
 
             try:
-                # Create session with dedicated environment
-                session_id, session_env = await self._create_session()
+                requested_session_id = websocket.query_params.get("session_id")
+                if requested_session_id:
+                    async with self._session_lock:
+                        attached_env = self._sessions.get(
+                            requested_session_id, _MISSING
+                        )
+                        if attached_env is _MISSING:
+                            raise RuntimeError(
+                                f"Unknown session_id: {requested_session_id}"
+                            )
+                        if attached_env is None:
+                            raise RuntimeError(
+                                f"Session {requested_session_id} is still initializing"
+                            )
+                        if requested_session_id in self._session_websocket_attachments:
+                            raise RuntimeError(
+                                f"Session {requested_session_id} already has "
+                                "an attached WebSocket"
+                            )
+                        self._session_websocket_attachments.add(requested_session_id)
+                        session_id = requested_session_id
+                        session_env = attached_env
+                        attached_session = True
+                    self._update_session_activity(session_id)
+                else:
+                    session_id, session_env = await self._create_session()
+                    owns_session = True
+
                 if session_env is None:
                     raise RuntimeError(
                         "Session environment not initialized for websocket"
@@ -1299,7 +1793,7 @@ all schema information needed to interact with the environment.
 
                 async with AsyncExitStack() as stack:
                     mcp_session_factory = getattr(session_env, "mcp_session", None)
-                    if callable(mcp_session_factory):
+                    if owns_session and callable(mcp_session_factory):
                         mcp_session_cm = cast(
                             AsyncContextManager[Any], mcp_session_factory()
                         )
@@ -1328,9 +1822,9 @@ all schema information needed to interact with the environment.
                                 case "reset":
                                     msg = WSResetMessage(**message_dict)
 
-                                    is_async = (
-                                        session_env.reset_async.__func__
-                                        is not Environment.reset_async
+                                    is_async = overrides_method(
+                                        session_env.reset_async,
+                                        Environment.reset_async,
                                     )
 
                                     if is_async:
@@ -1366,9 +1860,9 @@ all schema information needed to interact with the environment.
                                         msg.data, self.action_cls
                                     )
 
-                                    is_async = (
-                                        session_env.step_async.__func__
-                                        is not Environment.step_async
+                                    is_async = overrides_method(
+                                        session_env.step_async,
+                                        Environment.step_async,
                                     )
 
                                     if is_async:
@@ -1478,11 +1972,18 @@ all schema information needed to interact with the environment.
                 )
                 await websocket.send_text(error_resp.model_dump_json())
             finally:
-                if session_id:
+                if attached_session and session_id:
+                    # Do not await the lock before releasing ownership: a task
+                    # cancelled while waiting would leave this session
+                    # permanently exempt from idle reaping.
+                    self._session_websocket_attachments.discard(session_id)
+                    if session_id in self._session_pending_closes:
+                        await asyncio.shield(self._destroy_session(session_id))
+                elif owns_session and session_id:
                     await self._destroy_session(session_id)
                 try:
                     await websocket.close()
-                except RuntimeError:
+                except (RuntimeError, WebSocketDisconnect):
                     pass
 
 
@@ -1494,6 +1995,14 @@ def create_app(
     max_concurrent_envs: Optional[int] = None,
     concurrency_config: Optional[ConcurrencyConfig] = None,
     gradio_builder: Optional[Callable[..., Any]] = None,
+    custom_tab_name: str = "Custom",
+    custom_tab_primary: bool = False,
+    show_default_tab: bool = True,
+    title_override: Optional[str] = None,
+    state_cls: Type[State] = State,
+    *,
+    mode: Optional[ServerMode | str] = None,
+    reset_observation_cls: Optional[Type[Observation]] = None,
 ) -> FastAPI:
     """
     Create a FastAPI application with or without web interface.
@@ -1502,21 +2011,48 @@ def create_app(
     including README integration for better user experience.
 
     Args:
-        env: Environment factory (callable) that creates new instances
-        action_cls: The Action subclass this environment expects
-        observation_cls: The Observation subclass this environment returns
-        env_name: Optional environment name for README loading
-        max_concurrent_envs: Maximum concurrent WebSocket sessions.
-                             Mutually exclusive with concurrency_config.
-        concurrency_config: Optional ConcurrencyConfig for advanced concurrency settings.
-                            Mutually exclusive with max_concurrent_envs.
-        gradio_builder: Optional callable to build a custom Gradio UI at /web.
-            Signature: (web_manager, action_fields, metadata, is_chat_env, title,
-            quick_start_md) -> gr.Blocks. When None, the default Gradio app is used.
-            See docs/customizing-web-ui.md.
+        env (`Callable[[], Environment]`):
+            Environment factory (callable) that creates new instances.
+        action_cls (`Type[Action]`):
+            The Action subclass this environment expects.
+        observation_cls (`Type[Observation]`):
+            The Observation subclass returned by step.
+        reset_observation_cls (`Type[Observation]`, *optional*):
+            The reset observation model published in `/schema`. Defaults to
+            `observation_cls`; declare a distinct model when reset differs.
+        env_name (`str`, *optional*):
+            Environment name for README loading.
+        max_concurrent_envs (`int`, *optional*):
+            Maximum concurrent WebSocket sessions. Mutually exclusive with
+            `concurrency_config`.
+        concurrency_config (`ConcurrencyConfig`, *optional*):
+            Advanced concurrency settings. Mutually exclusive with
+            `max_concurrent_envs`.
+        gradio_builder (`Callable`, *optional*):
+            Callable to build a custom Gradio UI at /web. Signature:
+            `(web_manager, action_fields, metadata, is_chat_env, title, quick_start_md) -> gr.Blocks`.
+            When `None`, the default Gradio app is used.
+        custom_tab_name (`str`, *optional*, defaults to `"Custom"`):
+            Label for the env-specific tab when `gradio_builder` is provided.
+        custom_tab_primary (`bool`, *optional*, defaults to `False`):
+            When `True`, the env-specific tab is active first; the auto-generated
+            Playground becomes secondary.
+        show_default_tab (`bool`, *optional*, defaults to `True`):
+            When `False`, mount the env's `gradio_builder` output alone (no
+            auto-generated Playground, no tab chrome). Only meaningful when
+            `gradio_builder` is provided.
+        title_override (`str`, *optional*):
+            If set, used as the Gradio app title instead of the default
+            `"OpenEnv Agentic Environment: {name}"`.
+        state_cls (`Type[State]`, *optional*, defaults to `State`):
+            The `State` subclass this environment reports, used for the `/state`
+            response model and the `state` entry of `/schema`.
+        mode (`ServerMode` or `str`, *optional*):
+            Server mode. When `None`, resolved from the `OPENENV_MODE`
+            environment variable, defaulting to simulation.
 
     Returns:
-        FastAPI application instance with or without web interface and README integration
+        `FastAPI` application instance with or without web interface and README integration.
     """
     # Check if web interface should be enabled
     # This can be controlled via environment variable or build argument
@@ -1537,12 +2073,27 @@ def create_app(
             env_name,
             max_concurrent_envs,
             concurrency_config,
+            state_cls=state_cls,
+            reset_observation_cls=reset_observation_cls,
             gradio_builder=gradio_builder,
+            custom_tab_name=custom_tab_name,
+            custom_tab_primary=custom_tab_primary,
+            show_default_tab=show_default_tab,
+            title_override=title_override,
+            mode=mode,
         )
     else:
         # Use standard FastAPI app without web interface
         return create_fastapi_app(
-            env, action_cls, observation_cls, max_concurrent_envs, concurrency_config
+            env,
+            action_cls,
+            observation_cls,
+            max_concurrent_envs,
+            concurrency_config,
+            env_name=env_name,
+            state_cls=state_cls,
+            reset_observation_cls=reset_observation_cls,
+            mode=mode,
         )
 
 
@@ -1552,21 +2103,42 @@ def create_fastapi_app(
     observation_cls: Type[Observation],
     max_concurrent_envs: Optional[int] = None,
     concurrency_config: Optional[ConcurrencyConfig] = None,
+    env_name: Optional[str] = None,
+    state_cls: Type[State] = State,
+    *,
+    mode: Optional[ServerMode | str] = None,
+    reset_observation_cls: Optional[Type[Observation]] = None,
 ) -> FastAPI:
     """
     Create a FastAPI application with comprehensive documentation.
 
     Args:
-        env: Environment factory (callable) that creates new instances
-        action_cls: The Action subclass this environment expects
-        observation_cls: The Observation subclass this environment returns
-        max_concurrent_envs: Maximum concurrent WebSocket sessions.
-                             Mutually exclusive with concurrency_config.
-        concurrency_config: Optional ConcurrencyConfig for advanced concurrency settings.
-                            Mutually exclusive with max_concurrent_envs.
+        env (`Callable[[], Environment]`):
+            Environment factory (callable) that creates new instances.
+        action_cls (`Type[Action]`):
+            The Action subclass this environment expects.
+        observation_cls (`Type[Observation]`):
+            The Observation subclass returned by step.
+        reset_observation_cls (`Type[Observation]`, *optional*):
+            The reset observation model published in `/schema`. Defaults to
+            `observation_cls`; declare a distinct model when reset differs.
+        max_concurrent_envs (`int`, *optional*):
+            Maximum concurrent WebSocket sessions. Mutually exclusive with
+            `concurrency_config`.
+        concurrency_config (`ConcurrencyConfig`, *optional*):
+            Advanced concurrency settings. Mutually exclusive with
+            `max_concurrent_envs`.
+        env_name (`str`, *optional*):
+            Optional environment name for task/split endpoints.
+        state_cls (`Type[State]`, *optional*, defaults to `State`):
+            The `State` subclass this environment reports, used for the `/state`
+            response model and the `state` entry of `/schema`.
+        mode (`ServerMode` or `str`, *optional*):
+            Server mode. When `None`, resolved from the `OPENENV_MODE`
+            environment variable, defaulting to simulation.
 
     Returns:
-        FastAPI application instance
+        `FastAPI` application instance.
     """
     try:
         from fastapi import FastAPI
@@ -1627,11 +2199,11 @@ HTTP API for interacting with OpenEnv environments through a standardized interf
         openapi_url="/openapi.json",
         contact={
             "name": "OpenEnv Team",
-            "url": "https://github.com/meta-pytorch/OpenEnv",
+            "url": "https://github.com/huggingface/OpenEnv",
         },
         license_info={
             "name": "BSD-3-Clause",
-            "url": "https://github.com/meta-pytorch/OpenEnv/blob/main/LICENSE",
+            "url": "https://github.com/huggingface/OpenEnv/blob/main/LICENSE",
         },
     )
 
@@ -1641,6 +2213,11 @@ HTTP API for interacting with OpenEnv environments through a standardized interf
         observation_cls,
         max_concurrent_envs,
         concurrency_config=concurrency_config,
+        env_name=env_name,
+        state_cls=state_cls,
+        reset_observation_cls=reset_observation_cls,
     )
-    server.register_routes(app)
+    if mode is None:
+        mode = os.environ.get("OPENENV_MODE", ServerMode.SIMULATION.value)
+    server.register_routes(app, mode=mode)
     return app

src/core/env_server/interfaces.py

@@ -1,12 +1,18 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 import inspect
 from abc import ABC, abstractmethod
-from typing import Any, Generic, Optional, Protocol, TYPE_CHECKING, TypeVar
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    List,
+    Optional,
+    Protocol,
+    Tuple,
+    TYPE_CHECKING,
+    TypeVar,
+)
 
 from typing_extensions import TypedDict
 
@@ -48,13 +54,17 @@ class ModelTokenizer(Protocol):
         """Apply a chat template to format and optionally tokenize a conversation.
 
         Args:
-            conversation: List of message dictionaries with 'role' and 'content'
-            tokenize: Whether to tokenize the output
-            return_tensors: Format for returned tensors ('pt' for PyTorch)
-            **kwargs: Additional arguments
+            conversation (`list[Message]`):
+                List of message dictionaries with 'role' and 'content'.
+            tokenize (`bool`, *optional*, defaults to `True`):
+                Whether to tokenize the output.
+            return_tensors (`str`, *optional*):
+                Format for returned tensors ('pt' for PyTorch).
+            **kwargs:
+                Additional arguments.
 
         Returns:
-            Formatted and optionally tokenized conversation
+            Formatted and optionally tokenized conversation.
         """
         ...
 
@@ -64,12 +74,126 @@ class ModelTokenizer(Protocol):
         """Decode token IDs back to text.
 
         Args:
-            token_ids: Token IDs to decode
-            skip_special_tokens: Whether to skip special tokens in output
-            **kwargs: Additional arguments
+            token_ids (`Any`):
+                Token IDs to decode.
+            skip_special_tokens (`bool`, *optional*, defaults to `False`):
+                Whether to skip special tokens in output.
+            **kwargs:
+                Additional arguments.
 
         Returns:
-            Decoded text string
+            `str`: Decoded text string.
+        """
+        ...
+
+
+class TaskProvider(Protocol):
+    """
+    Optional task discovery API for dataset-backed environments.
+
+    An environment implements this protocol structurally — declare the methods on
+    an [`~openenv.core.env_server.interfaces.Environment`] subclass, without
+    inheriting from `TaskProvider`. When the methods are present,
+    [`~openenv.core.env_server.http_server.HTTPEnvServer`] exposes them as HTTP
+    routes under `/{env_name}/…`; when they are absent, those routes return
+    `501 Not Implemented`. Each method may be sync or async.
+
+    Task provider methods are for metadata/discovery only and should be
+    side-effect-free. They must be callable on a freshly constructed
+    environment instance because HTTP compatibility routes may create a
+    short-lived instance solely for task discovery.
+
+    Selecting a task is not part of this protocol — pass the chosen split and
+    index to `reset()` instead. See the
+    [Task API guide](https://huggingface.co/docs/openenv/guides/task-api).
+
+    Examples:
+
+    ```python
+    env.list_splits()          # ["train", "test"]
+    env.num_tasks("test")      # 7595
+    env.get_task("test", 12)   # {"id": "test-12", "index": 12}
+    env.reset(split="test", index=12)
+    ```
+    """
+
+    def list_splits(self) -> list[Any]:
+        """
+        Return task split descriptors supported by this environment.
+
+        Returns:
+            `list[Any]`: Split descriptors. Plain strings, dicts, and Pydantic
+                models are all accepted; the server normalizes each entry to
+                `{"name": ..., "type": ...}`.
+        """
+        ...
+
+    def list_tasks(self, split: str) -> list[Any]:
+        """
+        Return all task specs for a split.
+
+        Args:
+            split (`str`):
+                Task split name.
+
+        Returns:
+            `list[Any]`: Task specs for the split. Environments backed by very
+                large or streamed splits may return a bounded preview, but
+                `num_tasks` should still report the true total.
+        """
+        ...
+
+    def num_tasks(self, split: str) -> int:
+        """
+        Return the number of task specs in a split.
+
+        Args:
+            split (`str`):
+                Task split name.
+
+        Returns:
+            `int`: Number of task specs available in the split.
+        """
+        ...
+
+    def get_task(self, split: str, index: int) -> Any:
+        """
+        Return one task spec by split and index.
+
+        Args:
+            split (`str`):
+                Task split name.
+            index (`int`):
+                Task index within the split.
+
+        Returns:
+            `Any`: The task spec at that position.
+
+        Raises:
+            `IndexError`: If `index` is out of range for the split. The HTTP
+                route converts this to a `400 Bad Request`.
+        """
+        ...
+
+    def get_task_range(
+        self,
+        split: str,
+        start: Optional[int] = None,
+        stop: Optional[int] = None,
+    ) -> list[Any]:
+        """
+        Return task specs for Python slice-style range bounds.
+
+        Args:
+            split (`str`):
+                Task split name.
+            start (`int`, *optional*):
+                Inclusive start index. Defaults to the beginning of the split.
+            stop (`int`, *optional*):
+                Exclusive stop index. Defaults to the end of the split.
+
+        Returns:
+            `list[Any]`: Task specs in `[start, stop)`.
         """
         ...
 
@@ -87,10 +211,11 @@ class Transform(ABC, Generic[ObsT]):
         """Transform an observation.
 
         Args:
-            observation: The input observation
+            observation (`ObsT`):
+                The input observation.
 
         Returns:
-            The transformed observation
+            `ObsT`: The transformed observation.
         """
         pass
 
@@ -98,35 +223,41 @@ class Transform(ABC, Generic[ObsT]):
 class Environment(ABC, Generic[ActT, ObsT, StateT]):
     """Base class for all environment servers following Gym/Gymnasium API.
 
+    See [rfcs/004-rubrics.md](https://github.com/huggingface/OpenEnv/blob/main/rfcs/004-rubrics.md) for rubric design details.
+
     Args:
-        transform: Optional transform to apply to observations
-        rubric: Optional rubric for reward computation. When provided, the
+        transform (`Transform`, *optional*):
+            Optional transform to apply to observations.
+        rubric (`Rubric`, *optional*):
+            Optional rubric for reward computation. When provided, the
             rubric's output can be used to set the observation's reward in step().
 
-    Class Attributes:
-        SUPPORTS_CONCURRENT_SESSIONS: Whether this environment supports concurrent sessions.
-            When True, multiple WebSocket connections can each have their own
-            environment instance (up to max_concurrent_envs). When False (default),
-            the environment should only be used with a single session at a time.
-
-            Set this to True in your Environment subclass if:
-            - The environment uses proper session isolation (e.g., unique working dirs)
-            - No shared mutable state exists between instances
-            - External resources (databases, APIs) can handle concurrent access
-
     Attributes:
-        rubric: Optional rubric for computing rewards. Environments can set this
-            in __init__ and use it in step() to compute observation rewards.
-            Training infrastructure can access it for introspection:
-                for name, r in env.rubric.named_rubrics():
-                    print(f"{name}: {r.last_score}")
+        SUPPORTS_CONCURRENT_SESSIONS (`bool`):
+            Whether this environment supports concurrent sessions. When ``True``,
+            multiple WebSocket connections can each have their own environment
+            instance (up to ``max_concurrent_envs``). When ``False`` (default),
+            the environment should only be used with a single session at a time.
 
-    See RFC 004 for rubric design: rfcs/004-rubrics.md
+            Set this to ``True`` in your subclass if the environment uses proper
+            session isolation (unique working dirs, no shared mutable state, and
+            external resources that can handle concurrent access).
+        rubric (`Rubric`, *optional*):
+            Optional rubric for computing rewards. Set in ``__init__`` and use in
+            ``step()`` to compute observation rewards. Training infrastructure can
+            access it for introspection:
+
+            ```python
+            for name, r in env.rubric.named_rubrics():
+                print(f"{name}: {r.last_score}")
+            ```
     """
 
     # Class-level flag indicating whether this environment supports concurrent sessions
     SUPPORTS_CONCURRENT_SESSIONS: bool = False
 
+    REQUIRES_SINGLE_THREAD_EXECUTOR: bool = False
+
     # Optional rubric for reward computation
     rubric: Optional["Rubric"]
 
@@ -196,7 +327,7 @@ class Environment(ABC, Generic[ActT, ObsT, StateT]):
         Default implementation returns basic metadata derived from class name.
 
         Returns:
-            EnvironmentMetadata with environment information
+            [`EnvironmentMetadata`] with environment information.
         """
         return EnvironmentMetadata(
             name=self.__class__.__name__,
@@ -214,17 +345,22 @@ class Environment(ABC, Generic[ActT, ObsT, StateT]):
         """Apply rubric if one is provided.
 
         Args:
-            action: The action taken by the agent.
-            observation: The resulting observation.
+            action (`ActT`):
+                The action taken by the agent.
+            observation (`ObsT`):
+                The resulting observation.
 
         Returns:
-            Reward value from the rubric, or 0.0 if no rubric is set.
+            `float`: Reward value from the rubric, or 0.0 if no rubric is set.
 
-        Usage in step():
-            def step(self, action: MyAction, ...) -> MyObservation:
-                # ... execute action and create observation ...
-                observation.reward = self._apply_rubric(action, observation)
-                return observation
+        Call this in `step()` to compute and assign the reward:
+
+        ```python
+        def step(self, action: MyAction, ...) -> MyObservation:
+            # ... execute action and create observation ...
+            observation.reward = self._apply_rubric(action, observation)
+            return observation
+        ```
         """
         if self.rubric is not None:
             return self.rubric(action, observation)
@@ -234,17 +370,22 @@ class Environment(ABC, Generic[ActT, ObsT, StateT]):
         """Apply rubric asynchronously if one is provided.
 
         Args:
-            action: The action taken by the agent.
-            observation: The resulting observation.
+            action (`ActT`):
+                The action taken by the agent.
+            observation (`ObsT`):
+                The resulting observation.
 
         Returns:
-            Reward value from the rubric, or 0.0 if no rubric is set.
+            `float`: Reward value from the rubric, or 0.0 if no rubric is set.
+
+        Call this in `step_async()` to compute and assign the reward:
 
-        Usage in step_async():
-            async def step_async(self, action: MyAction, ...) -> MyObservation:
-                # ... execute action and create observation ...
-                observation.reward = await self._apply_rubric_async(action, observation)
-                return observation
+        ```python
+        async def step_async(self, action: MyAction, ...) -> MyObservation:
+            # ... execute action and create observation ...
+            observation.reward = await self._apply_rubric_async(action, observation)
+            return observation
+        ```
         """
         if self.rubric is not None:
             result = self.rubric(action, observation)
@@ -257,13 +398,14 @@ class Environment(ABC, Generic[ActT, ObsT, StateT]):
     def _reset_rubric(self) -> None:
         """Reset the rubric state if one is provided.
 
-        Call this in reset() to clear any trajectory state in the rubric.
+        Call this in `reset()` to clear any trajectory state in the rubric:
 
-        Usage in reset():
-            def reset(self, ...) -> MyObservation:
-                self._reset_rubric()
-                # ... create initial observation ...
-                return observation
+        ```python
+        def reset(self, ...) -> MyObservation:
+            self._reset_rubric()
+            # ... create initial observation ...
+            return observation
+        ```
         """
         if self.rubric is not None:
             self.rubric.reset()
@@ -271,13 +413,14 @@ class Environment(ABC, Generic[ActT, ObsT, StateT]):
     async def _reset_rubric_async(self) -> None:
         """Reset the rubric state asynchronously if one is provided.
 
-        Call this in reset_async() to clear any trajectory state in the rubric.
+        Call this in `reset_async()` to clear any trajectory state in the rubric:
 
-        Usage in reset_async():
-            async def reset_async(self, ...) -> MyObservation:
-                await self._reset_rubric_async()
-                # ... create initial observation ...
-                return observation
+        ```python
+        async def reset_async(self, ...) -> MyObservation:
+            await self._reset_rubric_async()
+            # ... create initial observation ...
+            return observation
+        ```
         """
         if self.rubric is not None:
             # Check if rubric has async reset method
@@ -288,6 +431,37 @@ class Environment(ABC, Generic[ActT, ObsT, StateT]):
             else:
                 self.rubric.reset()
 
+    def render_web(self, observation: Dict[str, Any]) -> Optional[str]:
+        """
+        HTML that draws an observation in the web interface at `/web`.
+
+        Override it to show the environment's state (a board, a page, a plot).
+        The default draws nothing and the interface lists the observation's fields.
+
+        Args:
+            observation (`dict`):
+                The serialized observation, as returned by `reset()` or `step()`,
+                plus top-level `reward` and `done`.
+        """
+        return None
+
+    def web_actions(
+        self, observation: Dict[str, Any]
+    ) -> List[Tuple[str, Dict[str, Any]]]:
+        """
+        One-click actions for the web interface at `/web`, as `(label, action)` pairs.
+
+        Override it when the valid actions are known from the observation (the legal
+        moves of a game). Each action is the dict `step()` receives. The default
+        offers none and the interface shows the action's form only.
+
+        Args:
+            observation (`dict`):
+                The serialized observation, as returned by `reset()` or `step()`,
+                plus top-level `reward` and `done`.
+        """
+        return []
+
     def close(self) -> None:
         """Clean up resources used by the environment.
 

src/core/env_server/mcp_environment.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 MCP Environment base class for OpenEnv.
@@ -19,7 +15,9 @@ Key features:
 - Mode-aware tool registration (production vs simulation)
 - Code mode support via get_callables() and execute_code()
 
-Usage:
+Examples:
+
+    ```python
     from fastmcp import FastMCP
     from openenv.core.env_server.mcp_environment import MCPEnvironment
 
@@ -50,6 +48,7 @@ Usage:
         def state(self):
             # Return current state
             ...
+    ```
 """
 
 import asyncio
@@ -104,6 +103,33 @@ def get_server_tools(mcp_server: Any) -> Dict[str, Any]:
     return {}
 
 
+def _tool_from_server_tool(tool: Any) -> Tool:
+    """Convert a FastMCP tool object into OpenEnv's Tool model."""
+    return Tool(
+        name=tool.name,
+        description=tool.description or "",
+        input_schema=tool.inputSchema if hasattr(tool, "inputSchema") else {},
+    )
+
+
+def _tool_from_mode_schema(schema: Dict[str, Any]) -> Tool:
+    """Convert a stored mode-aware schema into OpenEnv's Tool model."""
+    return Tool(
+        name=schema["name"],
+        description=schema["description"],
+        input_schema=schema["input_schema"],
+    )
+
+
+def _schema_for_mode(
+    mode_schemas: Dict[Optional[str], Dict[str, Any]], current_mode: Optional[str]
+) -> Optional[Dict[str, Any]]:
+    """Return the schema visible for the current mode, if any."""
+    if None in mode_schemas:
+        return mode_schemas[None]
+    return mode_schemas.get(current_mode)
+
+
 class MCPEnvironment(Environment):
     """
     Base class for environments that expose tools via MCP (Model Context Protocol).
@@ -127,16 +153,21 @@ class MCPEnvironment(Environment):
         ValueError: If any tool in the MCP server uses a reserved name
             (reset, step, state, close).
 
-    Example:
-        >>> from fastmcp import FastMCP
-        >>> mcp = FastMCP("calculator")
-        >>> @mcp.tool()
-        ... def add(a: int, b: int) -> int:
-        ...     return a + b
-        >>> env = MyMCPEnvironment(mcp)
-        >>> obs = env.step(ListToolsAction())
-        >>> obs.tools[0].name
-        'add'
+    Examples:
+
+        ```python
+        from fastmcp import FastMCP
+
+        mcp = FastMCP("calculator")
+
+        @mcp.tool()
+        def add(a: int, b: int) -> int:
+            return a + b
+
+        env = MyMCPEnvironment(mcp)
+        obs = env.step(ListToolsAction())
+        obs.tools[0].name  # 'add'
+        ```
     """
 
     def __init__(self, mcp_server: Any, transform: Optional[Any] = None) -> None:
@@ -398,19 +429,20 @@ class MCPEnvironment(Environment):
         the subclass's _step_impl() method.
 
         Args:
-            action: The action to execute. Can be:
-                - ListToolsAction: Returns available MCP tools
-                - CallToolAction: Invokes a specific MCP tool
-                - Any other Action: Delegated to _step_impl()
-            timeout_s: Optional timeout in seconds for the action.
-                Defaults to MCP_TOOL_CALL_TIMEOUT (30s) for MCP actions.
-            **kwargs: Additional arguments passed to handlers.
+            action (`Action`):
+                The action to execute. `ListToolsAction` returns available MCP tools,
+                `CallToolAction` invokes a specific MCP tool, and any other action
+                is delegated to _step_impl().
+            timeout_s (`float`, *optional*):
+                Timeout in seconds for the action. Defaults to MCP_TOOL_CALL_TIMEOUT
+                (30s) for MCP actions.
+            **kwargs (`Any`):
+                Additional arguments passed to handlers.
 
         Returns:
-            Observation appropriate to the action type:
-                - ListToolsObservation for ListToolsAction
-                - CallToolObservation for CallToolAction
-                - Subclass-defined Observation for other actions
+            `Observation`: `ListToolsObservation` for `ListToolsAction`,
+            `CallToolObservation` for `CallToolAction`, or a subclass-defined
+            Observation for other actions.
         """
         if isinstance(action, ListToolsAction):
             return self._handle_list_tools()
@@ -465,34 +497,11 @@ class MCPEnvironment(Environment):
             tools = []
             for tool in tools_result:
                 if tool.name not in self._mode_tool_schemas:
-                    tools.append(
-                        Tool(
-                            name=tool.name,
-                            description=tool.description or "",
-                            input_schema=tool.inputSchema
-                            if hasattr(tool, "inputSchema")
-                            else {},
-                        )
-                    )
-            for tool_name, mode_schemas in self._mode_tool_schemas.items():
-                if None in mode_schemas:
-                    schema = mode_schemas[None]
-                    tools.append(
-                        Tool(
-                            name=schema["name"],
-                            description=schema["description"],
-                            input_schema=schema["input_schema"],
-                        )
-                    )
-                elif current_mode in mode_schemas:
-                    schema = mode_schemas[current_mode]
-                    tools.append(
-                        Tool(
-                            name=schema["name"],
-                            description=schema["description"],
-                            input_schema=schema["input_schema"],
-                        )
-                    )
+                    tools.append(_tool_from_server_tool(tool))
+            for mode_schemas in self._mode_tool_schemas.values():
+                schema = _schema_for_mode(mode_schemas, current_mode)
+                if schema is not None:
+                    tools.append(_tool_from_mode_schema(schema))
             return ListToolsObservation(tools=tools)
         except Exception as e:
             return ListToolsObservation(

src/core/env_server/mcp_types.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 MCP (Model Context Protocol) type definitions for OpenEnv.
@@ -231,9 +227,8 @@ class ListToolsAction(Action):
     Request list of available tools from the environment.
 
     This action triggers MCP's tools/list operation and returns
-    all available tools with their schemas.
-
-    Note: Does NOT require reset() to be called first.
+    all available tools with their schemas. Does NOT require reset()
+    to be called first.
     """
 
     type: Literal["list_tools"] = Field(

src/core/env_server/route_config.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 Route configuration utilities for declarative FastAPI route registration.
@@ -30,28 +26,32 @@ class GetEndpointConfig:
     description: str
 
 
+def _make_get_endpoint(
+    handler: Callable[[], BaseModel | dict],
+) -> Callable[[], BaseModel | dict]:
+    """Wrap a sync GET handler in the async endpoint FastAPI expects."""
+
+    async def endpoint() -> BaseModel | dict:
+        return handler()
+
+    return endpoint
+
+
 def register_get_endpoints(app: FastAPI, configs: List[GetEndpointConfig]) -> None:
     """
     Register multiple GET endpoints from configuration.
 
     Args:
-        app: FastAPI application instance
-        configs: List of GET endpoint configurations
+        app ([`~fastapi.FastAPI`]):
+            FastAPI application instance.
+        configs (`List[GetEndpointConfig]`):
+            List of GET endpoint configurations.
     """
     for config in configs:
-        # Capture handler in a closure to avoid non-serializable default parameter
-        def make_endpoint(
-            handler: Callable[[], BaseModel | dict],
-        ) -> Callable[[], BaseModel | dict]:
-            async def endpoint() -> BaseModel | dict:
-                return handler()
-
-            return endpoint
-
         app.get(
             config.path,
             response_model=config.response_model,
             tags=[config.tag],
             summary=config.summary,
             description=config.description,
-        )(make_endpoint(config.handler))
+        )(_make_get_endpoint(config.handler))

src/core/env_server/serialization.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 Shared serialization and deserialization utilities for OpenEnv HTTP servers.
@@ -12,6 +8,7 @@ and Pydantic models (Action/Observation) to eliminate code duplication across
 HTTP server and web interface implementations.
 """
 
+import json
 from typing import Any, Dict, Type
 
 from .mcp_types import CallToolAction, ListToolsAction
@@ -27,6 +24,22 @@ _MCP_ACTION_TYPES: Dict[str, Type[Action]] = {
 }
 
 
+def _deserialize_mcp_action(
+    action_data: Dict[str, Any], action_cls: Type[Action]
+) -> Action | None:
+    # Only intercept when action_cls is the generic Action base or itself an
+    # MCP type. This keeps env-specific action validation authoritative.
+    action_type = action_data.get("type")
+    if action_type not in _MCP_ACTION_TYPES:
+        return None
+
+    mcp_cls = _MCP_ACTION_TYPES[action_type]
+    if action_cls is Action or action_cls in _MCP_ACTION_TYPES.values():
+        return mcp_cls.model_validate(action_data)
+
+    return None
+
+
 def deserialize_action(action_data: Dict[str, Any], action_cls: Type[Action]) -> Action:
     """
     Convert JSON dict to Action instance using Pydantic validation.
@@ -40,28 +53,20 @@ def deserialize_action(action_data: Dict[str, Any], action_cls: Type[Action]) ->
     use deserialize_action_with_preprocessing().
 
     Args:
-        action_data: Dictionary containing action data
-        action_cls: The Action subclass to instantiate
+        action_data (`dict`):
+            Dictionary containing action data.
+        action_cls (`type`):
+            The Action subclass to instantiate.
 
     Returns:
-        Action instance
+        `Action` instance.
 
     Raises:
-        ValidationError: If action_data is invalid for the action class
-
-    Note:
-        This uses Pydantic's model_validate() for automatic validation.
+        `ValidationError`: If `action_data` is invalid for the action class.
     """
-    # Route MCP action types before falling through to the env action_cls.
-    # Only intercept when action_cls is the generic Action base or itself an
-    # MCP type (i.e. the server hosts an MCP environment).  This avoids
-    # silently bypassing env-specific validation for non-MCP environments
-    # that happen to use "call_tool" / "list_tools" as a type discriminator.
-    action_type = action_data.get("type")
-    if action_type in _MCP_ACTION_TYPES:
-        mcp_cls = _MCP_ACTION_TYPES[action_type]
-        if action_cls is Action or action_cls in _MCP_ACTION_TYPES.values():
-            return mcp_cls.model_validate(action_data)
+    mcp_action = _deserialize_mcp_action(action_data, action_cls)
+    if mcp_action is not None:
+        return mcp_action
 
     return action_cls.model_validate(action_data)
 
@@ -73,28 +78,36 @@ def deserialize_action_with_preprocessing(
     Convert JSON dict to Action instance with preprocessing for special types.
 
     This version handles common type conversions needed for web interfaces:
+    - Converting JSON string arguments to dict for MCP call_tool actions
     - Converting lists/strings to tensors for 'tokens' field
     - Converting string action_id to int
     - Other custom preprocessing as needed
 
     Args:
-        action_data: Dictionary containing action data
-        action_cls: The Action subclass to instantiate
+        action_data (`dict`):
+            Dictionary containing action data.
+        action_cls (`type`):
+            The Action subclass to instantiate.
 
     Returns:
-        Action instance
+        `Action` instance.
 
     Raises:
-        ValidationError: If action_data is invalid for the action class
+        `ValidationError`: If `action_data` is invalid for the action class.
     """
-    # Route MCP action types before preprocessing (they don't need it).
-    # Same guard as deserialize_action: only intercept when action_cls is
-    # the generic Action base or itself an MCP type.
-    action_type = action_data.get("type")
-    if action_type in _MCP_ACTION_TYPES:
-        mcp_cls = _MCP_ACTION_TYPES[action_type]
-        if action_cls is Action or action_cls in _MCP_ACTION_TYPES.values():
-            return mcp_cls.model_validate(action_data)
+    mcp_data = action_data
+    if action_data.get("type") == "call_tool" and isinstance(
+        action_data.get("arguments"), str
+    ):
+        mcp_data = dict(action_data)
+        try:
+            mcp_data["arguments"] = json.loads(action_data["arguments"])
+        except Exception:
+            pass
+
+    mcp_action = _deserialize_mcp_action(mcp_data, action_cls)
+    if mcp_action is not None:
+        return mcp_action
 
     processed_data = {}
 
@@ -104,8 +117,6 @@ def deserialize_action_with_preprocessing(
             if isinstance(value, str):
                 # If it's a string, try to parse it as a list of numbers
                 try:
-                    import json
-
                     value = json.loads(value)
                 except Exception:
                     # If parsing fails, treat as empty list
@@ -138,34 +149,37 @@ def serialize_observation(observation: Observation) -> Dict[str, Any]:
     Convert Observation instance to JSON-compatible dict using Pydantic.
 
     Args:
-        observation: Observation instance
+        observation (`Observation`):
+            Observation instance to serialize.
 
     Returns:
-        Dictionary compatible with EnvClient._parse_result()
-
-    The format matches what EnvClient expects:
-    {
-        "observation": {...},  # Observation fields
-        "reward": float | None,
-        "done": bool,
-    }
+        `dict` compatible with `EnvClient._parse_result()`, with keys:
+            - `observation` (`dict`): Observation fields.
+            - `reward` (`float` or `None`): Reward value.
+            - `done` (`bool`): Whether the episode is done.
+            - `metadata` (`dict`, *optional*): Additional observation metadata.
     """
-    # Use Pydantic's model_dump() for serialization
+    # Keep metadata in the nested observation payload for backwards
+    # compatibility with typed clients, and also surface it as a top-level
+    # sibling for clients that read the generic wire format directly.
     obs_dict = observation.model_dump(
         exclude={
             "reward",
             "done",
-            "metadata",
         }  # Exclude these from observation dict
     )
 
-    # Extract reward and done directly from the observation
+    # Extract reward, done, and metadata directly from the observation.
     reward = observation.reward
     done = observation.done
+    metadata = observation.metadata
 
-    # Return in EnvClient expected format
-    return {
+    # Return in EnvClient expected format.
+    result = {
         "observation": obs_dict,
         "reward": reward,
         "done": done,
     }
+    if metadata:
+        result["metadata"] = metadata
+    return result

src/core/env_server/types.py

@@ -1,13 +1,9 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 from enum import Enum
 from typing import Annotated, Any, Dict, Literal, Optional, Union
 
-from pydantic import BaseModel, ConfigDict, Field, model_validator
+from pydantic import BaseModel, ConfigDict, Field, model_serializer, model_validator
 
 
 # Type aliases
@@ -121,6 +117,9 @@ class ResetResponse(BaseModel):
     done: bool = Field(
         default=False, description="Whether episode is already done (typically False)"
     )
+    metadata: Optional[Dict[str, Any]] = Field(
+        default=None, description="Additional metadata from the environment"
+    )
 
 
 class StepRequest(BaseModel):
@@ -164,6 +163,9 @@ class StepResponse(BaseModel):
         default=None, description="Reward signal from the action"
     )
     done: bool = Field(default=False, description="Whether the episode has terminated")
+    metadata: Optional[Dict[str, Any]] = Field(
+        default=None, description="Additional metadata from the environment"
+    )
 
 
 class BaseMessage(BaseModel):
@@ -234,6 +236,16 @@ class SchemaResponse(BaseMessage):
     state: Dict[str, Any] = Field(
         description="JSON schema for environment state objects"
     )
+    reset_observation: Optional[Dict[str, Any]] = Field(
+        default=None, description="JSON schema for observations returned by reset"
+    )
+
+    @model_serializer(mode="wrap")
+    def _omit_undeclared_reset_schema(self, handler):
+        data = handler(self)
+        if self.reset_observation is None:
+            data.pop("reset_observation", None)
+        return data
 
 
 class HealthResponse(BaseMessage):
@@ -245,6 +257,33 @@ class HealthResponse(BaseMessage):
     )
 
 
+class ListTasksRequest(BaseMessage):
+    """Request model for ORS-compatible task listing."""
+
+    split: str = Field(description="Task split name")
+
+
+class NumTasksRequest(BaseMessage):
+    """Request model for ORS-compatible task counts."""
+
+    split: str = Field(description="Task split name")
+
+
+class GetTaskRequest(BaseMessage):
+    """Request model for ORS-compatible task lookup."""
+
+    split: str = Field(description="Task split name")
+    index: int = Field(description="Task index within the split")
+
+
+class GetTaskRangeRequest(BaseMessage):
+    """Request model for ORS-compatible task range lookup."""
+
+    split: str = Field(description="Task split name")
+    start: Optional[int] = Field(default=None, description="Inclusive start index")
+    stop: Optional[int] = Field(default=None, description="Exclusive stop index")
+
+
 class WSResetMessage(BaseMessage):
     """WebSocket message to reset the environment."""
 

src/core/env_server/web_interface.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """
 Web interface for OpenEnv environments.
@@ -15,59 +11,79 @@ option (e.g. openenv push --enable-interface) or ENABLE_WEB_INTERFACE env var.
 from __future__ import annotations
 
 import asyncio
+import functools
 import inspect
 import json
 from concurrent.futures import ThreadPoolExecutor
 from datetime import datetime
-from typing import Any, Callable, Dict, List, Optional, Type
+from typing import Any, Callable, Dict, get_origin, List, Optional, Tuple, Type
 
 import gradio as gr
 from fastapi import Body, FastAPI, HTTPException, status, WebSocket, WebSocketDisconnect
 from fastapi.responses import RedirectResponse
 from pydantic import BaseModel, ConfigDict, Field
 
+from ._utils import overrides_method
 from .gradio_theme import OPENENV_GRADIO_CSS, OPENENV_GRADIO_THEME
 from .gradio_ui import build_gradio_app, get_gradio_display_title
 from .interfaces import Environment
 from .serialization import deserialize_action_with_preprocessing, serialize_observation
 from .types import Action, EnvironmentMetadata, Observation, State
 
-# Quick Start markdown template; placeholders match init suffixes (__ENV_NAME__, __ENV_CLASS_NAME__*).
-DEFAULT_QUICK_START_MARKDOWN = """
-### Connect to this environment
 
-Connect from Python using `__ENV_CLASS_NAME__Env`:
+_CLIENT_BASES = {"EnvClient", "MCPClientBase", "MCPToolClient"}
 
-```python
-from __ENV_NAME__ import __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Env
 
-with __ENV_CLASS_NAME__Env.from_env("<SPACE_ID>") as env:
-    result = await env.step(__ENV_CLASS_NAME__Action(message="..."))
-```
-
-Or connect directly to a running server:
-
-```python
-env = __ENV_CLASS_NAME__Env(base_url="http://localhost:8000")
-```
-
-### Contribute to this environment
-
-Submit improvements via pull request on the Hugging Face Hub.
-
-```bash
-openenv fork <SPACE_ID> --repo-id <your-username>/<your-repo-name>
-```
-
-Then make your changes and submit a pull request:
-
-```bash
-cd <forked-repo>
-openenv push <SPACE_ID> --create-pr
-```
+def _read_package(package: str) -> Tuple[Optional[str], set]:
+    """
+    The client class defined in the package's `client.py` and the names its
+    `__init__.py` exports. The files are parsed, not imported, so the server
+    never loads client code.
+    """
+    import ast
+    import importlib.util
+    from pathlib import Path
 
-For more information, see the [OpenEnv documentation](https://meta-pytorch.org/OpenEnv/).
-"""
+    spec = importlib.util.find_spec(package)
+    if spec is None or not spec.submodule_search_locations:
+        return None, set()
+    root = Path(next(iter(spec.submodule_search_locations)))
+
+    def parse(name: str) -> Optional[ast.Module]:
+        path = root / name
+        return ast.parse(path.read_text()) if path.exists() else None
+
+    client = None
+    tree = parse("client.py")
+    for node in tree.body if tree else []:
+        if isinstance(node, ast.ClassDef):
+            bases = {ast.unparse(b).split("[")[0].split(".")[-1] for b in node.bases}
+            if bases & _CLIENT_BASES:
+                client = node.name
+                break
+
+    exports = set()
+    tree = parse("__init__.py")
+    for node in ast.walk(tree) if tree else []:
+        if isinstance(node, ast.ImportFrom):
+            exports.update(alias.asname or alias.name for alias in node.names)
+        elif isinstance(node, ast.Constant) and isinstance(node.value, str):
+            exports.add(node.value)
+    return client, exports
+
+
+def _example_value(annotation: Any) -> str:
+    """A placeholder literal for a field of the given type."""
+    origin = get_origin(annotation) or annotation
+    literals = {
+        bool: "False",
+        int: "0",
+        float: "0.0",
+        str: '"..."',
+        list: "[]",
+        dict: "{}",
+    }
+    return literals.get(origin, "...")
 
 
 def get_quick_start_markdown(
@@ -76,35 +92,86 @@ def get_quick_start_markdown(
     observation_cls: Type[Observation],
 ) -> str:
     """
-    Build Quick Start markdown with class names replaced from current env (init-style suffixes).
+    Build the Quick Start markdown shown next to the web interface.
 
-    Uses the same placeholder names as the init template so that __ENV_CLASS_NAME__Env,
-    __ENV_CLASS_NAME__Action, __ENV_CLASS_NAME__Observation and __ENV_NAME__ are
-    replaced with the actual class/package names.
+    The client class is the one the environment package exports, the URL is the
+    Space's own when running on Hugging Face Spaces, and the example action uses
+    the action's real fields.
     """
     import os
 
-    # Prefix from action class (e.g. EchoAction -> Echo)
-    action_name = getattr(action_cls, "__name__", "Action")
-    if action_name.endswith("Action"):
-        prefix = action_name[: -len("Action")]
-    else:
-        prefix = action_name.replace("Action", "").strip() or "Env"
-
-    env_client_name = f"{prefix}Env"
-    obs_name = getattr(observation_cls, "__name__", "Observation")
-    pkg_name = (metadata.name if metadata else "env").replace(" ", "_").lower()
-
-    space_id = os.environ.get("SPACE_ID", "<hf-username>/<hf-repo-name>")
+    from .mcp_types import CallToolAction
+
+    package = (metadata.name if metadata else "env").replace(" ", "_").lower()
+    client, exports = _read_package(package)
+    if client is None:
+        # The env name can differ from its package (reasoning_gym vs reasoning_gym_env).
+        module_root = action_cls.__module__.split(".")[0]
+        if module_root not in (package, "openenv", "models", "server"):
+            package = module_root
+            client, exports = _read_package(package)
+    space_id = os.environ.get("SPACE_ID")
+    space_host = os.environ.get("SPACE_HOST")
+    base_url = f"https://{space_host}" if space_host else "http://localhost:8000"
+
+    if client is None:
+        return (
+            "### Connect to this environment\n\n"
+            f"The server is at `{base_url}`. See the environment's README for how to connect."
+        )
 
-    content = DEFAULT_QUICK_START_MARKDOWN
-    content = content.replace("__ENV_CLASS_NAME__Env", env_client_name)
-    content = content.replace("__ENV_CLASS_NAME__Action", action_name)
-    content = content.replace("__ENV_CLASS_NAME__Observation", obs_name)
-    content = content.replace("__ENV_CLASS_NAME__", prefix)
-    content = content.replace("__ENV_NAME__", pkg_name)
-    content = content.replace("<SPACE_ID>", space_id)
-    return content.strip()
+    lines = ["### Connect to this environment", ""]
+    if space_id:
+        lines += [
+            "```bash",
+            f"pip install git+https://huggingface.co/spaces/{space_id}",
+            "```",
+            "",
+        ]
+
+    if issubclass(action_cls, CallToolAction):
+        code = [
+            f"from {package} import {client}",
+            "",
+            f'with {client}(base_url="{base_url}").sync() as env:',
+            "    env.reset()",
+            "    print([tool.name for tool in env.list_tools()])",
+        ]
+        after = "Then call a tool with `env.call_tool(name, **arguments)`."
+    else:
+        action = action_cls.__name__
+        exported = action in exports
+        action_module = package if exported else action_cls.__module__
+        fields = ", ".join(
+            f"{name}={_example_value(field.annotation)}"
+            for name, field in action_cls.model_fields.items()
+            if field.is_required() and name != "metadata"
+        )
+        code = [
+            *(
+                [f"from {package} import {action}, {client}"]
+                if exported
+                else [
+                    f"from {action_module} import {action}",
+                    f"from {package} import {client}",
+                ]
+            ),
+            "",
+            f'with {client}(base_url="{base_url}").sync() as env:',
+            "    result = env.reset()",
+            f"    result = env.step({action}({fields}))",
+            "    print(result.observation, result.reward)",
+        ]
+        after = ""
+
+    lines += ["```python", *code, "```"]
+    if after:
+        lines += ["", after]
+    lines += [
+        "",
+        "More in the [OpenEnv documentation](https://huggingface.co/docs/openenv).",
+    ]
+    return "\n".join(lines)
 
 
 def load_environment_metadata(
@@ -114,39 +181,29 @@ def load_environment_metadata(
     Load environment metadata including README content.
 
     Args:
-        env: The environment instance, class, or factory function.
-             - If a class: used as a factory, won't call instance methods
-             - If a function: used as a factory, won't call instance methods
-             - If an instance: may call get_metadata() if available
-        env_name: Optional environment name for README file lookup
+        env: The environment instance or factory (class, function or
+            `functools.partial`). A factory is not instantiated. If an instance,
+            `get_metadata()` is called.
+        env_name (`str`, *optional*):
+            Optional environment name for README file lookup.
 
     Returns:
-        EnvironmentMetadata with loaded information
+        `EnvironmentMetadata` with loaded information.
     """
-    import inspect
-
-    # Determine what type of env we received:
-    # 1. A class (used as factory) - e.g., PythonCodeActEnv
-    # 2. A function (factory function) - e.g., create_chat_environment
-    # 3. An actual instance - e.g., SnakeEnvironment()
-    is_class = inspect.isclass(env)
-    is_function = inspect.isfunction(env) or inspect.ismethod(env)
-    is_factory = is_class or is_function
-
-    # Try to get metadata from environment if it's an instance with get_metadata
-    if not is_factory and hasattr(env, "get_metadata"):
+    # An Environment instance provides its own metadata. Anything else is a factory:
+    # a class (e.g. PythonCodeActEnv), a factory function (e.g.
+    # create_chat_environment) or a `functools.partial` of either.
+    if isinstance(env, Environment):
         return env.get_metadata()
 
+    if isinstance(env, functools.partial):
+        env = env.func
+
     # Determine the class name for default metadata
-    if is_class:
-        # env is the class itself
+    if inspect.isclass(env):
         class_name = env.__name__
-    elif is_function:
-        # env is a factory function - use its name or derive from env_name
-        class_name = env_name or env.__name__
     else:
-        # env is an instance
-        class_name = env.__class__.__name__
+        class_name = env_name or env.__name__
 
     # Default metadata
     metadata = EnvironmentMetadata(
@@ -167,22 +224,14 @@ def _load_readme_from_filesystem(env_name: Optional[str]) -> Optional[str]:
     """
     Load README content from the filesystem.
 
-    Tries multiple locations:
-    1. Container filesystem: /app/README.md
-    2. Local development: src/envs/{env_name}/README.md
-    3. Environment variable: ENV_README_PATH
+    Tries multiple locations in order: the path given by the `ENV_README_PATH`
+    environment variable, the repository path `envs/{env_name}/README.md`, and the
+    container paths `/app/env/README.md` (where Space images copy the environment)
+    and `/app/README.md`.
     """
     import os
     from pathlib import Path
 
-    # Try container filesystem first
-    container_readme = Path("/app/README.md")
-    if container_readme.exists():
-        try:
-            return container_readme.read_text(encoding="utf-8")
-        except Exception:
-            pass
-
     # Try environment variable path
     custom_path = os.environ.get("ENV_README_PATH")
     if custom_path and Path(custom_path).exists():
@@ -191,12 +240,12 @@ def _load_readme_from_filesystem(env_name: Optional[str]) -> Optional[str]:
         except Exception:
             pass
 
-    # Try local development path
-    if env_name:
-        local_readme = Path(f"src/envs/{env_name}/README.md")
-        if local_readme.exists():
+    candidates = [Path(f"envs/{env_name}/README.md")] if env_name else []
+    candidates += [Path("/app/env/README.md"), Path("/app/README.md")]
+    for readme in candidates:
+        if readme.exists():
             try:
-                return local_readme.read_text(encoding="utf-8")
+                return readme.read_text(encoding="utf-8")
             except Exception:
                 pass
 
@@ -248,13 +297,8 @@ class WebInterfaceManager:
         observation_cls: Type[Observation],
         metadata: Optional[EnvironmentMetadata] = None,
     ):
-        import inspect
-
-        # If env is a class or factory function, instantiate it
-        if inspect.isclass(env) or inspect.isfunction(env):
-            self.env = env()
-        else:
-            self.env = env
+        # Anything that is not an Environment instance is a factory, instantiate it
+        self.env = env if isinstance(env, Environment) else env()
         self.action_cls = action_cls
         self.observation_cls = observation_cls
         self.metadata = metadata or EnvironmentMetadata(
@@ -347,7 +391,7 @@ class WebInterfaceManager:
         """Reset the environment and update state."""
         reset_kwargs = reset_kwargs or {}
 
-        is_async = self.env.reset_async.__func__ is not Environment.reset_async
+        is_async = overrides_method(self.env.reset_async, Environment.reset_async)
         sig = inspect.signature(self.env.reset_async if is_async else self.env.reset)
         valid_kwargs = self._get_valid_kwargs(sig, reset_kwargs)
 
@@ -433,6 +477,14 @@ def create_web_interface_app(
     max_concurrent_envs: Optional[int] = None,
     concurrency_config: Optional[Any] = None,
     gradio_builder: Optional[Callable[..., Any]] = None,
+    custom_tab_name: str = "Custom",
+    custom_tab_primary: bool = False,
+    show_default_tab: bool = True,
+    title_override: Optional[str] = None,
+    state_cls: Type[State] = State,
+    *,
+    mode: Optional[Any] = None,
+    reset_observation_cls: Optional[Type[Observation]] = None,
 ) -> FastAPI:
     """
     Create a FastAPI application with web interface for the given environment.
@@ -440,13 +492,36 @@ def create_web_interface_app(
     Args:
         env: The Environment instance to serve
         action_cls: The Action subclass this environment expects
-        observation_cls: The Observation subclass this environment returns
+        observation_cls: The Observation subclass returned by step
+        reset_observation_cls: The reset observation model published in /schema.
+            Defaults to observation_cls.
         env_name: Optional environment name for README loading
         max_concurrent_envs: Maximum concurrent WebSocket sessions
         concurrency_config: Optional ConcurrencyConfig for advanced concurrency settings
         gradio_builder: Optional callable (web_manager, action_fields, metadata,
             is_chat_env, title, quick_start_md) -> gr.Blocks to use instead of the
             default Gradio UI. Lets envs replace or customize the /web interface.
+        custom_tab_name: Label shown on the env-specific tab when ``gradio_builder``
+            is provided. Defaults to ``"Custom"`` for backwards compatibility; envs
+            that ship a rich custom UI should pass a descriptive name
+            (e.g. ``"REPL"``). Ignored when ``show_default_tab=False`` (no tab chrome
+            is rendered).
+        custom_tab_primary: When True, the env-specific tab is rendered first and
+            selected by default; the auto-generated Playground becomes secondary.
+            Use this for envs whose custom tab is the real interaction surface
+            (so visitors don't land on a less informative schema form). Ignored
+            when ``show_default_tab=False``.
+        show_default_tab: When False, the auto-generated Playground tab is not
+            rendered and the env's ``gradio_builder`` output is mounted directly
+            (single-view UI, no tab chrome). Only meaningful when
+            ``gradio_builder`` is provided.
+        title_override: If set, used verbatim as the Gradio app/browser-tab
+            title instead of the default ``"OpenEnv Agentic Environment: {name}"``.
+        state_cls: The State subclass this environment reports. Used for the /state
+            response model and the state entry of /schema. Defaults to State.
+        mode: Server mode (``ServerMode`` or string). When ``None``, resolved
+            from the ``OPENENV_MODE`` environment variable, defaulting to
+            simulation.
 
     Returns:
         FastAPI application instance with web interface
@@ -455,7 +530,15 @@ def create_web_interface_app(
 
     # Create the base environment app
     app = create_fastapi_app(
-        env, action_cls, observation_cls, max_concurrent_envs, concurrency_config
+        env,
+        action_cls,
+        observation_cls,
+        max_concurrent_envs,
+        concurrency_config,
+        env_name=env_name,
+        state_cls=state_cls,
+        reset_observation_cls=reset_observation_cls,
+        mode=mode,
     )
 
     # Load environment metadata
@@ -484,7 +567,7 @@ def create_web_interface_app(
     async def websocket_ui_endpoint(websocket: WebSocket):
         """WebSocket endpoint for web UI real-time updates.
 
-        Note: Uses /ws/ui to avoid conflict with /ws in http_server.py
+        Uses `/ws/ui` to avoid conflict with `/ws` in `http_server.py`,
         which is used for concurrent environment sessions.
         """
         await web_manager.connect_websocket(websocket)
@@ -534,35 +617,19 @@ def create_web_interface_app(
     is_chat_env = _is_chat_env(action_cls)
     quick_start_md = get_quick_start_markdown(metadata, action_cls, observation_cls)
 
-    default_blocks = build_gradio_app(
+    display_title = title_override or get_gradio_display_title(metadata)
+    gradio_blocks = _build_gradio_blocks(
         web_manager,
         action_fields,
         metadata,
         is_chat_env,
-        title=metadata.name,
-        quick_start_md=quick_start_md,
+        quick_start_md,
+        display_title=display_title,
+        gradio_builder=gradio_builder,
+        custom_tab_name=custom_tab_name,
+        custom_tab_primary=custom_tab_primary,
+        show_default_tab=show_default_tab,
     )
-    if gradio_builder is not None:
-        custom_blocks = gradio_builder(
-            web_manager,
-            action_fields,
-            metadata,
-            is_chat_env,
-            metadata.name,
-            quick_start_md,
-        )
-        if not isinstance(custom_blocks, gr.Blocks):
-            raise TypeError(
-                f"gradio_builder must return a gr.Blocks instance, "
-                f"got {type(custom_blocks).__name__}"
-            )
-        gradio_blocks = gr.TabbedInterface(
-            [default_blocks, custom_blocks],
-            tab_names=["Playground", "Custom"],
-            title=get_gradio_display_title(metadata),
-        )
-    else:
-        gradio_blocks = default_blocks
     app = gr.mount_gradio_app(
         app,
         gradio_blocks,
@@ -574,6 +641,65 @@ def create_web_interface_app(
     return app
 
 
+def _build_gradio_blocks(
+    web_manager: WebInterfaceManager,
+    action_fields: List[Dict[str, Any]],
+    metadata: EnvironmentMetadata,
+    is_chat_env: bool,
+    quick_start_md: str,
+    *,
+    display_title: str,
+    gradio_builder: Optional[Callable[..., Any]],
+    custom_tab_name: str,
+    custom_tab_primary: bool,
+    show_default_tab: bool,
+) -> gr.Blocks:
+    """Build the Gradio block tree mounted at `/web`."""
+    default_blocks = build_gradio_app(
+        web_manager,
+        action_fields,
+        metadata,
+        is_chat_env,
+        title=metadata.name,
+        quick_start_md=quick_start_md,
+    )
+    default_blocks.title = display_title
+
+    if gradio_builder is None:
+        return default_blocks
+
+    custom_blocks = gradio_builder(
+        web_manager,
+        action_fields,
+        metadata,
+        is_chat_env,
+        display_title,
+        quick_start_md,
+    )
+    if not isinstance(custom_blocks, gr.Blocks):
+        raise TypeError(
+            f"gradio_builder must return a gr.Blocks instance, "
+            f"got {type(custom_blocks).__name__}"
+        )
+
+    if not show_default_tab:
+        # No TabbedInterface wrapper to carry the app title.
+        custom_blocks.title = display_title
+        return custom_blocks
+
+    if custom_tab_primary:
+        tab_blocks = [custom_blocks, default_blocks]
+        tab_labels = [custom_tab_name, "Playground"]
+    else:
+        tab_blocks = [default_blocks, custom_blocks]
+        tab_labels = ["Playground", custom_tab_name]
+
+    # The playground has its own header, so the title only names the browser tab.
+    tabbed = gr.TabbedInterface(tab_blocks, tab_names=tab_labels)
+    tabbed.title = display_title
+    return tabbed
+
+
 def _is_chat_env(action_cls: Type[Action]) -> bool:
     """Return True if the action class is a chat-style env (tokens field)."""
     if hasattr(action_cls, "model_fields"):

src/core/evals/__init__.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """Evaluation harness support for OpenEnv."""
 

src/core/evals/base.py

@@ -1,8 +1,4 @@
-# Copyright (c) Meta Platforms, Inc. and affiliates.
-# All rights reserved.
-#
-# This source code is licensed under the BSD-style license found in the
-# LICENSE file in the root directory of this source tree.
+# SPDX-License-Identifier: BSD-3-Clause
 
 """Base class for evaluation harnesses."""